#!/usr/bin/env python3

"""
Talk2KnowledgeGraphs: A Streamlit app for the Talk2KnowledgeGraphs graph.
"""

import os
import sys
import random
import streamlit as st
import pandas as pd
import hydra
from streamlit_feedback import streamlit_feedback
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_core.messages import ChatMessage
from langchain_core.tracers.context import collect_runs
from langchain.callbacks.tracers import LangChainTracer
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_ollama import OllamaEmbeddings, ChatOllama
from utils import streamlit_utils

sys.path.append("./")
from aiagents4pharma.talk2knowledgegraphs.agents.t2kg_agent import get_app
# from talk2knowledgegraphs.agents.t2kg_agent import get_app

st.set_page_config(
    page_title="Talk2KnowledgeGraphs",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="collapsed",
)

# Initialize configuration
hydra.core.global_hydra.GlobalHydra.instance().clear()
if "config" not in st.session_state:
    # Load Hydra configuration
    with hydra.initialize(
        version_base=None,
        config_path="../../aiagents4pharma/talk2knowledgegraphs/configs",
    ):
        cfg = hydra.compose(config_name="config", overrides=["app/frontend=default"])
        cfg = cfg.app.frontend
        st.session_state.config = cfg
else:
    cfg = st.session_state.config


# st.logo(
#     image='docs/VPE.png',
#     size='large',
#     link='https://github.com/VirtualPatientEngine'
# )

# Check if env variable OPENAI_API_KEY exists
if "OPENAI_API_KEY" not in os.environ:
    st.error(
        "Please set the OPENAI_API_KEY environment \
        variable in the terminal where you run the app."
    )
    st.stop()

# Initialize current user
if "current_user" not in st.session_state:
    st.session_state.current_user = cfg.default_user

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Initialize session state for SBML file uploader
# if "sbml_key" not in st.session_state:
#     st.session_state.sbml_key = 0

# Initialize session state for pre-clinical data package uploader
if "data_package_key" not in st.session_state:
    st.session_state.data_package_key = 0

# Initialize session state for patient gene expression data uploader
if "endotype_key" not in st.session_state:
    st.session_state.endotype_key = 0

# Initialize session state for uploaded files
if "uploaded_files" not in st.session_state:
    st.session_state.uploaded_files = []

    # Make directories if not exists
    os.makedirs(cfg.upload_data_dir, exist_ok=True)

# Initialize project_name for Langsmith
if "project_name" not in st.session_state:
    # st.session_state.project_name = str(st.session_state.user_name) + '@' + str(uuid.uuid4())
    st.session_state.project_name = "T2KG-" + str(random.randint(1000, 9999))

# Initialize run_id for Langsmith
if "run_id" not in st.session_state:
    st.session_state.run_id = None

# Initialize graph
if "unique_id" not in st.session_state:
    st.session_state.unique_id = random.randint(1, 1000)

# Initialize the LLM model
if "llm_model" not in st.session_state:
    st.session_state.llm_model = tuple(cfg.openai_llms + cfg.ollama_llms)[0]

# Initialize the app with default LLM model for the first time
if "app" not in st.session_state:
    # Initialize the app
    if st.session_state.llm_model in cfg.openai_llms:
        print("Using OpenAI LLM model")
        st.session_state.app = get_app(
            st.session_state.unique_id,
            llm_model=ChatOpenAI(
                model=st.session_state.llm_model, temperature=cfg.temperature
            ),
        )
    else:
        print("Using Ollama LLM model")
        st.session_state.app = get_app(
            st.session_state.unique_id,
            llm_model=ChatOllama(
                model=st.session_state.llm_model, temperature=cfg.temperature
            ),
        )

if "topk_nodes" not in st.session_state:
    # Subgraph extraction settings
    st.session_state.topk_nodes = cfg.reasoning_subgraph_topk_nodes
    st.session_state.topk_edges = cfg.reasoning_subgraph_topk_edges

# Get the app
app = st.session_state.app

# Apply custom CSS
streamlit_utils.apply_css()

