#!/usr/bin/env python3

"""

Talk2Scholars: A Streamlit app for the Talk2Scholars graph.

"""

import os

import random

import sys

import hydra

import pandas as pd

import streamlit as st

from langchain.callbacks.tracers import LangChainTracer

from langchain_core.messages import AIMessage, ChatMessage, HumanMessage, SystemMessage

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

from langchain_core.tracers.context import collect_runs

from langchain_openai import ChatOpenAI

from langsmith import Client

from langchain.callbacks.tracers import LangChainTracer

from streamlit_feedback import streamlit_feedback

from utils import streamlit_utils

sys.path.append("./")

# import get_app from main_agent

from aiagents4pharma.talk2scholars.agents.main_agent import get_app

# 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/talk2scholars/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.set_page_config(

    page_title=cfg.page.title, page_icon=cfg.page.icon, layout=cfg.page.layout

)

# Set the logo

st.logo(

    image="docs/assets/VPE.png",

    size="large",

    link="https://github.com/VirtualPatientEngine",

)

# Check if env variables OPENAI_API_KEY and/or NVIDIA_API_KEY exist

if cfg.api_keys.openai_key not in os.environ:

    st.error(

        "Please set the OPENAI_API_KEY "

        "environment variables in the terminal where you run "

        "the app. For more information, please refer to our "

        "[documentation](https://virtualpatientengine.github.io/AIAgents4Pharma/#option-2-git)."

    )

    st.stop()

# Create a chat prompt template

prompt = ChatPromptTemplate.from_messages(

    [

        ("system", "Welcome to Talk2Scholars!"),

        MessagesPlaceholder(variable_name="chat_history", optional=True),

        ("human", "{input}"),

        ("placeholder", "{agent_scratchpad}"),

    ]

)

# Initialize chat history

if "messages" not in st.session_state:

    st.session_state.messages = []

# 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 = "Talk2Scholars-" + 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)

if "app" not in st.session_state:

    if "llm_model" not in st.session_state:

        st.session_state.app = get_app(

            st.session_state.unique_id,

            llm_model=ChatOpenAI(model="gpt-4o-mini", temperature=0),

        )

    else:

        print(st.session_state.llm_model)

        st.session_state.app = get_app(

            st.session_state.unique_id,

            llm_model=streamlit_utils.get_base_chat_model(st.session_state.llm_model),

        )

# Get the app

app = st.session_state.app

def _submit_feedback(user_response):

    """

    Function to submit feedback to the developers.

    """

    client = Client()

    client.create_feedback(

        st.session_state.run_id,

        key="feedback",

        score=1 if user_response["score"] == "👍" else 0,

        comment=user_response["text"],

    )

    st.info("Your feedback is on its way to the developers. Thank you!", icon="🚀")

@st.fragment

def process_pdf_upload():

    """

    Process the uploaded PDF file automatically:

    Read the file as binary and store it in session state under "pdf_data".

    """

    pdf_file = st.file_uploader(

        "Upload an article",

        help="Upload an article in PDF format.",

        type=["pdf"],

        key="pdf_upload",

    )

    if pdf_file:

        import tempfile

        # print (pdf_file.name)

        with tempfile.NamedTemporaryFile(delete=False) as f:

            f.write(pdf_file.read())

            # print (f.name)

        st.session_state.pdf_data = {

            "pdf_object": f.name,  # binary formatted PDF

            "pdf_url": f.name,  # placeholder for URL if needed later

            "arxiv_id": None,  # placeholder for an arXiv id if applicable

        }

        # Create config for the agent

        config = {"configurable": {"thread_id": st.session_state.unique_id}}

        # Update the agent state with the selected LLM model

        app.update_state(config, {"pdf_data": st.session_state.pdf_data})

