#from climateqa.engine.vectorstore import get_pinecone_vectorstore, 
from climateqa.engine.vectorstore import build_vectores_stores
from climateqa.engine.embeddings import get_embeddings_function
from climateqa.engine.rag import make_rag_papers_chain
from climateqa.engine.keywords import make_keywords_chain
from climateqa.sample_questions import QUESTIONS
from climateqa.engine.text_retriever import ClimateQARetriever
from climateqa.engine.rag import make_rag_chain
from climateqa.engine.llm import get_llm
from utils import create_user_id
from datetime import datetime
import json
import re
import gradio as gr
from climateqa.papers.openalex import OpenAlex
from sentence_transformers import CrossEncoder

reranker = CrossEncoder("mixedbread-ai/mxbai-rerank-xsmall-v1")
oa = OpenAlex()

# Load environment variables in local mode
try:
    from dotenv import load_dotenv
    load_dotenv()
except Exception as e:
    pass

# Set up Gradio Theme
theme = gr.themes.Base(
    primary_hue="blue",
    secondary_hue="red",
    font=[gr.themes.GoogleFont("Poppins"), "ui-sans-serif",
          "system-ui", "sans-serif"],
)


init_prompt = ""

system_template = {
    "role": "system",
    "content": init_prompt,
}

user_id = create_user_id()


def parse_output_llm_with_sources(output):
    # Split the content into a list of text and "[Doc X]" references
    content_parts = re.split(r'\[(Doc\s?\d+(?:,\s?Doc\s?\d+)*)\]', output)
    parts = []
    for part in content_parts:
        if part.startswith("Doc"):
            subparts = part.split(",")

            subparts = [subpart.lower().replace("doc", "").strip()
                        for subpart in subparts]
            subparts = [f"""<a href="#doc{subpart}" class="a-doc-ref" target="_self"><span class='doc-ref'><sup style="color:#FFC000 !important;">({
                subpart})</sup></span></a>""" for subpart in subparts]
            parts.append("".join(subparts))
        else:
            parts.append(part)
    content_parts = "".join(parts)
    return content_parts


def serialize_docs(docs):
    new_docs = []
    for doc in docs:
        new_doc = {}
        new_doc["page_content"] = doc.page_content
        new_doc["metadata"] = doc.metadata
        new_docs.append(new_doc)
    return new_docs


# Create vectorstore and retriever
embeddings_function = get_embeddings_function()

#vectorstore = get_pinecone_vectorstore(embeddings_function)
vectorstore = build_vectores_stores("./sources")
llm = get_llm(provider="openai", max_tokens=1024, temperature=0.0)


async def chat(query, history):
    """taking a query and a message history, use a pipeline (reformulation, retriever, answering) to yield a tuple of:
    (messages in gradio format, messages in langchain format, source documents)"""

    print(f">> NEW QUESTION : {query}")

    retriever = ClimateQARetriever(
        vectorstore=vectorstore, sources=["Custom"], reports=[])
    rag_chain = make_rag_chain(retriever, llm)

    inputs = {"query": query, "audience": None}
    result = rag_chain.astream_log(inputs)

    path_reformulation = "/logs/reformulation/final_output"
    path_keywords = "/logs/keywords/final_output"
    path_retriever = "/logs/find_documents/final_output"
    path_answer = "/logs/answer/streamed_output_str/-"

    docs_html = ""
    output_query = ""
    output_language = ""
    output_keywords = ""
    gallery = []

    try:
        async for op in result:

            op = op.ops[0]

            if op['path'] == path_reformulation:  # reforulated question
                try:
                    output_language = op['value']["language"]  # str
                    output_query = op["value"]["question"]
                except Exception as e:
                    raise gr.Error(f"ClimateQ&A Error: {e} - The error has been noted, try another question and if the error remains, you can contact us :)")

            if op["path"] == path_keywords:
                try:
                    output_keywords = op['value']["keywords"]  # str
                    output_keywords = " AND ".join(output_keywords)
                except Exception as e:
                    pass

            elif op['path'] == path_retriever:  # documents
                try:
                    docs = op['value']['docs']  # List[Document]
                    docs_html = []
                    for i, d in enumerate(docs, 1):
                        docs_html.append(make_html_source(d, i))
                    docs_html = "".join(docs_html)
                except TypeError:
                    print("No documents found")
                    print("op: ", op)
                    continue

            elif op['path'] == path_answer:  # final answer
                new_token = op['value']  # str
                # time.sleep(0.01)
                previous_answer = history[-1][1]
                previous_answer = previous_answer if previous_answer is not None else ""
                answer_yet = previous_answer + new_token
                answer_yet = parse_output_llm_with_sources(answer_yet)
                history[-1] = (query, answer_yet)

            else:
                continue

            history = [tuple(x) for x in history]
            yield history, docs_html, output_query, output_language, gallery, output_query, output_keywords

    except Exception as e:
        raise gr.Error(f"{e}")

    timestamp = str(datetime.now().timestamp())
    log_file = "logs/" + timestamp + ".json"
    prompt = history[-1][0]
    logs = {
        "user_id": str(user_id),
        "prompt": prompt,
        "query": prompt,
        "question": output_query,
        "sources": ["Custom"],
        "docs": serialize_docs(docs),
        "answer": history[-1][1],
        "time": timestamp,
    }
    log_locally(log_file, logs)

    yield history, docs_html, output_query, output_language, gallery, output_query, output_keywords


