import streamlit as st
import os
import pandas as pd
from command_center import CommandCenter
from process_documents import process_documents
from embed_documents import create_retriever
import json
from langchain.callbacks import get_openai_callback
from langchain.chains import ConversationalRetrievalChain
from langchain_openai import ChatOpenAI
import base64

st.set_page_config(layout="wide")
os.environ["OPENAI_API_KEY"] = "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"
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "ls__aca2f2f97d2f4b9caef0ef75c3c33f9d"

get_references = lambda relevant_docs: " ".join(
    [f"[{ref}]" for ref in sorted([ref.metadata["chunk_id"] for ref in relevant_docs])]
)
session_state_2_llm_chat_history = lambda session_state: [
    ss[:2] for ss in session_state if not ss[0].startswith("/")
]
ai_message_format = lambda message, references: (
    f"{message}\n\n---\n\n{references}" if references != "" else message
)
welcome_message = """
Hi I'm Agent Zeta, your AI assistant, dedicated to making your journey through machine learning research papers as insightful and interactive as possible. Whether you're diving into the latest studies or brushing up on foundational papers, I'm here to help navigate, discuss, and analyze content with you.

Here's a quick guide to getting started with me:

| Command | Description |
|---------|-------------|
| `/upload` | Upload and process documents for our conversation. |
| `/index` | View an index of processed documents to easily navigate your research. |
| `/cost` | Calculate the cost of our conversation, ensuring transparency in resource usage. |
| `/download` | Download conversation data for your records or further analysis. |

<br>

Feel free to use these commands to enhance your research experience. Let's embark on this exciting journey of discovery together!

Use `/man` at any point of time to view this guide again.
"""


def process_documents_wrapper(inputs):
    snippets = process_documents(inputs)
    st.session_state.retriever = create_retriever(snippets)
    st.session_state.source_doc_urls = inputs
    st.session_state.index = [
        [snip.metadata["chunk_id"], snip.metadata["header"]] for snip in snippets
    ]
    response = f"Uploaded and processed documents {inputs}"
    st.session_state.messages.append((f"/upload {inputs}", response, ""))
    return response


def index_documents_wrapper(inputs=None):
    response = pd.DataFrame(
        st.session_state.index, columns=["id", "reference"]
    ).to_markdown()
    st.session_state.messages.append(("/index", response, ""))
    return response


def calculate_cost_wrapper(inputs=None):
    try:
        stats_df = pd.DataFrame(st.session_state.costing)
        stats_df.loc["total"] = stats_df.sum()
        response = stats_df.to_markdown()
    except ValueError:
        response = "No cost incurred yet"
    st.session_state.messages.append(("/cost", response, ""))
    return response


def download_conversation_wrapper(inputs=None):
    conversation_data = json.dumps(
        {
            "document_urls": (
                st.session_state.source_doc_urls
                if "source_doc_urls" in st.session_state
                else []
            ),
            "document_snippets": (
                st.session_state.index if "index" in st.session_state else []
            ),
            "conversation": [
                {"human": message[0], "ai": message[1], "references": message[2]}
                for message in st.session_state.messages
            ],
            "costing": (
                st.session_state.costing if "costing" in st.session_state else []
            ),
            "total_cost": (
                {
                    k: sum(d[k] for d in st.session_state.costing)
                    for k in st.session_state.costing[0]
                }
                if "costing" in st.session_state and len(st.session_state.costing) > 0
                else {}
            ),
        }
    )
    conversation_data = base64.b64encode(conversation_data.encode()).decode()
    st.session_state.messages.append(("/download", "Conversation data downloaded", ""))
    return f'<a href="data:text/csv;base64,{conversation_data}" download="conversation_data.json">Download Conversation</a>'


def query_llm_wrapper(inputs):
    retriever = st.session_state.retriever
    qa_chain = ConversationalRetrievalChain.from_llm(
        llm=ChatOpenAI(model="gpt-4-0125-preview", temperature=0),
        retriever=retriever,
        return_source_documents=True,
        chain_type="stuff",
    )
    relevant_docs = retriever.get_relevant_documents(inputs)
    with get_openai_callback() as cb:
        result = qa_chain(
            {
                "question": inputs,
                "chat_history": session_state_2_llm_chat_history(
                    st.session_state.messages
                ),
            }
        )
        stats = cb
    result = result["answer"]
    references = get_references(relevant_docs)
    st.session_state.messages.append((inputs, result, references))
    st.session_state.costing.append(
        {
            "prompt tokens": stats.prompt_tokens,
            "completion tokens": stats.completion_tokens,
            "cost": stats.total_cost,
        }
    )
    return result, references


def boot(command_center):
    st.write("# Agent Zeta")
    if "costing" not in st.session_state:
        st.session_state.costing = []
    if "messages" not in st.session_state:
        st.session_state.messages = []
    st.chat_message("ai").write(welcome_message, unsafe_allow_html=True)
    for message in st.session_state.messages:
        st.chat_message("human").write(message[0])
        st.chat_message("ai").write(
            ai_message_format(message[1], message[2]), unsafe_allow_html=True
        )
    if query := st.chat_input():
        st.chat_message("human").write(query)
        response = command_center.execute_command(query)
        if response is None:
            pass
        elif type(response) == tuple:
            result, references = response
            st.chat_message("ai").write(
                ai_message_format(result, references), unsafe_allow_html=True
            )
        else:
            st.chat_message("ai").write(response, unsafe_allow_html=True)


if __name__ == "__main__":
    all_commands = [
        ("/upload", list, process_documents_wrapper),
        ("/index", None, index_documents_wrapper),
        ("/cost", None, calculate_cost_wrapper),
        ("/download", None, download_conversation_wrapper),
        ("/man", None, lambda x: welcome_message),
    ]
    command_center = CommandCenter(
        default_input_type=str,
        default_function=query_llm_wrapper,
        all_commands=all_commands,
    )
    boot(command_center)