"""LangGraph-based agent runtime.

This is the core AI engine for EnCoach — every pipeline that used to call
``OpenAIService`` directly can instead go through an
:py:class:`AgentRuntime` loaded from a named :py:class:`encoach.ai.agent`
row. That buys us:

* A single place to reason about retries, logging, tool execution and
  self-review, instead of the same boilerplate inside every pipeline.
* Admins editing prompts, model choice, temperature or enabled tools in
  the UI without redeploys.
* A consistent shape (``invoke(variables, payload)``) across course
  planning, exam generation, LMS tutor, grading, etc.

Graph topologies
----------------

Each agent picks one of four graphs. They're all built on the same
:py:class:`AgentState` TypedDict so upgrading an agent from one topology
to another only means flipping a selection field.

``simple``
    ``START → llm → END``. The workhorse — used for deterministic
    JSON generation (course plan header, exam question batches).

``plan_review_revise``
    ``START → llm → review → [revise → llm]? → END``. ``review`` runs
    every configured quality tool (``quality.cefr_check`` etc.); if any
    returns ``ok=False`` we ask the LLM to revise once, capped by
    ``max_revisions`` on the agent. Keeps structured outputs (reading
    passages, listening scripts) inside their CEFR band.

``rag``
    ``START → retrieve → llm → END``. Runs ``resources.search`` before
    the LLM and injects the hits as extra system context. Used for
    curriculum-aware generation that must cite real library material.

``react``
    Classic tool-calling loop. The LLM is given the OpenAI-format tool
    list and can call tools in a loop until it emits a final answer.
    Used for the LMS tutor / study assistant.

We depend on ``langgraph`` for the orchestration (state machine,
conditional edges) but still call OpenAI through the existing
:py:class:`OpenAIService` so the API key wiring, retry behaviour and
``encoach.ai.log`` rows keep working unchanged.
"""

from __future__ import annotations

import json
import logging
import time
from typing import Any, TypedDict

from odoo.tools import config as odoo_config  # noqa: F401  (future use)

from . import agent_tools
from .openai_service import OpenAIService

_logger = logging.getLogger(__name__)


# =============================================================================
# State
# =============================================================================
class AgentState(TypedDict, total=False):
    """Shared state passed between every node of the graph."""

    messages: list[dict]            # chat history sent to the LLM
    output: Any                     # parsed final answer (dict or string)
    output_raw: str                 # the raw LLM text (for logging)
    tool_calls: list[dict]          # pending tool_calls emitted by the LLM
    tool_results: list[dict]        # results of executed tools, appended
    quality_issues: list[str]       # issues collected by the review node
    revisions_used: int             # revision counter (capped by max_revisions)
    variables: dict                 # caller-supplied variables (for prompt rendering)
    retrieval: list[dict]           # hits from the retrieval node (RAG)
    iterations: int                 # guard against runaway ReAct loops
    error: str                      # populated on fatal failure
    should_revise: bool             # set by review node when a revise pass is queued


