# agentefinal_gradio.py
# ---------------------
# Interfaz web con Gradio para tu agente basado en LangGraph + Trustcall.
# Toma como base tu agentefinal.py y expone un Chatbot en localhost.

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from typing import TypedDict, Literal
import uuid
from datetime import datetime
from typing import Optional

from pydantic import BaseModel, Field
from trustcall import create_extractor

from langchain_core.runnables import RunnableConfig
from langchain_core.messages import merge_message_runs, HumanMessage, SystemMessage

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, MessagesState, END, START
from langgraph.store.base import BaseStore
from langgraph.store.memory import InMemoryStore

# --- NUEVO: Gradio
import gradio as gr

# ---------------------------------------------------------------------
# CARGA DE VARIABLES DE ENTORNO (por ejemplo, OPENAI_API_KEY desde .env)
# ---------------------------------------------------------------------
load_dotenv()

# ---------------------------------------------------------------------
# MODELO BASE
# ---------------------------------------------------------------------
# Puedes ajustar el modelo/temperatura si lo necesitas.
model = ChatOpenAI(model="gpt-4.1-mini", temperature=0)

# ---------------------------------------------------------------------
# TOOLS / ESQUEMAS
# ---------------------------------------------------------------------
class UpdateMemory(TypedDict):
    """ Decision on what memory type to update """
    update_type: Literal['user', 'todo', 'instructions']

class Profile(BaseModel):
    """This is the profile of the user you are chatting with"""
    name: Optional[str] = Field(description="The user's name", default=None)
    location: Optional[str] = Field(description="The user's location", default=None)
    job: Optional[str] = Field(description="The user's job", default=None)
    connections: list[str] = Field(
        description="Personal connection of the user, such as family members, friends, or coworkers",
        default_factory=list
    )
    interests: list[str] = Field(
        description="Interests that the user has",
        default_factory=list
    )

class ToDo(BaseModel):
    task: str = Field(description="The task to be completed.")
    time_to_complete: Optional[int] = Field(description="Estimated time to complete the task (minutes).")
    deadline: Optional[datetime] = Field(
        description="When the task needs to be completed by (if applicable)",
        default=None
    )
    solutions: list[str] = Field(
        description="List of specific, actionable solutions (e.g., specific ideas, service providers, or concrete options relevant to completing the task)",
        min_items=1,
        default_factory=list
    )
    status: Literal["not started", "in progress", "done", "archived"] = Field(
        description="Current status of the task",
        default="not started"
    )

# Extractor para perfil
profile_extractor = create_extractor(
    model,
    tools=[Profile],
    tool_choice="Profile",
)

# ---------------------------------------------------------------------
# UTILIDAD PARA INSPECCIONAR LLAMADAS DE HERRAMIENTAS (Trustcall)
# ---------------------------------------------------------------------
class Spy:
    def __init__(self):
        self.called_tools = []

    def __call__(self, run):
        q = [run]
        while q:
            r = q.pop()
            if getattr(r, "child_runs", None):
                q.extend(r.child_runs)
            if getattr(r, "run_type", None) == "chat_model":
                try:
                    self.called_tools.append(
                        r.outputs["generations"][0][0]["message"]["kwargs"]["tool_calls"]
                    )
                except Exception:
                    pass

def extract_tool_info(tool_calls, schema_name="Memory"):
    """Extrae información útil de las tool calls (Trustcall)."""
    changes = []
    for call_group in tool_calls:
        for call in call_group:
            if call.get('name') == 'PatchDoc':
                if call.get('args', {}).get('patches'):
                    changes.append({
                        'type': 'update',
                        'doc_id': call['args'].get('json_doc_id'),
                        'planned_edits': call['args'].get('planned_edits'),
                        'value': call['args']['patches'][0].get('value')
                    })
            elif call.get('name') == schema_name:
                changes.append({'type': 'new', 'value': call.get('args')})

    result_parts = []
    for change in changes:
        if change['type'] == 'update':
            result_parts.append(
                f"Document {change['doc_id']} updated:\n"
                f"Plan: {change['planned_edits']}\n"
                f"Added content: {change['value']}"
            )
        else:
            result_parts.append(
                f"New {schema_name} created:\n"
                f"Content: {change['value']}"
            )
    return "\n\n".join(result_parts)