# Sidebar
with st.sidebar:
    st.markdown("**⚙️ Subgraph Extraction Settings**")
    topk_nodes = st.slider(
        "Top-K (Nodes)",
        cfg.reasoning_subgraph_topk_nodes_min,
        cfg.reasoning_subgraph_topk_nodes_max,
        st.session_state.topk_nodes,
        key="st_slider_topk_nodes",
    )
    st.session_state.topk_nodes = topk_nodes
    topk_edges = st.slider(
        "Top-K (Edges)",
        cfg.reasoning_subgraph_topk_nodes_min,
        cfg.reasoning_subgraph_topk_nodes_max,
        st.session_state.topk_edges,
        key="st_slider_topk_edges",
    )
    st.session_state.topk_edges = topk_edges

# Main layout of the app split into two columns
main_col1, main_col2 = st.columns([3, 7])
# First column
with main_col1:
    with st.container(border=True):
        # Title
        st.write(
            """
            <h3 style='margin: 0px; padding-bottom: 10px; font-weight: bold;'>
            🤖 Talk2KnowledgeGraphs
            </h3>
            """,
            unsafe_allow_html=True,
        )

        # LLM panel (Only at the front-end for now)
        # llms = ["gpt-4o-mini", "gpt-4-turbo", "gpt-3.5-turbo"]
        llms = tuple(cfg.openai_llms + cfg.ollama_llms)
        st.selectbox(
            "Pick an LLM to power the agent",
            llms,
            index=0,
            key="llm_model",
            on_change=streamlit_utils.update_llm_model,
        )

        # Upload files
        streamlit_utils.get_uploaded_files(cfg)

        # Help text
        # st.button("Know more ↗",
        #         #   icon="ℹ️",
        #           on_click=streamlit_utils.help_button,
        #           use_container_width=False)

    with st.container(border=False, height=500):
        prompt = st.chat_input("Say something ...", key="st_chat_input")

