"""LangChain adapter — records and gates an agent's model AND tool calls.

LangChain's integration seam is a *callback handler*, not a wrappable client,
so this adapter is a handler you pass at invocation time::

    from abom.adapters.langchain import AbomLangChainHandler

    handler = AbomLangChainHandler(manifest="abom.json", gate=True, sink="auto")
    agent.invoke({"input": "..."}, config={"callbacks": [handler]})

    handler.abom.chain      # model + tool records, hash-chained, persisted

Two properties make this a *control point* rather than a tap:

* ``raise_error = True`` — LangChain's callback manager re-raises a handler's
  exception instead of logging it. Combined with gating in ``on_tool_start``
  (which LangChain dispatches *before* executing the tool), an undeclared tool
  raises ``ActionDenied`` and the tool body never runs.
* the handler shares the exact primitives of the SDK adapters (``start_run``,
  the gate's ``kind="tool"``/``kind="model"`` rules, the sink), so its records
  verify with ``abom verify --run`` like any other chain.

The ``langchain`` package is NOT a hard dependency: if ``langchain_core`` is
importable the handler subclasses its ``BaseCallbackHandler`` (inheriting no-op
defaults for every event we don't use); otherwise a minimal stand-in keeps the
class importable and testable.
"""
from __future__ import annotations

import hashlib
from pathlib import Path
from typing import Any

from .base import ModelCall, start_run

try:  # real base class when LangChain is installed — never a hard dependency
    from langchain_core.callbacks import BaseCallbackHandler as _Base
except ImportError:  # pragma: no cover - exercised implicitly in CI without langchain

    class _Base:  # minimal stand-in with the attributes LangChain would provide
        """Placeholder for langchain_core's BaseCallbackHandler."""


def _sha256(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _model_name(serialized: Any, kwargs: dict) -> str | None:
    """Best-effort model id from what LangChain hands ``on_*_model_start``.

    Looks in ``invocation_params`` (the reliable spot), then the serialized
    constructor kwargs. Returns None when unresolvable — the caller decides
    whether that is a deny (strict) or a record-as-unknown (default).
    """
    for source in (kwargs.get("invocation_params"), (serialized or {}).get("kwargs")):
        if isinstance(source, dict):
            for key in ("model", "model_name", "model_id"):
                if source.get(key):
                    return str(source[key])
    return None


class AbomLangChainHandler(_Base):
    """LangChain callback handler that gates and records model + tool calls.

    Args:
        manifest: signed Composition Manifest (dict or path), as for ``instrument``.
        gate: when True, deny-by-default — an undeclared tool (or resolvable
            undeclared model) raises ``ActionDenied`` before it executes.
        strict_models: when True, a model whose name cannot be resolved from
            the callback payload is also denied. Default False: unresolvable
            names are recorded as ``"unknown"`` but not denied, because tool
            gating (always strict) is the enforcement surface here and model
            identity is better enforced by the SDK adapters.
        agent_ref / sink: as for ``instrument``.
    """

    # LangChain re-raises handler exceptions only when this is True; it is the
    # documented mechanism that lets a gate denial abort the action.
    raise_error = True
    run_inline = True  # dispatch on the caller's thread, so a raise lands there

    def __init__(
        self,
        *,
        manifest: dict | str | Path,
        gate: bool = False,
        strict_models: bool = False,
        agent_ref: str | None = None,
        sink: Any = None,
    ):
        super().__init__()
        self.abom = start_run(manifest, gate=gate, agent_ref=agent_ref, sink=sink)
        self.strict_models = strict_models
        self._pending_models: dict[str, str] = {}  # lc run_id -> model name
        self._pending_tools: dict[str, tuple[str, str | None]] = {}  # -> (name, args_sha)

    # -- models ----------------------------------------------------------------
    def _model_start(self, serialized: Any, *, run_id: Any = None, **kwargs: Any) -> None:
        model = _model_name(serialized, kwargs)
        if model is None and self.strict_models:
            model = "unknown"  # unresolvable + strict -> gate it (and be denied)
        if model is not None:
            self.abom.enforce(ModelCall(model=model))
        self._pending_models[str(run_id)] = model or "unknown"

    def on_chat_model_start(self, serialized: Any, messages: Any, **kwargs: Any) -> None:
        self._model_start(serialized, **kwargs)

    def on_llm_start(self, serialized: Any, prompts: Any, **kwargs: Any) -> None:
        self._model_start(serialized, **kwargs)

    def on_llm_end(self, response: Any, *, run_id: Any = None, **kwargs: Any) -> None:
        llm_output = getattr(response, "llm_output", None)
        if llm_output is None and isinstance(response, dict):
            llm_output = response.get("llm_output")
        llm_output = llm_output or {}

        # prefer the name that was gated at start; fall back to what the
        # provider echoed in llm_output when the start event couldn't resolve it
        model = self._pending_models.pop(str(run_id), None)
        if not model or model == "unknown":
            model = llm_output.get("model_name") or "unknown"
        usage = llm_output.get("token_usage") or {}
        tokens = usage.get("total_tokens")

        texts = []
        generations = getattr(response, "generations", None)
        if generations is None and isinstance(response, dict):
            generations = response.get("generations")
        for batch in generations or []:
            for gen in batch:
                text = getattr(gen, "text", None)
                if text is None and isinstance(gen, dict):
                    text = gen.get("text")
                if text:
                    texts.append(text)
        joined = "".join(texts)

        self.abom.record(
            ModelCall(
                model=model,
                tokens=tokens,
                output_sha256=_sha256(joined) if joined else None,
            )
        )

    # -- tools -------------------------------------------------------------------
    def on_tool_start(
        self, serialized: Any, input_str: Any, *, run_id: Any = None, **kwargs: Any
    ) -> None:
        name = (serialized or {}).get("name") or "unknown"
        # Gate BEFORE the tool executes: LangChain dispatches this event first,
        # and raise_error=True propagates the denial, aborting the tool call.
        self.abom.enforce_tool(name)
        args_sha = _sha256(str(input_str)) if input_str is not None else None
        self._pending_tools[str(run_id)] = (name, args_sha)

    def on_tool_end(self, output: Any, *, run_id: Any = None, **kwargs: Any) -> None:
        name, args_sha = self._pending_tools.pop(str(run_id), ("unknown", None))
        self.abom.record_tool(name, args_sha256=args_sha)

    def on_tool_error(self, error: Any, *, run_id: Any = None, **kwargs: Any) -> None:
        # the tool ran and failed: still evidence — record it as an errored call
        name, args_sha = self._pending_tools.pop(str(run_id), ("unknown", None))
        self.abom.record_tool(name, args_sha256=args_sha, decision="tool_error")
