"""Bounded LangGraph agent used for the agent-loop workshop demo."""

from __future__ import annotations

from collections.abc import Iterator
from pathlib import Path
from typing import Annotated, Any, NotRequired, TypedDict

from langchain_core.messages import (
    AIMessage,
    AnyMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langsmith import traceable

from src.config import build_langchain_llm
from src.schemas import SystemOutput
from src.tools import LANGCHAIN_TOOLS, TOOLS_BY_NAME

PROMPTS_DIR = Path(__file__).resolve().parent.parent / "prompts"
SYSTEM_PROMPT = (PROMPTS_DIR / "support_system_prompt.txt").read_text()
MAX_ITERATIONS = 4


class AgentState(TypedDict):
    """Store LangGraph agent state across model and tool steps."""

    messages: Annotated[list[AnyMessage], add_messages]
    metadata: dict[str, Any]
    used_tools: list[str]
    iterations: int
    final_output: NotRequired[SystemOutput]


def _build_model() -> Any:
    """Build the tool-bound LLM used inside the agent loop."""
    llm = build_langchain_llm(temperature=0)
    return llm.bind_tools(LANGCHAIN_TOOLS)


@traceable(name="agent.llm_call", run_type="chain")
def llm_call(state: AgentState) -> dict[str, Any]:
    """Ask the model for the next agent step."""
    model = _build_model()
    response = model.invoke(
        [
            SystemMessage(
                content=(
                    SYSTEM_PROMPT
                    + "\nYou are running inside a bounded helpdesk agent loop."
                    "\nUse tools only when needed, avoid redundant calls, "
                    "and stop once you can answer."
                    "\nNever exceed the available information."
                )
            ),
            *state["messages"],
        ]
    )
    return {"messages": [response], "iterations": state["iterations"] + 1}


@traceable(name="agent.tool_node", run_type="chain")
def tool_node(state: AgentState) -> dict[str, Any]:
    """Execute the tools requested by the latest AI message."""
    last_message = state["messages"][-1]
    if not isinstance(last_message, AIMessage):
        raise ValueError("Tool node expected an AI message.")

    tool_messages: list[ToolMessage] = []
    used_tools = list(state["used_tools"])
    for tool_call in last_message.tool_calls:
        tool_name = tool_call["name"]
        tool_ = TOOLS_BY_NAME[tool_name]
        tool_result = tool_.invoke(tool_call)
        used_tools.append(tool_name)
        tool_messages.append(
            ToolMessage(
                content=str(tool_result),
                tool_call_id=tool_call["id"],
                name=tool_name,
            )
        )
    return {"messages": tool_messages, "used_tools": used_tools}


def should_continue(state: AgentState) -> str:
    """Choose whether the graph should continue with tools or finalize."""
    last_message = state["messages"][-1]
    if state["iterations"] >= MAX_ITERATIONS:
        return "finalize_output"
    if isinstance(last_message, AIMessage) and last_message.tool_calls:
        return "tool_node"
    return "finalize_output"


@traceable(name="agent.finalize_output", run_type="chain")
def finalize_output(state: AgentState) -> dict[str, SystemOutput]:
    """Convert the conversation trace into the final structured output."""
    llm = build_langchain_llm(temperature=0)
    structured_llm = llm.with_structured_output(SystemOutput)
    response = structured_llm.invoke(
        [
            SystemMessage(
                content=(
                    "Convert the prior conversation into the final support system output."
                    " Use the provided metadata ticket_id."
                )
            ),
            *state["messages"],
            HumanMessage(
                content=(
                    f"Metadata: {state['metadata']}\n"
                    f"Used tools: {state['used_tools']}\n"
                    "Return the final structured answer now."
                )
            ),
        ]
    )
    response.used_tools = state["used_tools"]
    response.tool_call_count = len(state["used_tools"])
    response.ticket_id = state["metadata"]["ticket_id"]
    return {"final_output": response}


def build_agent_graph() -> Any:
    """Compile and return the LangGraph agent application."""
    graph = StateGraph(AgentState)
    graph.add_node("llm_call", llm_call)
    graph.add_node("tool_node", tool_node)
    graph.add_node("finalize_output", finalize_output)
    graph.add_edge(START, "llm_call")
    graph.add_conditional_edges(
        "llm_call",
        should_continue,
        {
            "tool_node": "tool_node",
            "finalize_output": "finalize_output",
        },
    )
    graph.add_edge("tool_node", "llm_call")
    graph.add_edge("finalize_output", END)
    return graph.compile()


def _build_agent_input(ticket_text: str, metadata: dict[str, Any]) -> dict[str, Any]:
    """Build the common input payload passed into the compiled agent graph."""
    return {
        "messages": [
            HumanMessage(
                content=f"Ticket metadata: {metadata}\nTicket text: {ticket_text}",
            )
        ],
        "metadata": metadata,
        "used_tools": [],
        "iterations": 0,
    }


@traceable(name="agent.run_agent", run_type="chain")
def run_agent(ticket_text: str, metadata: dict[str, Any]) -> SystemOutput:
    """Run the bounded agent and return the final structured output."""
    app = build_agent_graph()
    result = app.invoke(_build_agent_input(ticket_text, metadata))
    return result["final_output"]


@traceable(name="agent.stream_agent_steps", run_type="chain")
def stream_agent_steps(ticket_text: str, metadata: dict[str, Any]) -> Iterator[dict[str, Any]]:
    """Stream intermediate agent graph updates for notebook display."""
    app = build_agent_graph()
    return app.stream(_build_agent_input(ticket_text, metadata), stream_mode="updates")
