import streamlit as st
import random
import time
import hmac
import os
import json
import requests
from llm_reasoner import LLMReasoner
from prompts import templates
from typing import Any
from string import Template

def safe_parse_json(model_answer):
    """.."""
    try:
        parsed_answer = json.loads(model_answer)
        return parsed_answer
    except json.JSONDecodeError as e:
        print("Failed to parse JSON: %s", e)
        print(f"Raw model answer is: {model_answer}")
        model_answer = model_answer.replace("```json", "")
        model_answer = model_answer.replace("```", "")
    try:
        parsed_answer = json.loads(model_answer)
        print(f"Hit parsing error but cleaned json as markdown code snippet and it can be parsed now")
        return parsed_answer
    except json.JSONDecodeError as e:
        print("Failed to parse JSON: %s", e)
        print(f"Tried to parse cleaned model answer which was: {model_answer}")
        return None

def check_password():
    """Returns `True` if the user had a correct password."""

    def login_form():
        """Form with widgets to collect user information"""
        with st.form("Credentials"):
            st.text_input("Username", key="username")
            st.text_input("Password", type="password", key="password")
            st.form_submit_button("Log in", on_click=password_entered)

    def password_entered():
        """Checks whether a password entered by the user is correct."""

        stored_password = os.getenv(st.session_state["username"])
        if stored_password == st.session_state["password"]:
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # Remove credentials from session
            del st.session_state["username"]
            return

          # If authentication fails
        st.session_state["password_correct"] = False

    # Return True if the username + password is validated.
    if st.session_state.get("password_correct", False):
        return True

    # Show inputs for username + password.
    login_form()
    if "password_correct" in st.session_state:
        st.error("😕 User not known or password incorrect")
    return False

def select_models():
    """Returns only when a valid option is selected from both dropdowns."""

    retriever_options = ["Choose one...", "BM25 Retriever", "Off-the-shelf Retriever", "Finetuned Retriever", "No Retriever"]
    reasoner_options = ["Choose one...", "Claude Sonnet", "GPT-4o", "o3-mini"]

    #selectboxes

    retriever = st.selectbox(
        "Select the Retriever Model",
        retriever_options,
        key="retriever"
    )

    reasoner = st.selectbox(
        "Select the Reasoner Model",
        reasoner_options,
        key="reasoner"
    )

    # next button

    if st.button("Next"):
        # Check that both selections are not the placeholder.
        if retriever == "Choose one..." or reasoner == "Choose one...":
            st.info("Please select both a retriever and a reasoner.")
            return None, None
        else:
            # Store the valid selections in session state
            st.session_state["selected_models"] = (retriever, reasoner)
            return retriever, reasoner
    else:
        st.info("Click 'Next' once you have made your selections.")
        return None, None

def retriever(query: str, selected_retriever: str):
    """Simulate a 'retriever' step, searching for relevant information."""
    with st.chat_message("assistant"):
        placeholder = st.empty()
        text = ""
        
        if selected_retriever == "BM25 Retriever":
            message = "Using the BM25 retriever to search for documents related to your query..."
            retriever_endpoint = "bm25"
        elif selected_retriever == "Off-the-shelf Retriever":
            message = "Using the off-the-shelf retriever to fetch detailed documents relevant to your query..."
            retriever_endpoint = "ots"
        elif selected_retriever == "Finetuned Retriever":
            message = "Using the finetuned retriever to fetch detailed documents relevant to your query..."
            retriever_endpoint = "ft"
        else:
            message = "No retriever selected. Skipping document retrieval."
            retriever_endpoint = "None"

        for chunk in message.split():
            text += chunk + " "
            time.sleep(0.05)
            # Add a blinking cursor to simulate typing
            placeholder.markdown(text + "▌")
        placeholder.markdown(text)

    if retriever_endpoint == "None":
        return ""

    headers = {
        'Content-Type': 'application/json',
    }

    json_data = {
        'claim': query,
    }

    url = "http://18.227.0.215"
    port = "8000"
    response = requests.post(f'{url}:{port}/{retriever_endpoint}', headers=headers, json=json_data)
    documents = response.json()["Documents"]

    k = 3
    topk_documents = documents[:k]

    corpus = '\n\n'.join(topk_documents)

    if "last_expander_id" not in st.session_state:
        st.session_state.last_expander_id = 0

    # Function to create a unique key for each expander
    def get_expander_key():
        st.session_state.last_expander_id += 1
        return f"expander_{st.session_state.last_expander_id}"

    for idx, document in enumerate(topk_documents):
        try:
            a_idx = document.lower().index("abstract:")
            title = document[:a_idx].replace('\n', ' ')
            doc_details = document[a_idx:]
        except:
            print(f"Document parsing error for title and other contents-\n{document}")
            title = "Document title"
            doc_details = document
        try:
            with st.expander(f"Rank {ix+1} - {title}", expanded=False):
                st.markdown(doc_details)
        except:
            print(f"Error in expander {idx} logic")

    # retrieved_document_titles = ""
    # for document in topk_documents:
    #     try:
    #         idx = document.lower().index("abstract:")
    #         title = document[:idx].replace('\n', ' ')
    #         retrieved_document_titles = retrieved_document_titles + title + ' ||| '
    #     except:
    #         retrieved_document_titles = retrieved_document_titles + document + ' ||| '

    # display_retrieved_documents = f"Retrieved Documents - ||| {retrieved_document_titles}"

    # with st.chat_message("assistant"):
    #     placeholder = st.empty()
    #     text = ""

    #     print(display_retrieved_documents)

    #     for chunk in display_retrieved_documents.split():
    #         if chunk == "|||":
    #             text += "  \n" + "  \n"
    #         else:
    #             text += chunk + " "
    #         time.sleep(0.05)
    #         # Add a blinking cursor to simulate typing
    #         placeholder.markdown(text + "▌")
    #     placeholder.markdown(text)

    # You could return retrieved info here.
    return corpus