# Second column
with main_col2:
    # Chat history panel
    with st.container(border=True, height=575):
        st.write("#### 💬 Chat History")

        # Display chat messages
        for count, message in enumerate(st.session_state.messages):
            if message["type"] == "message":
                with st.chat_message(
                    message["content"].role,
                    avatar="🤖" if message["content"].role != "user" else "👩🏻‍💻",
                ):
                    st.markdown(message["content"].content)
                    st.empty()
            elif message["type"] == "plotly":
                streamlit_utils.render_plotly(
                    message["content"],
                    key=message["key"],
                    title=message["title"],
                    #   tool_name=message["tool_name"],
                    save_chart=False,
                )
                st.empty()
            elif message["type"] == "toggle":
                streamlit_utils.render_toggle(
                    key=message["key"],
                    toggle_text=message["content"],
                    toggle_state=message["toggle_state"],
                    save_toggle=False,
                )
                st.empty()
            elif message["type"] == "dataframe":
                streamlit_utils.render_table(
                    message["content"],
                    key=message["key"],
                    # tool_name=message["tool_name"],
                    save_table=False,
                )
                st.empty()
            elif message["type"] == "graph":
                streamlit_utils.render_graph(
                    message["content"], key=message["key"], save_graph=False
                )
                st.empty()

        # When the user asks a question
        if prompt:
            # Display user prompt
            prompt_msg = ChatMessage(prompt, role="user")
            st.session_state.messages.append({"type": "message", "content": prompt_msg})
            with st.chat_message("user", avatar="👩🏻‍💻"):
                st.markdown(prompt)
                st.empty()

            # Auxiliary visualization-related variables
            graphs_visuals = []
            with st.chat_message("assistant", avatar="🤖"):
                # with st.spinner("Fetching response ..."):
                with st.spinner():
                    # Get chat history
                    history = [
                        (m["content"].role, m["content"].content)
                        for m in st.session_state.messages
                        if m["type"] == "message"
                    ]
                    # Convert chat history to ChatMessage objects
                    chat_history = [
                        SystemMessage(content=m[1])
                        if m[0] == "system"
                        else HumanMessage(content=m[1])
                        if m[0] == "human"
                        else AIMessage(content=m[1])
                        for m in history
                    ]

                    # Prepare LLM and embedding model for updating the agent
                    if st.session_state.llm_model in cfg.openai_llms:
                        llm_model = ChatOpenAI(
                            model=st.session_state.llm_model,
                            temperature=cfg.temperature,
                        )
                        emb_model = OpenAIEmbeddings(model=cfg.openai_embeddings[0])
                    else:
                        llm_model = ChatOllama(
                            model=st.session_state.llm_model,
                            temperature=cfg.temperature,
                        )
                        emb_model = OllamaEmbeddings(model=cfg.ollama_embeddings[0])

                    # Create config for the agent
                    config = {"configurable": {"thread_id": st.session_state.unique_id}}
                    app.update_state(
                        config,
                        {
                            "llm_model": llm_model,
                            "embedding_model": emb_model,
                            "uploaded_files": st.session_state.uploaded_files,
                            "topk_nodes": st.session_state.topk_nodes,
                            "topk_edges": st.session_state.topk_edges,
                            "dic_source_graph": [
                                {
                                    "name": st.session_state.config["kg_name"],
                                    "kg_pyg_path": st.session_state.config["kg_pyg_path"],
                                    "kg_text_path": st.session_state.config["kg_text_path"],
                                }
                            ],
                        },
                    )

                    # Update the agent states
                    current_state = app.get_state(config)

                    ERROR_FLAG = False
                    with collect_runs() as cb:
                        # Add Langsmith tracer
                        tracer = LangChainTracer(
                            project_name=st.session_state.project_name
                        )
                        # Get response from the agent
                        response = app.invoke(
                            {"messages": [HumanMessage(content=prompt)]},
                            config=config | {"callbacks": [tracer]},
                        )
                        st.session_state.run_id = cb.traced_runs[-1].id
                    current_state = app.get_state(config)

                    # Add response to chat history
                    assistant_msg = ChatMessage(
                        response["messages"][-1].content, role="assistant"
                    )
                    st.session_state.messages.append(
                        {"type": "message", "content": assistant_msg}
                    )
                    # Display the response in the chat
                    st.markdown(response["messages"][-1].content)
                    st.empty()

                    # Get the current state of the graph
                    current_state = app.get_state(config)

                    # # Get the messages from the current state
                    # # and reverse the order
                    reversed_messages = current_state.values["messages"][::-1]

                    # Loop through the reversed messages until a
                    # HumanMessage is found i.e. the last message
                    # from the user. This is to display the results
                    # of the tool calls made by the agent since the
                    # last message from the user.
                    for msg in reversed_messages:
                        # print (msg)
                        # Break the loop if the message is a HumanMessage
                        # i.e. the last message from the user
                        if isinstance(msg, HumanMessage):
                            break
                        # Skip the message if it is an AIMessage
                        # i.e. a message from the agent. An agent
                        # may make multiple tool calls before the
                        # final response to the user.
                        if isinstance(msg, AIMessage):
                            continue
                        # Work on the message if it is a ToolMessage
                        # These may contain additional visuals that
                        # need to be displayed to the user.
                        # print("ToolMessage", msg)
                        # Skip the Tool message if it is an error message
                        if msg.status == "error":
                            continue

                        # Create a unique message id to identify the tool call
                        # msg.name is the name of the tool
                        # msg.tool_call_id is the unique id of the tool call
                        # st.session_state.run_id is the unique id of the run
                        uniq_msg_id = (
                            msg.name
                            + "_"
                            + msg.tool_call_id
                            + "_"
                            + str(st.session_state.run_id)
                        )
                        if msg.name in ["subgraph_extraction"]:
                            print(
                                "-",
                                len(current_state.values["dic_extracted_graph"]),
                                "subgraph_extraction",
                            )
                            # Add the graph into the visuals list
                            latest_graph = current_state.values["dic_extracted_graph"][
                                -1
                            ]
                            if current_state.values["dic_extracted_graph"]:
                                graphs_visuals.append(
                                    {
                                        "content": latest_graph["graph_dict"],
                                        "key": "subgraph_" + uniq_msg_id,
                                    }
                                )

            # Visualize the graph
            if len(graphs_visuals) > 0:
                for count, graph in enumerate(graphs_visuals):
                    streamlit_utils.render_graph(
                        graph_dict=graph["content"], key=graph["key"], save_graph=True
                    )

        # Collect feedback and display the thumbs feedback
        if st.session_state.get("run_id"):
            feedback = streamlit_feedback(
                feedback_type="thumbs",
                optional_text_label="[Optional] Please provide an explanation",
                on_submit=streamlit_utils.submit_feedback,
                key=f"feedback_{st.session_state.run_id}",
            )