from langgraph.graph import StateGraph, START, END
from langchain_core.messages import SystemMessage, ToolMessage, BaseMessage
from typing import Literal, Callable, Any
from dataclasses import dataclass, field
from .schemas import MessagesState


@dataclass
class AgentConfig:
    """Configuración inmutable del agente para fácil testing."""
    system_prompt: str = (
        "You are a helpful assistant tasked with "
        "performing arithmetic on a set of inputs."
    )
    max_iterations: int = 10
    timeout_seconds: float = 30.0


class Agent:
    """
    Agente conversacional basado en LangGraph.
    """

    def __init__(
        self,
        model: Any,
        config: AgentConfig | None = None,
        tools: list | None = None
    ):
        self._model = model
        self._config = config or AgentConfig()
        self._tools = tools or []
        self._tools_by_name: dict[str, Any] = {
            tool.name: tool for tool in self._tools
        }
        self._compiled_agent = self._compile_agent()

        # Hooks para testing (monkey-patching friendly)
        self._pre_llm_hook: Callable | None = None
        self._post_tool_hook: Callable | None = None

    @property
    def tools(self) -> list:
        """Retorna copia de herramientas registradas."""
        return self._tools.copy()

    @property
    def config(self) -> AgentConfig:
        """Retorna configuración actual."""
        return self._config

    @property
    def graph(self) -> StateGraph:
        """Expone el grafo para inspección en tests."""
        return self._build_agent()

    def invoke(self, state: dict) -> dict:
        """
        Ejecuta el agente con el estado inicial.

        Args:
            state: Dict con key "messages" y opcionalmente "llm_calls"

        Returns:
            Estado final después de la ejecución
        """

        if "messages" not in state:
            raise ValueError("State must contain 'messages' key")

        return self._compiled_agent.invoke(state)

    def stream(self, state: dict):
        """Streaming de la ejecución para debugging."""
        yield from self._compiled_agent.stream(state)

    # ============ Hooks para Testing ============

    def set_pre_llm_hook(self, hook: Callable[[dict], None]):
        """Registra hook pre-llm para testing."""
        self._pre_llm_hook = hook

    def set_post_tool_hook(self, hook: Callable[[dict, Any], None]):
        """Registra hook post-tool para testing."""
        self._post_tool_hook = hook

    def _compile_agent(self):
        """Compila el grafo una sola vez."""
        agent = self._build_agent()
        return agent.compile()

    def _build_agent(self) -> StateGraph:
        """Construye el grafo de estado completo."""
        agent_builder = StateGraph(MessagesState)

        # Nodos
        agent_builder.add_node("llm_call", self._llm_call_node)
        agent_builder.add_node("tool_node", self._tool_node)

        # Flujo
        agent_builder.add_edge(START, "llm_call")
        agent_builder.add_conditional_edges(
            "llm_call",
            self._should_continue,
            {
                "tool_node": "tool_node",
                END: END
            }
        )
        agent_builder.add_edge("tool_node", "llm_call")

        return agent_builder

    # ============ Nodos del Grafo ============

    def _llm_call_node(self, state: MessagesState) -> dict:
        """
        Nodo LLM: Decide siguiente acción o responde.

        State transitions:
        - Si tool_calls → tool_node
        - Si no → END
        """
        if self._pre_llm_hook:
            self._pre_llm_hook(state)

        messages = [
            SystemMessage(content=self._config.system_prompt)
        ] + state["messages"]

        model_with_tools = self._model.bind_tools(self._tools)
        response = model_with_tools.invoke(messages)

        current_calls = state.get("llm_calls", 0)
        new_calls = current_calls + 1

        if new_calls > self._config.max_iterations:
            raise RuntimeError(
                f"Max iterations ({self._config.max_iterations}) exceeded"
            )

        return {
            "messages": state["messages"] + [response],
            "llm_calls": new_calls
        }

    def _tool_node(self, state: MessagesState) -> dict:
        """
        Nodo Tools: Ejecuta herramientas solicitadas por el LLM.
        """
        last_message = state["messages"][-1]

        if not hasattr(last_message, "tool_calls"):
            return {"messages": state["messages"]}

        tool_results: list[BaseMessage] = []

        for tool_call in last_message.tool_calls:
            result = self._execute_tool(tool_call)
            tool_results.append(result)

            if self._post_tool_hook:
                self._post_tool_hook(state, result)

        return {
            "messages": state["messages"] + tool_results
        }

    def _execute_tool(self, tool_call: dict) -> ToolMessage:
        """
        Ejecuta una herramienta individual con manejo de errores.
        """
        tool_name = tool_call.get("name")
        tool_id = tool_call.get("id", "unknown")
        tool_args = tool_call.get("args", {})

        if tool_name not in self._tools_by_name:
            return ToolMessage(
                content=f"Error: Tool '{tool_name}' not found",
                tool_call_id=tool_id,
                status="error"
            )

        try:
            tool = self._tools_by_name[tool_name]
            observation = tool.invoke(tool_args)
            return ToolMessage(
                content=str(observation),
                tool_call_id=tool_id,
                status="success"
            )
        except Exception as e:
            return ToolMessage(
                content=f"Error executing tool: {str(e)}",
                tool_call_id=tool_id,
                status="error"
            )

    def _should_continue(
            self, state: MessagesState
    ) -> Literal["tool_node", END]:
        """
        Condición de transición: ¿El LLM solicitó herramientas?
        """
        messages = state.get("messages", [])
        if not messages:
            return END

        last_message = messages[-1]

        if hasattr(last_message, "tool_calls") and last_message.tool_calls:
            return "tool_node"

        return END