"""Provider-agnostic core for ABOM runtime adapters.

An adapter's only provider-specific job is to (1) recognize its client type and
(2) know how to read a model call's ``model``, ``tokens``, and ``output`` out of
that provider's request/response. Everything else — appending to the provenance
chain, calling the gate, binding to the manifest — lives here and is shared.
"""
from __future__ import annotations

import hashlib
import json
import threading
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Protocol

from .. import bom
from ..gate import Action, Gate
from .sink import Sink, resolve_sink, safe_slug


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


@dataclass
class ModelCall:
    """One captured model call, provider-normalized."""

    model: str
    tokens: int | None = None
    output_sha256: str | None = None
    endpoint: str | None = None


class Adapter(Protocol):
    """A per-provider adapter. ``provider`` names it; ``matches`` recognizes a
    client instance; ``wrap`` returns an instrumented client; ``parse`` reads a
    response into a :class:`ModelCall`."""

    provider: str

    def matches(self, client: Any) -> bool: ...
    def wrap(self, client: Any, run: "AbomRun") -> Any: ...


@dataclass
class AbomRun:
    """Holds the live provenance chain + optional gate for an instrumented run.

    Attached to the wrapped client as ``client.abom`` so callers can read
    ``client.abom.chain`` after the run. When a ``sink`` is set, every record is
    also streamed to it the moment it is appended, so the evidence survives the
    process (see :mod:`abom.adapters.sink`).
    """

    composition_sha256: str
    agent_ref: str
    gate: Gate | None = None
    chain: list = field(default_factory=list)
    run_id: str = ""
    sink: Sink | None = None
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False)

    def record(self, call: ModelCall, *, decision: str = "model_call") -> dict:
        """Append a captured model call to the Action Provenance chain.

        Thread-safe (concurrent calls get unique, ordered ``seq`` values). If a
        sink is configured, the record is written through before returning — a
        sink failure raises rather than silently dropping evidence.
        """
        with self._lock:
            row = bom.append_action(
                self.chain,
                composition_sha256=self.composition_sha256,
                agent_ref=self.agent_ref,
                decision=decision,
                model_calls=[
                    {
                        "model": call.model,
                        "tokens": call.tokens,
                        "output_sha256": call.output_sha256,
                    }
                ],
            )
            if self.sink is not None:
                self.sink.append(row)
            return row

    def record_tool(
        self,
        name: str,
        *,
        endpoint: str | None = None,
        args_sha256: str | None = None,
        decision: str = "tool_call",
    ) -> dict:
        """Append a captured tool invocation to the Action Provenance chain."""
        with self._lock:
            row = bom.append_action(
                self.chain,
                composition_sha256=self.composition_sha256,
                agent_ref=self.agent_ref,
                decision=decision,
                tools_invoked=[
                    {"name": name, "endpoint": endpoint, "args_sha256": args_sha256}
                ],
            )
            if self.sink is not None:
                self.sink.append(row)
            return row

    def enforce(self, call: ModelCall) -> None:
        """If a gate is configured, deny-by-default a call outside the envelope.

        A model call is treated as a tool named after the model, with the
        provider endpoint as its egress endpoint. Raises ``ActionDenied`` (from
        the gate) if the model isn't in the signed manifest.
        """
        if self.gate is None:
            return
        self.gate.enforce(
            Action(tool=call.model, endpoint=call.endpoint, kind="model")
        )

    def enforce_tool(self, name: str, *, endpoint: str | None = None) -> None:
        """Deny-by-default gate for a tool invocation (no-op without a gate).

        Raises ``ActionDenied`` if the tool is not declared in the signed
        manifest, or its endpoint is outside the tool's allowlist.
        """
        if self.gate is None:
            return
        self.gate.enforce(Action(tool=name, endpoint=endpoint, kind="tool"))


def _load_manifest(manifest: dict | str | Path) -> dict:
    if isinstance(manifest, (str, Path)):
        return json.loads(Path(manifest).read_text())
    return dict(manifest)


# Registry of available adapters, populated by each provider module.
_ADAPTERS: list[Adapter] = []


def register(adapter: Adapter) -> None:
    _ADAPTERS.append(adapter)


def start_run(
    manifest: dict | str | Path,
    *,
    gate: bool = False,
    agent_ref: str | None = None,
    sink: Any = None,
) -> AbomRun:
    """Create an :class:`AbomRun` bound to a manifest, with optional gate + sink.

    The shared entry point for every integration style: ``instrument`` uses it
    before wrapping an SDK client, and callback-based integrations (e.g. the
    LangChain handler) use it directly since they have no client to wrap.
    """
    m = _load_manifest(manifest)
    comp_sha = m.get("composition_sha256", "")
    ref = agent_ref or m.get("agent", {}).get("name", "run")
    run_id = (
        f"{safe_slug(ref)}-{datetime.now(timezone.utc):%Y%m%dT%H%M%SZ}-{uuid.uuid4().hex[:8]}"
    )
    return AbomRun(
        composition_sha256=comp_sha,
        agent_ref=ref,
        gate=Gate(m, run_id=ref) if gate else None,
        run_id=run_id,
        sink=resolve_sink(sink, run_id=run_id),
    )


def instrument(
    client: Any,
    *,
    manifest: dict | str | Path,
    gate: bool = False,
    provider: str | None = None,
    agent_ref: str | None = None,
    sink: Any = None,
) -> Any:
    """Wrap ``client`` so its model calls are recorded (and optionally gated).

    Args:
        client: a provider SDK client (e.g. ``anthropic.Anthropic()``).
        manifest: the signed Composition Manifest, as a dict or a path to JSON.
        gate: when True, enforce the manifest — a model not in it is denied.
        provider: force a provider instead of auto-detecting from the client.
        agent_ref: run identifier; defaults to the manifest's ``agent.name``.
        sink: where to persist the chain as it grows. ``None`` (default) keeps
            it in memory only; ``True``/``"auto"`` streams each record to
            ``.abom/runs/<run_id>.jsonl``; a path selects the file (``.jsonl``)
            or directory; any object with ``.append(record)`` is used as-is.
            Inspect afterwards with ``abom runs list`` / ``abom runs show``, or
            audit with ``abom verify abom.json --run <run_id> --policy p.json``.

    Returns the same client object, instrumented in place, with an ``abom``
    attribute exposing the live provenance chain (``client.abom.chain``).
    """
    run = start_run(manifest, gate=gate, agent_ref=agent_ref, sink=sink)

    adapter = _select_adapter(client, provider)
    wrapped = adapter.wrap(client, run)
    # expose the run on the returned client for post-hoc inspection
    try:
        setattr(wrapped, "abom", run)
    except (AttributeError, TypeError):
        pass
    return wrapped


def _select_adapter(client: Any, provider: str | None) -> Adapter:
    if provider is not None:
        for a in _ADAPTERS:
            if a.provider == provider:
                return a
        raise ValueError(
            f"no ABOM adapter registered for provider '{provider}' "
            f"(available: {[a.provider for a in _ADAPTERS]})"
        )
    for a in _ADAPTERS:
        if a.matches(client):
            return a
    raise ValueError(
        f"could not auto-detect a provider for {type(client).__module__}."
        f"{type(client).__name__}; pass provider=… explicitly "
        f"(available: {[a.provider for a in _ADAPTERS]})"
    )
