import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')))
import pandas as pd
import streamlit as st
import time


from src.rag.pipeline import RAGPipeline
from src.utils.data_v2 import (
    build_filter,
    get_meta,
    load_json,
    load_css,
)
from src.utils.writer import typewriter


st.set_page_config(layout="wide")

EMBEDDING_MODEL = "sentence-transformers/distiluse-base-multilingual-cased-v1"
PROMPT_TEMPLATE = os.path.join("src", "rag", "prompt_template.yaml")


@st.cache_data
def load_css_style(path: str) -> None:
    load_css(path)


@st.cache_data
def get_meta_data() -> pd.DataFrame:
    return pd.read_csv(
        os.path.join("database", "meta_data.csv"), dtype=({"retriever_id": str})
    )


@st.cache_data
def get_df() -> pd.DataFrame:
    return pd.read_csv(
        os.path.join("data", "inc_df.csv"), dtype=({"retriever_id": str})
    )[["retriever_id", "draft_labs", "author", "href", "round"]]


@st.cache_data
def get_authors_taxonomy() -> list[str]:
    taxonomy = load_json(os.path.join("data", "authors_taxonomy.json"))
    countries = []
    members = taxonomy["Members"]
    for key, value in members.items():
        if key == "Countries" or key == "International and Regional State Associations":
            countries.extend(value)
    return countries


@st.cache_data
def get_draft_cat_taxonomy() -> dict[str, list[str]]:
    taxonomy = load_json(os.path.join("data", "draftcat_taxonomy_filter.json"))
    draft_labels = []
    for _, subpart in taxonomy.items():
        for label in subpart:
            draft_labels.append(label)
    return draft_labels


@st.cache_data
def get_example_prompts() -> list[str]:
    return [
        example["question"]
        for example in load_json(os.path.join("data", "example_prompts.json"))
    ]


@st.cache_data
def set_trigger_state_values() -> tuple[bool, bool]:
    trigger_filter = st.session_state.setdefault("trigger", False)
    trigger_ask = st.session_state.setdefault("trigger", False)
    return trigger_filter, trigger_ask


@st.cache_resource
def load_pipeline() -> RAGPipeline:
    return RAGPipeline(
        embedding_model=EMBEDDING_MODEL,
        prompt_template=PROMPT_TEMPLATE,
    )


@st.cache_data
def load_app_init() -> None:
    # Define the title of the app
    st.title("INC Plastic Pollution Country Profile Analysis")

    # add warning emoji and style

    st.markdown(
        """
    <div class="remark">
        <div class="remark-content">
            <p class="remark-text" style="font-size: 20px;"> ⚠️ The app is a beta version that serves as a basis for further development. We are aware that the performance is not yet sufficient and that the data basis is not yet complete. We are grateful for any feedback that contributes to the further development and improvement of the app!</p>
        </div>
    </div>
    """,
        unsafe_allow_html=True,
    )

    st.markdown(
        """
        <a href="mailto:teresa.kroesen@giz.de" class="feedback-link">Send feedback</a>
        """,
        unsafe_allow_html=True,
    )

    # add explanation to the app
    st.markdown(
        """
        <p class="description">
        The app is tailored to enhance the efficiency of finding and accessing information on the UN Plastics Treaty Negotiations. It hosts a comprehensive database of relevant documents submitted by the members available <a href="https://www.unep.org/inc-plastic-pollution"> here</a>, which users can explore through an intuitive chatbot interface as well as simple filtering options.
        The app excels in querying specific information about countries and their positions in  the negotiations, providing targeted and precise answers. However, it can process only up to 8 relevant documents at a time, which may limit responses to more complex inquiries. Filter options by authors and sections of the negotiation draft ensure the accuracy of the answers. Each document found via these filters is also directly accessible via a link, ensuring complete and easy access to the desired information.
        </p>
        """,
        unsafe_allow_html=True,
    )


load_css_style("style/style.css")


# Load the data
df = get_df()
df_transformed = get_meta_data()
countries = get_authors_taxonomy()
draft_labels = get_draft_cat_taxonomy()
example_prompts = get_example_prompts()
trigger_filter, trigger_ask = set_trigger_state_values()