def reasoner(query: str, documents: list[str], selected_reasoner: str, llm_client: Any):
    """Simulate a 'reasoner' step, thinking about how to answer."""
    with st.chat_message("assistant"):
        placeholder = st.empty()
        text = ""

        if selected_reasoner == "Claude Sonnet":
            message = "Using Claude Sonnet to reason and verify the claim..."
        elif selected_reasoner == "GPT-4o":
            message = "Using GPT-4o to analyze and verify the claim in detail..."
        elif selected_reasoner == "o3-mini":
            message = "Using o3-mini to quickly analyze the claim..."

        if not documents or len(documents) == 0:
            prompt_template = Template(templates["no_evidence"])
            prompt = prompt_template.substitute(claim=query)
            print(prompt)
        else:
            prompt_template = Template(templates["with_evidence"])
            prompt = prompt_template.substitute(claim=query, corpus_text=documents)

        print(prompt)
        
        llm_response = llm_client.run_inference(prompt)

        answer_dict = safe_parse_json(llm_response)
        try:
            decision = answer_dict.get("decision", "")
            reasoning = answer_dict.get("reasoning", "")
        except:
            print(f"Error with parsing the returned {answer_dict}")
            decision, reasoning = "", ""

        for chunk in message.split():
            text += chunk + " "
            time.sleep(0.05)
            # Add a blinking cursor to simulate typing
            placeholder.markdown(text + "▌")
        placeholder.markdown(text)
    # You could return reasoning info here.
    return reasoning, decision