# ---------------------------------------------------------------------
# PROMPTS DEL AGENTE
# ---------------------------------------------------------------------
MODEL_SYSTEM_MESSAGE = """You are a helpful chatbot.

You are designed to be a companion to a user, helping them keep track of their ToDo list.

You have a long term memory which keeps track of three things:
1. The user's profile (general information about them)
2. The user's ToDo list
3. General instructions for updating the ToDo list

Here is the current User Profile (may be empty if no information has been collected yet):
<user_profile>
{user_profile}
</user_profile>

Here is the current ToDo List (may be empty if no tasks have been added yet):
<todo>
{todo}
</todo>

Here are the current user-specified preferences for updating the ToDo list (may be empty if no preferences have been specified yet):
<instructions>
{instructions}
</instructions>

Here are your instructions for reasoning about the user's messages:

1. Reason carefully about the user's messages as presented below.

2. Decide whether any of the your long-term memory should be updated:
- If personal information was provided about the user, update the user's profile by calling UpdateMemory tool with type `user`
- If tasks are mentioned, update the ToDo list by calling UpdateMemory tool with type `todo`
- If the user has specified preferences for how to update the ToDo list, update the instructions by calling UpdateMemory tool with type `instructions`

3. Tell the user that you have updated your memory, if appropriate:
- Do not tell the user you have updated the user's profile
- Tell the user them when you update the todo list
- Do not tell the user that you have updated instructions

4. Err on the side of updating the todo list. No need to ask for explicit permission.

5. Respond naturally to user user after a tool call was made to save memories, or if no tool call was made."""

TRUSTCALL_INSTRUCTION = """Reflect on following interaction.

Use the provided tools to retain any necessary memories about the user.

Use parallel tool calling to handle updates and insertions simultaneously.

System Time: {time}"""

CREATE_INSTRUCTIONS = """Reflect on the following interaction.

Based on this interaction, update your instructions for how to update ToDo list items.

Use any feedback from the user to update how they like to have items added, etc.

Your current instructions are:

<current_instructions>
{current_instructions}
</current_instructions>"""

# ---------------------------------------------------------------------
# NODOS DEL GRAFO
# ---------------------------------------------------------------------
def task_mAIstro(state: MessagesState, config: RunnableConfig, store: BaseStore):
    """Carga memorias y responde con el modelo, decidiendo si llamar UpdateMemory."""
    user_id = config["configurable"]["user_id"]

    # Profile
    namespace = ("profile", user_id)
    memories = store.search(namespace)
    user_profile = memories[0].value if memories else None

    # ToDo
    namespace = ("todo", user_id)
    memories = store.search(namespace)
    todo = "\n".join(f"{mem.value}" for mem in memories)

    # Instrucciones
    namespace = ("instructions", user_id)
    memories = store.search(namespace)
    instructions = memories[0].value if memories else ""

    system_msg = MODEL_SYSTEM_MESSAGE.format(
        user_profile=user_profile,
        todo=todo,
        instructions=instructions
    )

    response = model.bind_tools([UpdateMemory], parallel_tool_calls=False).invoke(
        [SystemMessage(content=system_msg)] + state["messages"]
    )
    return {"messages": [response]}