def make_html_source(source, i):
    # Prépare le contenu HTML pour un fichier texte
    text_content = source.page_content.strip()
    meta = source.metadata
    # Nom de la source
    name = f"<b>Document {i}</b>"

    # Contenu HTML de la carte
    card = f"""
    <div class="card" id="doc{i}">
        <div class="card-content">
            <div>
                <div style="float:right;width 10%;position:relative;top:0px">
                    <a href='{meta['ax_url']}'><img style="width:20px" src='/file/assets/download.png' /></a>
                </div>
                <div>
                    <h2>Extrait {i}</h2>
                    <h2> {meta['ax_name']} - Page {int(meta['ax_page'])}</h2>
                </div>
            </div>
            <p>{text_content}</p>

        </div>
        <div class="card-footer">
            <span>{name}</span>
        </div>
    </div>
    """

    return card


def log_locally(file, logs):
    # Convertit les logs en format JSON
    logs_json = json.dumps(logs)

    # Écrit les logs dans un fichier local
    with open(file, 'w') as f:
        f.write(logs_json)


def generate_keywords(query):
    chain = make_keywords_chain(llm)
    keywords = chain.invoke(query)
    keywords = " AND ".join(keywords["keywords"])
    return keywords


papers_cols_widths = {
    "doc": 50,
    "id": 100,
    "title": 300,
    "doi": 100,
    "publication_year": 100,
    "abstract": 500,
    "rerank_score": 100,
    "is_oa": 50,
}

papers_cols = list(papers_cols_widths.keys())
papers_cols_widths = list(papers_cols_widths.values())


async def find_papers(query, keywords, after):

    summary = ""

    df_works = oa.search(keywords, after=after)
    df_works = df_works.dropna(subset=["abstract"])
    df_works = oa.rerank(query, df_works, reranker)
    df_works = df_works.sort_values("rerank_score", ascending=False)
    G = oa.make_network(df_works)

    height = "750px"
    network = oa.show_network(
        G, color_by="rerank_score", notebook=False, height=height)
    network_html = network.generate_html()

    network_html = network_html.replace("'", "\"")
    css_to_inject = "<style>#mynetwork { border: none !important; } .card { border: none !important; }</style>"
    network_html = network_html + css_to_inject

    network_html = f"""<iframe style="width: 100%; height: {height};margin:0 auto" name="result" allow="midi; geolocation; microphone; camera;
    display-capture; encrypted-media;" sandbox="allow-modals allow-forms
    allow-scripts allow-same-origin allow-popups
    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
    allowpaymentrequest="" frameborder="0" srcdoc='{network_html}'></iframe>"""

    docs = df_works["content"].head(15).tolist()

    df_works = df_works.reset_index(
        drop=True).reset_index().rename(columns={"index": "doc"})
    df_works["doc"] = df_works["doc"] + 1
    df_works = df_works[papers_cols]

    yield df_works, network_html, summary

    chain = make_rag_papers_chain(llm)
    result = chain.astream_log(
        {"question": query, "docs": docs, "language": "English"})
    path_answer = "/logs/StrOutputParser/streamed_output/-"

    async for op in result:

        op = op.ops[0]

        if op['path'] == path_answer:  # reforulated question
            new_token = op['value']  # str
            summary += new_token
        else:
            continue
        yield df_works, network_html, summary


# --------------------------------------------------------------------
# Gradio
# --------------------------------------------------------------------


