from typing import TypedDict, List, Union
from langgraph.graph import StateGraph, START, END
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage, AIMessage
from langchain_groq import ChatGroq
from dotenv import load_dotenv, find_dotenv
import gradio as gr
import os

# Load environment variables
load_dotenv(find_dotenv())

# Define Agent State
class AgentState(TypedDict):
    messages: List[Union[HumanMessage, AIMessage]]

# Initialize Gemini
# llm = ChatGoogleGenerativeAI(
#     model="models/gemini-1.5-flash-latest",
#     temperature=1
# )
llm = ChatGroq(
    model="llama-3.1-8b-instant",
    temperature=0.6,
    max_tokens=64,
    # api_key=groq_api_key
)

# LangGraph processing node
def process(state: AgentState) -> AgentState:
    response = llm.invoke(state["messages"])
    state["messages"].append(AIMessage(content=response.content))
    return state

# Create the graph
graph = StateGraph(AgentState)
graph.add_node("process", process)
graph.add_edge(START, "process")
graph.add_edge("process", END)
agent = graph.compile()

# Global message store
conversation_history = []

# Chat handler
def chat_handler(user_input, history):
    global conversation_history
    conversation_history.append(HumanMessage(content=user_input))

    # Trim old HumanMessages if more than 10
    human_msgs = [m for m in conversation_history if isinstance(m, HumanMessage)]
    if len(human_msgs) > 10:
        for i, m in enumerate(conversation_history):
            if isinstance(m, HumanMessage):
                del conversation_history[i]
                break

    result = agent.invoke({"messages": conversation_history})
    conversation_history = result["messages"]
    reply = conversation_history[-1].content

    history = history or []
    history.append({"role": "user", "content": user_input})
    history.append({"role": "assistant", "content": reply})
    return history, history

# Save chat log
def save_chat_log():
    filename = "logging.txt"
    with open(filename, "w") as file:
        file.write("Your Conversation Log:\n\n")
        for msg in conversation_history:
            if isinstance(msg, HumanMessage):
                file.write(f"You: {msg.content}\n")
            elif isinstance(msg, AIMessage):
                file.write(f"AI: {msg.content}\n\n")
        file.write("End of Conversation\n")
    return gr.File.update(value=filename, visible=True)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align: center;'>LangGraph Chatbot</h1>")

    chatbot = gr.Chatbot(label="Chat", type="messages")
    msg = gr.Textbox(placeholder="Type your message...", label="Your Message")

    with gr.Row():
        send_btn = gr.Button("Send")
        save_btn = gr.Button("Download Chat")
        file_output = gr.File(label="Click to download logging.txt", visible=False)

    state = gr.State([])

    send_btn.click(chat_handler, inputs=[msg, state], outputs=[chatbot, state])
    msg.submit(chat_handler, inputs=[msg, state], outputs=[chatbot, state])
    save_btn.click(save_chat_log, outputs=file_output)

if __name__ == "__main__":
    demo.launch()