def update_profile(state: MessagesState, config: RunnableConfig, store: BaseStore):
    """Actualiza memoria de perfil con Trustcall."""
    user_id = config["configurable"]["user_id"]
    namespace = ("profile", user_id)

    existing_items = store.search(namespace)
    tool_name = "Profile"
    existing_memories = ([(existing_item.key, tool_name, existing_item.value)
                          for existing_item in existing_items] if existing_items else None)

    TRUSTCALL_INSTRUCTION_FORMATTED = TRUSTCALL_INSTRUCTION.format(time=datetime.now().isoformat())
    updated_messages = list(merge_message_runs(
        messages=[SystemMessage(content=TRUSTCALL_INSTRUCTION_FORMATTED)] + state["messages"][:-1]
    ))

    result = profile_extractor.invoke({"messages": updated_messages, "existing": existing_memories})

    for r, rmeta in zip(result["responses"], result["response_metadata"]):
        store.put(namespace, rmeta.get("json_doc_id", str(uuid.uuid4())), r.model_dump(mode="json"))

    tool_calls = state['messages'][-1].tool_calls
    return {"messages": [{"role": "tool", "content": "updated profile", "tool_call_id": tool_calls[0]['id']}]}

def update_todos(state: MessagesState, config: RunnableConfig, store: BaseStore):
    """Actualiza ToDos con Trustcall (inserciones + parches)."""
    user_id = config["configurable"]["user_id"]
    namespace = ("todo", user_id)

    existing_items = store.search(namespace)
    tool_name = "ToDo"
    existing_memories = ([(existing_item.key, tool_name, existing_item.value)
                          for existing_item in existing_items] if existing_items else None)

    TRUSTCALL_INSTRUCTION_FORMATTED = TRUSTCALL_INSTRUCTION.format(time=datetime.now().isoformat())
    updated_messages = list(merge_message_runs(
        messages=[SystemMessage(content=TRUSTCALL_INSTRUCTION_FORMATTED)] + state["messages"][:-1]
    ))

    spy = Spy()
    todo_extractor = create_extractor(
        model,
        tools=[ToDo],
        tool_choice=tool_name,
        enable_inserts=True
    ).with_listeners(on_end=spy)

    result = todo_extractor.invoke({"messages": updated_messages, "existing": existing_memories})

    for r, rmeta in zip(result["responses"], result["response_metadata"]):
        store.put(namespace, rmeta.get("json_doc_id", str(uuid.uuid4())), r.model_dump(mode="json"))

    tool_calls = state['messages'][-1].tool_calls
    todo_update_msg = extract_tool_info(spy.called_tools, tool_name)
    return {"messages": [{"role": "tool", "content": todo_update_msg or "updated todos", "tool_call_id": tool_calls[0]['id']}]}

def update_instructions(state: MessagesState, config: RunnableConfig, store: BaseStore):
    """Actualiza instrucciones personalizadas del usuario."""
    user_id = config["configurable"]["user_id"]
    namespace = ("instructions", user_id)

    existing_memory = store.get(namespace, "user_instructions")
    system_msg = CREATE_INSTRUCTIONS.format(
        current_instructions=existing_memory.value if existing_memory else None
    )

    new_memory = model.invoke(
        [SystemMessage(content=system_msg)] +
        state['messages'][:-1] +
        [HumanMessage(content="Please update the instructions based on the conversation")]
    )

    store.put(namespace, "user_instructions", {"memory": new_memory.content})
    tool_calls = state['messages'][-1].tool_calls
    return {"messages": [{"role": "tool", "content": "updated instructions", "tool_call_id": tool_calls[0]['id']}]}

def route_message(state: MessagesState, config: RunnableConfig, store: BaseStore) -> Literal[END, "update_todos", "update_instructions", "update_profile"]:
    """Decide qué colección actualizar según la tool call del modelo."""
    message = state['messages'][-1]
    if len(getattr(message, "tool_calls", []) or []) == 0:
        return END
    tool_call = message.tool_calls[0]
    ut = tool_call['args']['update_type']
    if ut == "user":
        return "update_profile"
    elif ut == "todo":
        return "update_todos"
    elif ut == "instructions":
        return "update_instructions"
    else:
        raise ValueError("Unknown update_type")