# Load pipeline
pipeline = load_pipeline()

# Load app init
load_app_init()


application_col = st.columns(1)


with application_col[0]:
    st.markdown("""<p class="header"> 1️⃣ Select countries""", unsafe_allow_html=True)
    st.markdown(
        """
        <p class="description">
        Please select the countries of interest. Your selection will refine the database to include documents submitted by these countries or recognized groupings such as Small Developing States, the African States Group, etc. </p>
        """,
        unsafe_allow_html=True,
    )
    selected_authors = st.multiselect(
        label="country",
        options=countries,
        label_visibility="collapsed",
        placeholder="Select country/countries",
    )

    st.write("\n")
    st.write("\n")

    st.markdown(
        """<p class="header">  2️⃣ Select parts of the negotiation draft""",
        unsafe_allow_html=True,
    )
    st.markdown(
        """
        <p class="description">
        Please select the parts of the negotiation draft of interest. The negotiation draft can be accessed <a href="https://www.unep.org/inc-plastic-pollution/session-4/documents"> here</a>. </p>
        """,
        unsafe_allow_html=True,
    )
    selected_draft_cats = st.multiselect(
        label="Subpart",
        options=draft_labels,
        label_visibility="collapsed",
        placeholder="Select draft category/draft categories",
    )

    st.write("\n")
    st.write("\n")

    st.markdown(
        """<p class="header"> 3️⃣ Ask a question or show documents based on selected filters""",
        unsafe_allow_html=True,
    )

    asking, filtering = st.tabs(["Ask a question", "Filter documents"])

    with filtering:
        application_col_filter, output_col_filter = st.columns([1, 1.5])
        # make the buttons text smaller
        with application_col_filter:
            st.markdown(
                """
                <p class="description">
                This filter function allows you to see all documents that match the selected filters. The documents can be accessed via a link. \n
                """,
                unsafe_allow_html=True,
            )
            if st.button("Filter documents"):
                filters, status = build_filter(
                    meta_data=df_transformed,
                    authors_filter=selected_authors,
                    draft_cats_filter=selected_draft_cats,
                )
                if status == "no filters selected":
                    st.info("No filters selected. All documents will be shown.")
                    df_filtered = df[
                        ["author", "href", "draft_labs", "round"]
                    ].sort_values(by="author")
                    trigger_filter = True
                if status == "no results found":
                    st.info(
                        "No documents found for the combination of filters you've chosen. All countries are represented at least once in the data. Remove the draft categories to see all documents for the countries selected or try other draft categories."
                    )
                if status == "success":
                    df_filtered = df[df["retriever_id"].isin(filters["retriever_id"])][
                        ["author", "href", "draft_labs", "round"]
                    ].sort_values(by="author")
                    trigger_filter = True

    with asking:
        application_col_ask, output_col_ask = st.columns([1, 1.5])
        with application_col_ask:
            st.markdown(
                """
                <p class="description"> Ask a question, noting that the database has been restricted by filters and that your question should pertain to the selected data. \n
                """,
                unsafe_allow_html=True,
            )
            if "prompt" not in st.session_state:
                prompt = st.text_area("Enter a question")
            if (
                "prompt" in st.session_state
                and st.session_state.prompt in example_prompts  # noqa: E501
            ):  # noqa: E501
                prompt = st.text_area(
                    "Enter a question", value=st.session_state.prompt
                )  # noqa: E501
            if (
                "prompt" in st.session_state
                and st.session_state.prompt not in example_prompts  # noqa: E501
            ):  # noqa: E501
                del st.session_state["prompt"]
                prompt = st.text_area("Enter a question")

            trigger_ask = st.session_state.setdefault("trigger", False)

            if st.button("Ask"):
                if prompt == "":
                    st.error(
                        "Please enter a question. Reloading the app in few seconds"
                    )
                    time.sleep(3)
                    st.rerun()
                with st.spinner("Filtering data...") as status:
                    filter_selection_transformed, status = build_filter(
                        meta_data=df_transformed,
                        authors_filter=selected_authors,
                        draft_cats_filter=selected_draft_cats,
                    )

                    if status == "no filters selected":
                        st.info(
                            "No filters selcted.This will increase the prcessing time significantly. Please select at least one filter."
                        )
                        # st.error(
                        #     "Selecting a filter is mandatory. We especially recommend to select countries you are interested in. Selecting at least one filter is mandatory, because otherwise the model would have to analyze all available documents which results in inaccurate answers and long processing times. Please select at least one filter."
                        # )
                        # st.stop()

                    documents = pipeline.document_store.get_all_documents(
                        filters=filter_selection_transformed
                    )

                    st.success("Filtering data completed.")
                with st.spinner("Answering question...") as status:
                    if filter_selection_transformed == {}:
                        st.warning(
                            "The combination of filters you've chosen does not match any documents. Giving answer based on all documents. Please note that the answer might not be accurate. We highly recommend to use a combination of filters that match the data. All countries are represented at least once in the data. Thus, for example, you could remove the draft categories to match the documents. Or you could check with the Filter documents function which documents are available for the selected countries by removing the draft categories and filter the documents."
                        )

                    result = pipeline.run(
                        prompt=prompt, filters=filter_selection_transformed
                    )
                    trigger_ask = True
                    st.success("Answering question completed.")

            st.markdown("### Examples")
            for i, prompt in enumerate(example_prompts):
                # with col[i % 4]:
                if st.button(prompt):
                    if "key" not in st.session_state:
                        st.session_state["prompt"] = prompt
            st.markdown(
                """
                <ul class="description" style="font-size: 20px;">
                    <li style="font-size: 17px;">These are example prompts that can be used to ask questions to the model</li>
                    <li style="font-size: 17px;">Click on a prompt to use it as a question. You can also type your own question in the text area above.</li>
                    <li style="font-size: 17px;">For questions like "How do country a, b and c [...]" please make sure to select the countries in the filter section. Otherwise the answer will not be accurate. In general we highly recommend to use the filter functions to narrow down the data.</li>
                </ul>
                """,
                unsafe_allow_html=True,
            )

            # for i, prompt in enumerate(example_prompts):
            #     # with col[i % 4]:
            #     if st.button(prompt):
            #         if "key" not in st.session_state:
            #             st.session_state["prompt"] = prompt
            # Define the button

    if trigger_ask:
        with output_col_ask:
            if result is None:
                st.error(
                    "Open AI rate limit exceeded. Please try again in a few seconds."
                )
                st.stop()
            meta_data = get_meta(result=result)
            answer = result["answers"][0].answer

            meta_data_cleaned = []
            seen_retriever_ids = set()

            for data in meta_data:
                retriever_id = data["retriever_id"]
                content = data["content"]
                if retriever_id not in seen_retriever_ids:
                    meta_data_cleaned.append(
                        {
                            "retriever_id": retriever_id,
                            "href": data["href"],
                            "content": [content],
                        }
                    )
                    seen_retriever_ids.add(retriever_id)
                else:
                    for i, item in enumerate(meta_data_cleaned):
                        if item["retriever_id"] == retriever_id:
                            meta_data_cleaned[i]["content"].append(content)

            references = ["\n"]
            for data in meta_data_cleaned:
                retriever_id = data["retriever_id"]
                href = data["href"]
                references.append(f"-[{retriever_id}]: {href} \n")
            st.write("#### 📌 Answer")
            typewriter(
                text=answer,
                references=references,
                speed=100,
            )

            with st.expander("Show more information to the documents"):
                for data in meta_data_cleaned:
                    markdown_text = f"- Document: {data['retriever_id']}\n"
                    markdown_text += "    - Text passages\n"
                    for content in data["content"]:
                        content = (
                            content.replace("[", "").replace("]", "").replace("'", "")
                        )
                        content = " ".join(content.split())
                        markdown_text += f"        - {content}\n"
                    st.write(markdown_text)

        trigger = 0

    if trigger_filter:
        with output_col_filter:
            st.markdown("### Overview of all filtered documents")
            st.dataframe(
                df_filtered,
                hide_index=True,
                column_config={
                    "author": st.column_config.ListColumn("Authors"),
                    "href": st.column_config.LinkColumn("Link to Document"),
                    "draft_labs": st.column_config.ListColumn("Draft Categories"),
                    "round": st.column_config.NumberColumn("Round"),
                },
            )