#import json
import os
#import shutil
#import requests

import gradio as gr

from transformers.utils import logging
from langchain.embeddings import HuggingFaceInstructEmbeddings, GooglePalmEmbeddings
import pinecone
from langchain.vectorstores import Pinecone

logging.set_verbosity_debug()

instructor_embeddings = HuggingFaceInstructEmbeddings(model_name="hkunlp/instructor-xl", model_kwargs={"device": "cpu"})

HF_TOKEN = os.environ.get("HF_TOKEN", None)
PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY", None)
PINECONE_ENV = os.environ.get("PINECONE_ENV", None)
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", None)

pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENV)

from langchain.llms import GooglePalm
from langchain.chains import RetrievalQAWithSourcesChain
llm=GooglePalm(google_api_key=GOOGLE_API_KEY, temperature=1, max_output_tokens=2048)
vectorStore = Pinecone.from_existing_index('tennis', instructor_embeddings)
retriever = vectorStore.as_retriever(search_kwargs={"k": 3})
qa_chain_instrucEmbed = RetrievalQAWithSourcesChain.from_chain_type(llm=llm,
                                  chain_type="stuff",
                                  retriever=retriever,
                                  return_source_documents=True,
                                  verbose=True
                                  )

theme = gr.themes.Monochrome(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    radius_size=gr.themes.sizes.radius_sm,
    font=[
        gr.themes.GoogleFont("Open Sans"),
        "ui-sans-serif",
        "system-ui",
        "sans-serif",
    ],
)

def generate(question):
    ret = qa_chain_instrucEmbed(question)
    print(str(ret))
    answer = ret['answer']
    sources = ret['sources']
    embed_video_html = ''
    if sources is not None and len(sources) > 0:
        sources = [s.strip() for s in sources.split(',')]
        for source in sources:
            embed_video_html += f'''
                <iframe width="560" height="315" src="https://www.youtube.com/embed/{source}"
                title="YouTube video player" frameborder="0" allow="accelerometer; autoplay;
                clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
                '''
    return answer, embed_video_html

examples = [
    "Tell me step by step how to find out my dominant eye when I play tennis.",
    "hat do we look for in a great tennis player? Write out the essential attributes."
]

def process_example(args):
    for x in generate(args):
        pass
    return x


css = ".generating {visibility: hidden}"

monospace_css = """
#q-input textarea {
    font-family: monospace, 'Consolas', Courier, monospace;
}
"""

css += monospace_css + ".gradio-container {color: black}"

description = """
<div style="text-align: center;">
    <h1>Ask Coach Patrick Mouratoglou</h1>
</div>
<div style="text-align: left;">
    <p>This is a demo to answer some popular questions from tennis fans to Coach Patrick. The information is being extracted from his official <a href="https://www.youtube.com/@patrickmouratoglou_official" style='color: #e6b800;'>Youtube channel</a>. It's using the following technologies:</p>
    <ul>
        <li>Google PALM</li>
        <li>Gradio</li>
        <li>hkunlp/instructor-xl</li>
        <li>HuggingFace</li>
        <li>Langchain</li>
        <li>Pinecone</li>
    </ul>
</div>
"""
disclaimer = """⚠️<b>This is an unofficial website.</b>\
 <br>**Intended Use**: this app for demonstration purposes; not to serve as replacement for Coach Patrick official media channels or personal expertise."""

with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
    with gr.Column():
        gr.Markdown(description)
        gr.Markdown(disclaimer)
        with gr.Row():
            with gr.Column():
                instruction = gr.Textbox(
                    placeholder="Enter your question here",
                    lines=5,
                    label="Input",
                    elem_id="q-input",
                )
                submit = gr.Button("Ask", variant="primary")
                output = gr.Code(elem_id="q-output", lines=10, label="Output")
                video = gr.HTML('')
                gr.Examples(
                    examples=examples,
                    inputs=[instruction],
                    cache_examples=False,
                    fn=process_example,
                    outputs=[output, video],
                )

    submit.click(
        generate,
        inputs=[instruction],
        outputs=[output, video],
    )
demo.queue(concurrency_count=16).launch(debug=True)