# ---------------------------------------------------------------------
# COMPILACIÓN DEL GRAFO + MEMORIA
# ---------------------------------------------------------------------
def build_graph():
    builder = StateGraph(MessagesState)
    builder.add_node(task_mAIstro)
    builder.add_node(update_todos)
    builder.add_node(update_profile)
    builder.add_node(update_instructions)

    builder.add_edge(START, "task_mAIstro")
    builder.add_conditional_edges("task_mAIstro", route_message)
    builder.add_edge("update_todos", "task_mAIstro")
    builder.add_edge("update_profile", "task_mAIstro")
    builder.add_edge("update_instructions", "task_mAIstro")

    across_thread_memory = InMemoryStore()   # memoria largo plazo (en RAM)
    within_thread_memory = MemorySaver()     # checkpointing corto plazo

    graph = builder.compile(checkpointer=within_thread_memory, store=across_thread_memory)
    return graph, across_thread_memory, within_thread_memory

GRAPH, STORE, CHECKPOINTER = build_graph()

# ---------------------------------------------------------------------
# FUNCIÓN DE CHAT PARA GRADIO
# ---------------------------------------------------------------------
def chat_fn(user_input, history, user_id, thread_id):
    """
    - user_input: texto del usuario
    - history: historial [(user, bot), ...] mostrado en Gradio
    - user_id: id lógico para memoria a largo plazo (e.g., nombre)
    - thread_id: id del hilo para memoria de corto plazo
    """
    # Config para LangGraph
    config = {"configurable": {"thread_id": str(thread_id or "1"), "user_id": str(user_id or "default")}}
    input_messages = [HumanMessage(content=user_input or "")]

    # Ejecutar grafo por streaming y quedarnos con el último mensaje
    response_text = ""
    try:
        for chunk in GRAPH.stream({"messages": input_messages}, config, stream_mode="values"):
            msg = chunk["messages"][-1]
            # msg puede ser un ChatMessage, ToolMessage, etc.
            content = getattr(msg, "content", None)
            if content:
                response_text = content
    except Exception as e:
        response_text = f"Oops, hubo un error procesando tu mensaje: {e}"

    # Actualizamos historial para el componente Chatbot
    history = (history or []) + [(user_input, response_text)]
    return history, history

def clear_fn():
    return [], []

# ---------------------------------------------------------------------
# UI DE GRADIO
# ---------------------------------------------------------------------
def build_ui():
    with gr.Blocks(title="Agente con Memoria (LangGraph + Trustcall)") as demo:
        gr.Markdown("## 🧠 Agente ToDo con memoria (LangGraph + Trustcall) + Gradio")

        with gr.Row():
            user_id = gr.Textbox(label="User ID (memoria largo plazo)", value="Lance")
            thread_id = gr.Textbox(label="Thread ID (memoria corto plazo)", value="1")

        chatbot = gr.Chatbot(label="Chat")
        msg = gr.Textbox(label="Escribe tu mensaje", placeholder="Hola, me llamo... Añade 'reservar clases...' etc.", lines=2)
        with gr.Row():
            send = gr.Button("Enviar", variant="primary")
            clear = gr.Button("Limpiar historial")

        state = gr.State([])  # historial

        # Acciones
        msg.submit(chat_fn, [msg, state, user_id, thread_id], [chatbot, state])
        send.click(chat_fn, [msg, state, user_id, thread_id], [chatbot, state])
        clear.click(lambda: ([], []), None, [chatbot, state])

        gr.Markdown(
            "Consejo: usa un **User ID** constante para que la memoria de perfil y ToDos "
            "se mantenga entre mensajes. Cambia el **Thread ID** para conversaciones paralelas."
        )
    return demo

# ---------------------------------------------------------------------
# MAIN
# ---------------------------------------------------------------------
if __name__ == "__main__":
    demo = build_ui()
    demo.queue().launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )