import os
import re
from dotenv import load_dotenv
load_dotenv()

from langchain.agents.openai_assistant import OpenAIAssistantRunnable
from langchain.schema import HumanMessage, AIMessage

import gradio as gr

# Load API key and assistant IDs
api_key = os.getenv('OPENAI_API_KEY')
extractor_agents = {
    "Solution Specifier A": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_A'),
    "Solution Specifier B": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_B'),
    "Solution Specifier C": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_C'),
    "Solution Specifier D": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_D'),
}

# Function to create a new extractor LLM instance
def get_extractor_llm(agent_id):
    return OpenAIAssistantRunnable(assistant_id=agent_id, api_key=api_key, as_agent=True)

# Utility function to remove citations
def remove_citation(text):
    # Define the regex pattern to match the citation format 【number†text】
    pattern = r"【\d+†\w+】"
    # Replace the pattern with an empty string
    return re.sub(pattern, "📚", text)

# Prediction function
def predict(message, history, selected_agent):
    # Get the extractor LLM for the selected agent
    agent_id = extractor_agents[selected_agent]
    extractor_llm = get_extractor_llm(agent_id)
    
    # Prepare the chat history
    history_langchain_format = []
    for human, ai in history:
        history_langchain_format.append(HumanMessage(content=human))
        history_langchain_format.append(AIMessage(content=ai))
    history_langchain_format.append(HumanMessage(content=message))
    
    # Get the response
    gpt_response = extractor_llm.invoke({"content": message})
    output = gpt_response.return_values["output"]
    non_cited_output = remove_citation(output)
    return non_cited_output

# Define the Gradio interface
def app_interface():
    dropdown = gr.Dropdown(choices=list(extractor_agents.keys()), value="Solution Specifier A", label="Choose Extractor Agent")
    chat = gr.ChatInterface(
        fn=lambda message, history, selected_agent: predict(message, history, selected_agent),
        inputs=[dropdown],
        title="Solution Specifier Chat",
        description="Test with different solution specifiers"
    )
    return chat

# Launch the app
chat_interface = app_interface()
chat_interface.launch(share=True)