def main():
    st.header(" Scientific Claim Verification ")
    st.caption("Team UMBC-SBU-UT")

    if not check_password():
        st.stop()

    if "selected_models" not in st.session_state:
        selected_retriever, selected_reasoner = select_models()
        # If valid selections are returned, store them and reset the change flag.
        if selected_retriever is not None and selected_reasoner is not None:
            st.session_state.selected_models = (selected_retriever, selected_reasoner)
            st.rerun()
        else:
            st.stop()  # Halt further execution until valid selections are made.
    else:
        selected_retriever, selected_reasoner = st.session_state.selected_models

    # START OF AGENTIC DEMO

    column1, column2 = st.columns(2)
    column1.caption(f"Retriever Selected: {selected_retriever}")
    column2.caption(f"Reasoner Selected: {selected_reasoner}")

    if st.button("Change Selection", key="change_selection_btn"):
        st.session_state.pop("selected_models", None)
        st.session_state.pop("retriever", None)
        st.session_state.pop("reasoner", None)
        st.session_state.messages = [{"role": "assistant", "content": "Let's start verifying the claims here! 👇"}]
        st.rerun()

    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state.messages = [{"role": "assistant", "content": "Let's start verifying the claims here! 👇"}]

    # Display chat messages from history on app rerun
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Model configurations on agentic demo page
    options = {}
    options["max_tokens"] = 500
    options["temperature"] = 0.0

    if selected_reasoner == "Claude Sonnet":
        api_key = os.getenv("claude_key")
        options["model_family"] = "Anthropic"
        options["model_name"] = "claude-3-5-sonnet-20240620"

    elif selected_reasoner == "GPT-4o":
        api_key = os.getenv("openai_key")
        options["model_family"] = "OpenAI"
        options["model_name"] = "gpt-4o-2024-11-20"

    elif selected_reasoner == "o3-mini":
        api_key = os.getenv("openai_key")
        options["model_family"] = "OpenAI"
        options["model_name"] = "o3-mini-2025-01-31"

    options["API_KEY"] = api_key

    llm_client = LLMReasoner(options)

    def send_preset_message(text):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": text})
        # Set the auto_submit flag to process this message
        st.session_state.auto_submit = True
        st.session_state.auto_submit_text = text

    # Quick input buttons - place these above the chat input
    st.sidebar.subheader("Quick Inputs")
    
    # Create buttons in the sidebar
    if st.sidebar.button("Example 1 C"):
        send_preset_message("The density of cytokine receptor bearing cells has no effect on the distance over which cytokines act.")
    
    if st.sidebar.button("Example 2 C"):
        send_preset_message("A total of 1,000 people in the UK are asymptomatic carriers of vCJD infection.")
    
    if st.sidebar.button("Example 3 C"):
        send_preset_message("Sepsis related mortality has risen from 2009 to 2014.")
    
    if st.sidebar.button("Example 4 S"):
        send_preset_message("IL-6 signaling plays a major role in atherosclerotic cardiovascular disease.")
    
    if st.sidebar.button("Example 5 S"):
        send_preset_message("The severity of cardiac involvement in amyloidosis can be described by the degree of transmurality of late gadolinium enhancement in MRI.")

    if st.sidebar.button("Example 6 S"):
        send_preset_message("There was an estimated 30 million cases of pneumonia in young children worldwide in 2010.")

    if st.sidebar.button("New Assessment"):
        st.session_state.messages = [{"role": "assistant", "content": "Let's start verifying the claims here! 👇"}]
        st.rerun()

    # Handle auto-submission if a button was clicked
    if "auto_submit" in st.session_state and st.session_state.auto_submit:
        prompt = st.session_state.auto_submit_text

        display_message = prompt + "  \n" + "  \n" + f"Retriever: {selected_retriever}, Reasoner: {selected_reasoner}"
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.session_state.messages.append({"role": "summary", "content": display_message})

        # Display user message in chat message container
        with st.chat_message("user"):
            st.markdown(display_message)

        retrieved_documents = retriever(prompt, selected_retriever)
        reasoning, decision = reasoner(prompt, retrieved_documents, selected_reasoner, llm_client)

        # Display assistant response in chat message container
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""
            if decision.lower() == 'support':
                assistant_response = f'The claim is CORRECT because {reasoning}'
            elif decision.lower() == 'contradict':
                assistant_response = f'The claim is INCORRECT because {reasoning}'
            else:
                assistant_response = f'Sorry, the query failed due to an issue with connecting to the LLM service.'

            # Simulate stream of response with milliseconds delay
            for chunk in assistant_response.split():
                full_response += chunk + " "
                time.sleep(0.05)
                # Add a blinking cursor to simulate typing
                message_placeholder.markdown(full_response + "▌")
            message_placeholder.markdown(full_response)
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": full_response})

        # Reset the auto_submit flag
        st.session_state.auto_submit = False
        # Rerun to update the UI immediately - may not be needed
        # st.rerun()

    # Accept user input
    prompt = st.chat_input("Type your claim here")
    if prompt:
        # Add user message to chat history
        display_message = prompt + "  \n"+ "  \n"+ f"Retriever: {selected_retriever}, Reasoner: {selected_reasoner}"
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.session_state.messages.append({"role": "summary", "content": display_message})
        # Display user message in chat message container
        with st.chat_message("user"):
            st.markdown(display_message)

        retrieved_documents = retriever(prompt, selected_retriever)
        reasoning, decision = reasoner(prompt, retrieved_documents, selected_reasoner, llm_client)

        # Display assistant response in chat message container
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""
            if decision.lower() == 'support':
                assistant_response = f'The claim is CORRECT because {reasoning}'
            elif decision.lower() == 'contradict':
                assistant_response = f'The claim is INCORRECT because {reasoning}'
            else:
                assistant_response = f'Sorry, the query failed due to an issue with connecting to the LLM service.'

            # Simulate stream of response with milliseconds delay
            for chunk in assistant_response.split():
                full_response += chunk + " "
                time.sleep(0.05)
                # Add a blinking cursor to simulate typing
                message_placeholder.markdown(full_response + "▌")
            message_placeholder.markdown(full_response)
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": full_response})

if __name__ == '__main__':
    main()