import os
from typing import Annotated
from typing_extensions import TypedDict

import gradio as gr
from langchain_groq import ChatGroq
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages


# ── State ─────────────────────────────────────────────

class State(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]


# ── Build Graph ───────────────────────────────────────

def build_graph(api_key: str, model: str, system_prompt: str):
    llm = ChatGroq(
        groq_api_key=api_key,
        model_name=model,
        temperature=0.7,
    )

    def chatbot_node(state: State) -> dict:
        full_messages = [SystemMessage(content=system_prompt)] + state["messages"]
        response = llm.invoke(full_messages)
        return {"messages": [response]}

    builder = StateGraph(State)
    builder.add_node("chatbot", chatbot_node)
    builder.add_edge(START, "chatbot")
    builder.add_edge("chatbot", END)
    return builder.compile()


# ── Chat Function ─────────────────────────────────────

def respond(message, history, api_key, model, system_prompt, max_history):
    if not api_key.strip():
        return "⚠️ Please enter your Groq API key."

    if not message.strip():
        return ""

    lc_messages = []

    # Convert history
    for entry in history[-(max_history * 2):]:
        if entry["role"] == "user":
            lc_messages.append(HumanMessage(content=entry["content"]))
        elif entry["role"] == "assistant":
            lc_messages.append(AIMessage(content=entry["content"]))

    lc_messages.append(HumanMessage(content=message))

    graph = build_graph(api_key, model, system_prompt)

    try:
        result = graph.invoke({"messages": lc_messages})
        return result["messages"][-1].content
    except Exception as e:
        return f"❌ Error: {str(e)}"


# ── Constants ─────────────────────────────────────────

MODELS = [
    "llama-3.3-70b-versatile",
    "llama-3.1-8b-instant",
    "mixtral-8x7b-32768",
    "gemma2-9b-it",
]

DEFAULT_SYSTEM = (
    "You are a helpful, friendly, and knowledgeable assistant. "
    "Answer clearly and concisely."
)


# ── UI ────────────────────────────────────────────────

with gr.Blocks(
    title="LangGraph × Groq Assistant",
    theme=gr.themes.Soft(primary_hue="violet"),
) as demo:

    gr.Markdown("# 🤖 LangGraph × Groq Assistant")
    gr.Markdown("Chatbot using LangGraph + Groq")

    chatbot = gr.Chatbot(type="messages")

    with gr.Row():
        msg = gr.Textbox(placeholder="Type your message...", scale=4)
        send = gr.Button("Send ➤")

    clear = gr.Button("🗑️ Clear Chat")

    # Settings panel
    with gr.Accordion("⚙️ Settings", open=False):
        api_key = gr.Textbox(label="Groq API Key", type="password")
        model = gr.Dropdown(choices=MODELS, value=MODELS[0], label="Model")
        system_prompt = gr.Textbox(value=DEFAULT_SYSTEM, label="System Prompt")
        max_history = gr.Slider(1, 20, value=10, step=1, label="Max history")

    # Examples (FIXED ✅)
    examples = [
        ["What is LangGraph?", "", MODELS[0], DEFAULT_SYSTEM, 10],
        ["Explain AI in simple terms", "", MODELS[0], DEFAULT_SYSTEM, 10],
        ["Write Python code for Fibonacci", "", MODELS[0], DEFAULT_SYSTEM, 10],
    ]

    gr.Examples(
        examples=examples,
        inputs=[msg, api_key, model, system_prompt, max_history],
    )

    # Send logic
    def user_input(user_message, chat_history):
        return "", chat_history + [{"role": "user", "content": user_message}]

    def bot_response(chat_history, api_key, model, system_prompt, max_history):
        user_message = chat_history[-1]["content"]
        response = respond(user_message, chat_history[:-1], api_key, model, system_prompt, max_history)
        chat_history.append({"role": "assistant", "content": response})
        return chat_history

    send.click(user_input, [msg, chatbot], [msg, chatbot]) \
        .then(bot_response, [chatbot, api_key, model, system_prompt, max_history], chatbot)

    msg.submit(user_input, [msg, chatbot], [msg, chatbot]) \
        .then(bot_response, [chatbot, api_key, model, system_prompt, max_history], chatbot)

    clear.click(lambda: [], None, chatbot)

demo.launch()