init_prompt = """
Hello, I am Clara, an AI Assistant created by Axionable. My purpose is to answer your questions using the provided extracted passages, context, and guidelines.

❓ How to interact with Clara

Ask your question: You can ask me anything you want to know. I'll provide an answer based on the extracted passages and other relevant sources.
Response structure: I aim to provide clear and structured answers using the given data.
Guidelines: I follow specific guidelines to ensure that my responses are accurate and useful.
⚠️ Limitations
Though I do my best to help, there might be times when my responses are incorrect or incomplete. If that happens, please feel free to ask for more information or provide feedback to help improve my performance.

What would you like to know today?
"""


with gr.Blocks(title="CLARA", css="style.css", theme=theme, elem_id="main-component") as demo:

    with gr.Tab("CLARA"):

        with gr.Row(elem_id="chatbot-row"):
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(
                    value=[(None, init_prompt)],
                    show_copy_button=True, show_label=False, elem_id="chatbot", layout="panel",
                    avatar_images=(None, "assets/logo4.png"))

                with gr.Row(elem_id="input-message"):
                    textbox = gr.Textbox(placeholder="Posez votre question", show_label=False,
                                         scale=7, lines=1, interactive=True, elem_id="input-textbox")

            with gr.Column(scale=1, variant="panel", elem_id="right-panel"):

                with gr.Tabs() as tabs:

                    with gr.Tab("Sources", elem_id="tab-citations", id=1):
                        sources_textbox = gr.HTML(
                            show_label=False, elem_id="sources-textbox")
                        docs_textbox = gr.State("")

# ---------------------------------------------------------------------------------------
# OTHER TABS
# ---------------------------------------------------------------------------------------

    with gr.Tab("Figures", elem_id="tab-images", elem_classes="max-height other-tabs"):
        gallery_component = gr.Gallery()

    with gr.Tab("Papers (beta)", elem_id="tab-papers", elem_classes="max-height other-tabs"):

        with gr.Row():
            with gr.Column(scale=1):
                query_papers = gr.Textbox(
                    placeholder="Question", show_label=False, lines=1, interactive=True, elem_id="query-papers")
                keywords_papers = gr.Textbox(
                    placeholder="Keywords", show_label=False, lines=1, interactive=True, elem_id="keywords-papers")
                after = gr.Slider(minimum=1950, maximum=2023, step=1, value=1960,
                                  label="Publication date", show_label=True, interactive=True, elem_id="date-papers")
                search_papers = gr.Button(
                    "Search", elem_id="search-papers", interactive=True)

            with gr.Column(scale=7):

                with gr.Tab("Summary", elem_id="papers-summary-tab"):
                    papers_summary = gr.Markdown(
                        visible=True, elem_id="papers-summary")

                with gr.Tab("Relevant papers", elem_id="papers-results-tab"):
                    papers_dataframe = gr.Dataframe(
                        visible=True, elem_id="papers-table", headers=papers_cols)

                with gr.Tab("Citations network", elem_id="papers-network-tab"):
                    citations_network = gr.HTML(
                        visible=True, elem_id="papers-citations-network")

    with gr.Tab("À propos", elem_classes="max-height other-tabs"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown(
                    "CLARA (Climate LLM for Adaptation & Risks Answers) by [Axionable](https://www.axionable.com/)"
                    "– Fork de [ClimateQ&A](https://huggingface.co/spaces/Ekimetrics/climate-question-answering/tree/main)")

    def start_chat(query, history):
        history = history + [(query, None)]
        history = [tuple(x) for x in history]
        return (gr.update(interactive=False), gr.update(selected=1), history)

    def finish_chat():
        return (gr.update(interactive=True, value=""))

    (textbox
        .submit(start_chat, [textbox, chatbot], [textbox, tabs, chatbot], queue=False, api_name="start_chat_textbox")
        .then(chat, [textbox, chatbot], [chatbot, sources_textbox], concurrency_limit=8, api_name="chat_textbox")
        .then(finish_chat, None, [textbox], api_name="finish_chat_textbox")
     )



    def change_sample_questions(key):
        index = list(QUESTIONS.keys()).index(key)
        visible_bools = [False] * len(samples)
        visible_bools[index] = True
        return [gr.update(visible=visible_bools[i]) for i in range(len(samples))]

 #   dropdown_samples.change(change_sample_questions,dropdown_samples,samples)

    query_papers.submit(generate_keywords, [query_papers], [keywords_papers])
    search_papers.click(find_papers, [query_papers, keywords_papers, after], [
                        papers_dataframe, citations_network, papers_summary])

    demo.queue()

demo.launch(allowed_paths=["assets/download.png",
            "assets/logo4.png"], 
            favicon_path="assets/logo4.png")