# =============================================================================
# AgentRuntime
# =============================================================================
class AgentRuntime:
    """Wraps an :py:class:`encoach.ai.agent` row with a compiled LangGraph."""

    MAX_REACT_ITERATIONS = 6  # hard cap on tool-calling loops

    # ------------------------------------------------------------------
    # Factories
    # ------------------------------------------------------------------
    def __init__(self, env, agent, *, language: str | None = None):
        # INVARIANT: every AgentRuntime is per-request and constructs a
        # fresh OpenAIService, which reads ir.config_parameter on every
        # __init__. Combined with MediaService → provider_router (also
        # uncached), this guarantees that flipping a provider in the
        # admin UI takes effect on the very next request — no Odoo
        # restart, no cache invalidation. Don't introduce any
        # module-level or class-level provider caches here.
        self.env = env
        self.agent = agent
        self.language = language
        self.ai = OpenAIService(env, language=language)
        self._graph = None  # lazily compiled

    @classmethod
    def from_key(cls, env, key: str, *, language: str | None = None):
        """Factory: load the active agent with ``key`` and build its runtime.

        Returns ``None`` if no agent is configured for that key — callers
        can decide whether to fall back to their legacy direct-SDK path.
        """
        agent = env["encoach.ai.agent"].sudo().get_by_key(key)
        if not agent:
            return None
        return cls(env, agent, language=language)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def invoke(self, variables: dict | None = None, payload: Any = None,
               *, extra_system: str = "") -> AgentState:
        """Run the agent's graph end-to-end and return the terminal state.

        ``variables`` are substituted into the system prompt (if the agent
        is bound to a prompt_key). ``payload`` becomes the first user
        message; pass a dict to have it rendered as JSON, or a string to
        pass it through verbatim. ``extra_system`` lets callers tack on
        a context block without touching the agent's stored prompt.
        """
        t0 = time.time()
        graph = self._compile()
        initial = self._initial_state(variables or {}, payload, extra_system)
        try:
            # LangGraph is sync here — we're already inside an Odoo
            # request worker so sticking with sync keeps the control
            # flow simple.
            final: AgentState = graph.invoke(initial)
        except Exception as exc:
            _logger.exception("agent %s crashed", self.agent.key)
            final = {**initial, "error": str(exc)}

        latency_ms = int((time.time() - t0) * 1000)
        self._log(final, latency_ms)
        return final

    # ------------------------------------------------------------------
    # Graph construction
    # ------------------------------------------------------------------
    def _compile(self):
        if self._graph is not None:
            return self._graph
        try:
            from langgraph.graph import StateGraph, START, END
        except ImportError as exc:
            raise RuntimeError(
                "LangGraph is not installed. "
                "Add `langgraph>=0.2.0` to backend/requirements.txt and pip install."
            ) from exc

        g = StateGraph(AgentState)

        if self.agent.graph_type == "simple":
            g.add_node("llm", self._node_llm)
            g.add_edge(START, "llm")
            g.add_edge("llm", END)

        elif self.agent.graph_type == "rag":
            g.add_node("retrieve", self._node_retrieve)
            g.add_node("llm", self._node_llm)
            g.add_edge(START, "retrieve")
            g.add_edge("retrieve", "llm")
            g.add_edge("llm", END)

        elif self.agent.graph_type == "plan_review_revise":
            g.add_node("llm", self._node_llm)
            g.add_node("review", self._node_review)
            g.add_edge(START, "llm")
            g.add_edge("llm", "review")
            g.add_conditional_edges(
                "review",
                self._route_after_review,
                {"revise": "llm", "done": END},
            )

        elif self.agent.graph_type == "react":
            g.add_node("llm", self._node_llm_tools)
            g.add_node("tools", self._node_tools)
            g.add_edge(START, "llm")
            g.add_conditional_edges(
                "llm",
                self._route_after_llm_tools,
                {"tools": "tools", "done": END},
            )
            g.add_edge("tools", "llm")

        else:
            raise ValueError(f"Unknown graph_type: {self.agent.graph_type}")

        self._graph = g.compile()
        return self._graph

    # ------------------------------------------------------------------
    # Initial state
    # ------------------------------------------------------------------
    def _initial_state(self, variables: dict, payload: Any, extra_system: str) -> AgentState:
        system_prompt = self.agent.resolved_system_prompt(variables)
        messages: list[dict] = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        if extra_system:
            messages.append({"role": "system", "content": extra_system})
        if payload is not None:
            user_content = (
                payload if isinstance(payload, str)
                else json.dumps(payload, ensure_ascii=False)
            )
            messages.append({"role": "user", "content": user_content})
        return {
            "messages": messages,
            "output": None,
            "output_raw": "",
            "tool_calls": [],
            "tool_results": [],
            "quality_issues": [],
            "revisions_used": 0,
            "variables": variables,
            "retrieval": [],
            "iterations": 0,
            "error": "",
        }

    # ------------------------------------------------------------------
    # Nodes
    # ------------------------------------------------------------------
    def _node_llm(self, state: AgentState) -> AgentState:
        """Plain LLM call respecting the agent's model + response_format."""
        messages = list(state.get("messages") or [])
        model = self.agent.model
        action = f"agent.{self.agent.key}"
        try:
            if self.agent.response_format == "json":
                content = self.ai.chat_json(
                    messages,
                    model=model,
                    temperature=self.agent.temperature,
                    max_tokens=self.agent.max_tokens,
                    action=action,
                )
                raw = json.dumps(content, ensure_ascii=False)
            else:
                raw = self.ai.chat(
                    messages,
                    model=model,
                    temperature=self.agent.temperature,
                    max_tokens=self.agent.max_tokens,
                    action=action,
                )
                content = raw
        except Exception as exc:
            # Try the fallback model exactly once before giving up.
            if self.agent.fallback_model and model != self.agent.fallback_model:
                _logger.warning(
                    "agent %s primary model %s failed (%s); retrying with %s",
                    self.agent.key, model, exc, self.agent.fallback_model,
                )
                try:
                    if self.agent.response_format == "json":
                        content = self.ai.chat_json(
                            messages, model=self.agent.fallback_model,
                            temperature=self.agent.temperature,
                            max_tokens=self.agent.max_tokens,
                            action=f"{action}.fallback",
                        )
                        raw = json.dumps(content, ensure_ascii=False)
                    else:
                        raw = self.ai.chat(
                            messages, model=self.agent.fallback_model,
                            temperature=self.agent.temperature,
                            max_tokens=self.agent.max_tokens,
                            action=f"{action}.fallback",
                        )
                        content = raw
                except Exception as exc2:
                    return {**state, "error": str(exc2)}
            else:
                return {**state, "error": str(exc)}

        new_messages = messages + [{"role": "assistant", "content": raw}]
        return {
            **state,
            "messages": new_messages,
            "output": content,
            "output_raw": raw,
        }

    def _node_retrieve(self, state: AgentState) -> AgentState:
        """RAG node: call resources.search with the user's payload as query."""
        variables = state.get("variables") or {}
        query = ""
        # The last user message is our best default query.
        for m in reversed(state.get("messages") or []):
            if m.get("role") == "user":
                query = m.get("content") or ""
                break
        query = variables.get("query") or query
        hits = agent_tools.invoke(self.env, "resources.search", {
            "query": query[:1000],
            "limit": 5,
        })
        items = hits.get("items") or []
        context = self._format_retrieval(items)
        messages = list(state.get("messages") or [])
        if context:
            # Insert the context block *after* the system prompt(s).
            last_sys = -1
            for i, m in enumerate(messages):
                if m.get("role") == "system":
                    last_sys = i
            insert_at = last_sys + 1 if last_sys >= 0 else 0
            messages.insert(insert_at, {
                "role": "system",
                "content": (
                    "Relevant content from the library (use it when accurate, "
                    "cite ids; do not fabricate):\n\n" + context
                ),
            })
        return {**state, "messages": messages, "retrieval": items}

    def _node_review(self, state: AgentState) -> AgentState:
        """Run quality tools against the LLM output and (maybe) queue a revision.

        We do all state mutations here — adding the critique message and
        bumping ``revisions_used`` — and only signal the router with a
        boolean ``should_revise``. LangGraph routing functions are
        treated as pure: any state changes made there are discarded, so
        keeping the router pure prevents an infinite revise loop where
        the counter never actually increments.
        """
        # If the LLM step already errored out, don't waste a quality
        # check on the empty/garbage output and don't try to "revise" —
        # another LLM call would just hit the same permanent failure.
        if state.get("error"):
            return {**state, "quality_issues": [], "should_revise": False}

        text = state.get("output_raw") or ""
        if isinstance(state.get("output"), dict):
            # Flatten the dict to text so the quality tools see something
            # meaningful (most just want prose).
            text = json.dumps(state["output"], ensure_ascii=False)
        issues: list[str] = []
        checks = [
            k.strip() for k in (self.agent.quality_checks or "").split(",")
            if k.strip()
        ]
        variables = state.get("variables") or {}
        target_cefr = (
            variables.get("cefr_level")
            or variables.get("target_cefr")
            or "b1"
        )
        for key in checks:
            res = agent_tools.invoke(self.env, key, {
                "text": text,
                "target_cefr": target_cefr,
                "cefr_level": target_cefr,
            })
            if res.get("ok") is False:
                issues.extend(res.get("issues") or [res.get("error") or key])

        revisions_used = state.get("revisions_used") or 0
        max_rev = max(0, int(self.agent.max_revisions or 0))
        if issues and revisions_used < max_rev:
            critique = (
                "Your previous draft was rejected for the following reasons:\n- "
                + "\n- ".join(issues)
                + "\n\nProduce an improved version that addresses every issue. "
                "Keep the same JSON schema if one was requested."
            )
            messages = list(state.get("messages") or []) + [
                {"role": "system", "content": critique}
            ]
            return {
                **state,
                "messages": messages,
                "quality_issues": issues,
                "revisions_used": revisions_used + 1,
                "should_revise": True,
            }
        return {
            **state,
            "quality_issues": issues,
            "should_revise": False,
        }

    def _route_after_review(self, state: AgentState) -> str:
        # Pure router: only inspect state, never mutate. The decision
        # was prepared in ``_node_review``.
        if state.get("error"):
            return "done"
        return "revise" if state.get("should_revise") else "done"

    # ReAct / tool-calling -------------------------------------------------
    def _node_llm_tools(self, state: AgentState) -> AgentState:
        """ReAct step: ask the LLM, exposing all enabled tools."""
        if state.get("iterations", 0) >= self.MAX_REACT_ITERATIONS:
            return {**state, "error": "react_iteration_limit_exceeded"}

        client = self.ai.client
        if client is None:
            return {**state, "error": "openai_not_configured"}

        tools = [t.to_openai_tool() for t in self.agent.tool_ids]
        try:
            resp = client.chat.completions.create(
                model=self.agent.model,
                messages=state.get("messages") or [],
                temperature=self.agent.temperature,
                max_tokens=self.agent.max_tokens,
                tools=tools or None,
                tool_choice="auto" if tools else None,
                timeout=self.ai.request_timeout,
            )
        except Exception as exc:
            return {**state, "error": str(exc)}

        choice = resp.choices[0].message
        assistant_msg: dict[str, Any] = {
            "role": "assistant",
            "content": choice.content or "",
        }
        tool_calls = []
        if getattr(choice, "tool_calls", None):
            # Preserve the OpenAI-shaped tool_calls list on the message so
            # the next round references them by id.
            assistant_msg["tool_calls"] = [
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                }
                for tc in choice.tool_calls
            ]
            for tc in choice.tool_calls:
                try:
                    args = json.loads(tc.function.arguments or "{}")
                except Exception:
                    args = {}
                tool_calls.append({
                    "id": tc.id,
                    "name": tc.function.name,
                    "args": args,
                })

        new_messages = list(state.get("messages") or []) + [assistant_msg]
        return {
            **state,
            "messages": new_messages,
            "tool_calls": tool_calls,
            "output": choice.content or state.get("output"),
            "output_raw": choice.content or state.get("output_raw") or "",
            "iterations": state.get("iterations", 0) + 1,
        }

    def _route_after_llm_tools(self, state: AgentState) -> str:
        if state.get("error"):
            return "done"
        if state.get("tool_calls"):
            return "tools"
        return "done"

    def _node_tools(self, state: AgentState) -> AgentState:
        """Execute every queued tool_call and append results to the chat."""
        allowed = {t.key: t for t in self.agent.tool_ids}
        messages = list(state.get("messages") or [])
        results: list[dict] = list(state.get("tool_results") or [])
        for call in state.get("tool_calls") or []:
            # Tools are stored with dotted keys but OpenAI flattens dots to
            # double-underscores (because function names must match [A-Za-z0-9_]).
            key = (call.get("name") or "").replace("__", ".")
            if key not in allowed:
                result = {"error": f"tool_not_allowed:{key}"}
            else:
                result = agent_tools.invoke(self.env, key, call.get("args") or {})
            results.append({"tool": key, "args": call.get("args"), "result": result})
            messages.append({
                "role": "tool",
                "tool_call_id": call.get("id"),
                "name": call.get("name") or key,
                "content": json.dumps(result, ensure_ascii=False, default=str)[:6000],
            })
        return {
            **state,
            "messages": messages,
            "tool_calls": [],
            "tool_results": results,
        }

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _format_retrieval(items: list[dict]) -> str:
        parts = []
        for r in items or []:
            label = f"[{r.get('type','?')}#{r.get('id','?')}]"
            title = r.get("title") or ""
            snippet = (r.get("snippet") or "")[:400]
            parts.append(f"{label} {title}\n{snippet}")
        return "\n---\n".join(parts)

    def _log(self, final: AgentState, latency_ms: int):
        try:
            self.env["encoach.ai.log"].sudo().create({
                "service": "openai",
                "action": f"agent.{self.agent.key}",
                "model_used": self.agent.model,
                "latency_ms": latency_ms,
                "status": "error" if final.get("error") else "success",
                "error_message": final.get("error") or "",
                "input_preview": json.dumps(final.get("variables") or {})[:500],
                "output_preview": (final.get("output_raw") or "")[:500],
            })
        except Exception:
            _logger.warning("agent %s log write failed", self.agent.key, exc_info=True)