# 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;'>

            🤖 Talk2Scholars

            </h3>

            """,

            unsafe_allow_html=True,

        )

        # LLM model panel

        st.selectbox(

            "Pick an LLM to power the agent",

            list(cfg.llms.available_models),

            index=0,

            key="llm_model",

            on_change=streamlit_utils.update_llm_model,

            help="Used for tool calling and generating responses.",

        )

        # Text embedding model panel

        text_models = [

            "OpenAI/text-embedding-ada-002",

            "NVIDIA/llama-3.2-nv-embedqa-1b-v2",

        ]

        st.selectbox(

            "Pick a text embedding model",

            text_models,

            index=0,

            key="text_embedding_model",

            on_change=streamlit_utils.update_text_embedding_model,

            kwargs={"app": app},

            help="Used for Retrival Augmented Generation (RAG)",

        )

        # Upload files (placeholder)

        process_pdf_upload()

    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=775):

        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"] == "button":

                if st.button(message["content"], key=message["key"]):

                    # Trigger the question

                    prompt = message["question"]

                    st.empty()

            elif message["type"] == "dataframe":

                if "tool_name" in message:

                    if message["tool_name"] in [

                        "display_results",

                    ]:

                        df_papers = message["content"]

                        st.dataframe(

                            df_papers,

                            use_container_width=True,

                            key=message["key"],

                            hide_index=True,

                            column_config={

                                "URL": st.column_config.LinkColumn(

                                    display_text="Open",

                                ),

                            },

                        )

                # else:

                #     streamlit_utils.render_table(message["content"],

                #                     key=message["key"],

                #                     # tool_name=message["tool_name"],

                #                     save_table=False)

                st.empty()

        # Display intro message only the first time

        # i.e. when there are no messages in the chat

        if not st.session_state.messages:

            with st.chat_message("assistant", avatar="🤖"):

                with st.spinner("Initializing the agent ..."):

                    config = {"configurable": {"thread_id": st.session_state.unique_id}}

                    # Update the agent state with the selected LLM model

                    current_state = app.get_state(config)

                    app.update_state(

                        config,

                        {

                            "llm_model": streamlit_utils.get_base_chat_model(

                                st.session_state.llm_model

                            )

                        },

                    )

                    intro_prompt = "Greet and tell your name and about yourself."

                    intro_prompt += " Also, tell about the agents you can access and ther short description."

                    intro_prompt += " We have provided starter questions (separately) outisde your response."

                    intro_prompt += " Do not provide any questions by yourself. Let the users know that they can"

                    intro_prompt += " simply click on the questions to execute them."

                    # intro_prompt += " Let them know that they can check out the use cases"

                    # intro_prompt += " and FAQs described in the link below. Be friendly and helpful."

                    # intro_prompt += "\n"

                    # intro_prompt += "Here is the link to the use cases: [Use Cases](https://virtualpatientengine.github.io/AIAgents4Pharma/talk2biomodels/cases/Case_1/)"

                    # intro_prompt += "\n"

                    # intro_prompt += "Here is the link to the FAQs: [FAQs](https://virtualpatientengine.github.io/AIAgents4Pharma/talk2biomodels/faq/)"

                    response = app.stream(

                        {"messages": [HumanMessage(content=intro_prompt)]},

                        config=config,

                        stream_mode="messages",

                    )

                    st.write_stream(streamlit_utils.stream_response(response))

                    current_state = app.get_state(config)

                    # Add response to chat history

                    assistant_msg = ChatMessage(

                        current_state.values["messages"][-1].content, role="assistant"

                    )

                    st.session_state.messages.append(

                        {"type": "message", "content": assistant_msg}

                    )

                    st.empty()

        if len(st.session_state.messages) <= 1:

            for count, question in enumerate(streamlit_utils.sample_questions_t2s()):

                if st.button(

                    f"Q{count+1}. {question}", key=f"sample_question_{count+1}"

                ):

                    # Trigger the question

                    prompt = question

                # Add button click to chat history

                st.session_state.messages.append(

                    {

                        "type": "button",

                        "question": question,

                        "content": f"Q{count+1}. {question}",

                        "key": f"sample_question_{count+1}",

                    }

                )

        # When the user asks a question

        if prompt:

            # Create a key 'uploaded_file' to read the uploaded file

            # if uploaded_file:

            #     st.session_state.article_pdf = uploaded_file.read().decode("utf-8")

            # 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()

            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

                    ]

                    # # Create config for the agent

                    config = {"configurable": {"thread_id": st.session_state.unique_id}}

                    # Update the LLM model

                    app.update_state(

                        config,

                        {

                            "llm_model": streamlit_utils.get_base_chat_model(

                                st.session_state.llm_model

                            ),

                            "text_embedding_model": streamlit_utils.get_text_embedding_model(

                                st.session_state.text_embedding_model

                            ),

                        },

                    )

                    current_state = app.get_state(config)

                    print("PDF_DATA", len(current_state.values["pdf_data"]))

                    streamlit_utils.get_response("T2S", None, app, st, prompt)

                    # # Create config for the agent

                    # config = {"configurable": {"thread_id": st.session_state.unique_id}}

                    # # Update the LLM model

                    # app.update_state(

                    #     config,

                    #     {

                    #         "llm_model": streamlit_utils.get_base_chat_model(

                    #             st.session_state.llm_model

                    #         )

                    #     },

                    # )

                    # # Update the agent state with the selected LLM model

                    # current_state = app.get_state(config)

                    # with collect_runs() as cb:

                    #     # Add Langsmith tracer

                    #     tracer = LangChainTracer(

                    #         project_name=st.session_state.project_name

                    #     )

                    #     # Get response from the agent with Langsmith tracing enabled

                    #     # response = app.invoke(

                    #     #     {"messages": [HumanMessage(content=prompt)]},

                    #     #     config=config | {"callbacks": [tracer]},

                    #     # )

                    #     response = app.stream(

                    #         {"messages": [HumanMessage(content=prompt)]},

                    #         config=config|{"callbacks": [tracer]},

                    #         stream_mode="messages"

                    #     )

                    #     st.write_stream(streamlit_utils.stream_response(response))

                    #     # Assign the traced run ID to session state

                    #     if cb.traced_runs:

                    #         st.session_state.run_id = cb.traced_runs[-1].id

                    # # # Get the latest agent state after the response

                    # # current_state = app.get_state(config)

                    # #

                    # # response = app.invoke(

                    # #     {"messages": [HumanMessage(content=prompt)]},

                    # #     config=config,

                    # # )

                    # current_state = app.get_state(config)

                    # # print (response["messages"])

                    # # Add assistant 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()

                    # 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):

                    #         # print ('AIMessage', msg)

                    #         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

                    #     # print("ToolMessage", msg)

                    #     uniq_msg_id = "_".join(

                    #         [msg.name, msg.tool_call_id, str(st.session_state.run_id)]

                    #     )

                    # if msg.name in ['search_tool',

                    #                 'get_single_paper_recommendations',

                    #                 'get_multi_paper_recommendations']:

                    # if msg.name in ["display_results"]:

                    #     # Display the results of the tool call

                    #     # for msg_artifact in msg.artifact:

                    #     # dic_papers = msg.artifact[msg_artifact]

                    #     dic_papers = msg.artifact

                    #     if not dic_papers:

                    #         continue

                    #     df_papers = pd.DataFrame.from_dict(

                    #         dic_papers, orient="index"

                    #     )

                    #     # Add index as a column "key"

                    #     df_papers["Key"] = df_papers.index

                    #     # Drop index

                    #     df_papers.reset_index(drop=True, inplace=True)

                    #     # Drop colum abstract

                    #     df_papers.drop(columns=["Abstract", "Key"], inplace=True)

                    #     if "Year" in df_papers.columns:

                    #         df_papers["Year"] = df_papers["Year"].apply(

                    #             lambda x: (

                    #                 str(int(x))

                    #                 if pd.notna(x) and str(x).isdigit()

                    #                 else None

                    #             )

                    #         )

                    #     if "Date" in df_papers.columns:

                    #         df_papers["Date"] = df_papers["Date"].apply(

                    #             lambda x: (

                    #                 pd.to_datetime(x, errors="coerce").strftime(

                    #                     "%Y-%m-%d"

                    #                 )

                    #                 if pd.notna(pd.to_datetime(x, errors="coerce"))

                    #                 else None

                    #             )

                    #         )

                    #     st.dataframe(

                    #         df_papers,

                    #         hide_index=True,

                    #         column_config={

                    #             "URL": st.column_config.LinkColumn(

                    #                 display_text="Open",

                    #             ),

                    #         },

                    #     )

                    #     # Add data to the chat history

                    #     st.session_state.messages.append(

                    #         {

                    #             "type": "dataframe",

                    #             "content": df_papers,

                    #             "key": "dataframe_" + uniq_msg_id,

                    #             "tool_name": msg.name,

                    #         }

                    #     )

                    #     st.empty()

        # 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=_submit_feedback,

                key=f"feedback_{st.session_state.run_id}",

            )