"""Optional adapter: wrap real LLM calls (Anthropic Messages API) as pipeline steps.

Nothing here is imported by the core package, and the ``anthropic`` SDK is
imported lazily on the first real call -- tests use a fake client and never
touch the network.

How real LLM calls fit counterfactual replay
--------------------------------------------
Counterfactual blame needs the *same exogenous noise* in the factual run and
in every counterfactual run. A hosted LLM samples on the server, so its
sampling noise is not observable or controllable from the client. Three
practical regimes:

1. **Temperature 0.** The call is (approximately) a deterministic function of
   its prompt, so there is no sampling noise to share. Replay = cache lookup.
   Caveat: hosted inference is not guaranteed bit-for-bit deterministic even
   at temperature 0; caching the factual response pins it.

2. **Sampled (temperature > 0) free text.** We key the response cache on
   ``(step, variant, inputs, seed)`` where ``seed`` comes from the step's
   noise handle. Replaying the factual world hits the cache and reproduces
   the factual output exactly. But a *changed* step (new prompt/model) or
   changed inputs produce a new key and hence an *independent* sample: the
   noise is not shared between factual and counterfactual for that call.
   Exact Gumbel-max replay is impossible without per-token log-probabilities
   and control over the sampler (the Anthropic Messages API does not return
   token log-probabilities). Blame for such steps is therefore
   "interventional" rather than truly counterfactual; average over several
   seeds (``mode="prior"`` with ``n_worlds > 1``) to control the variance.

3. **Choices sampled client-side (recommended where possible).** If a step
   picks one of a fixed set of options (route, tool, document, verdict), ask
   the model for a *score per option* at temperature 0 and sample the choice
   locally with ``noise.categorical`` (Gumbel-max). The randomness now lives
   in agentblame's noise handle, so counterfactual replay is exact: a better
   prompt changes the scores (logits) while the Gumbel noise stays fixed.
   See :func:`llm_choice_step`.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import threading
from typing import Any, Callable, Mapping, Sequence

from ..pipeline import Step

__all__ = ["AnthropicLLM", "ResponseCache", "llm_choice_step", "llm_text_step"]


class ResponseCache:
    """A tiny JSON-file-backed cache of LLM responses (in-memory if ``path`` is None)."""

    def __init__(self, path: str | None = None):
        self.path = path
        self._lock = threading.Lock()
        self._data: dict[str, str] = {}
        if path and os.path.exists(path):
            with open(path, encoding="utf-8") as f:
                self._data = json.load(f)

    @staticmethod
    def key(**parts: Any) -> str:
        blob = json.dumps(parts, sort_keys=True, default=repr)
        return hashlib.sha256(blob.encode()).hexdigest()

    def get(self, key: str) -> str | None:
        return self._data.get(key)

    def put(self, key: str, value: str) -> None:
        with self._lock:
            self._data[key] = value
            if self.path:
                tmp = self.path + ".tmp"
                with open(tmp, "w", encoding="utf-8") as f:
                    json.dump(self._data, f)
                os.replace(tmp, self.path)

    def __len__(self) -> int:
        return len(self._data)


class AnthropicLLM:
    """Cached wrapper around ``anthropic.Anthropic().messages.create``.

    Parameters
    ----------
    model:
        Model id to call (required; pass whichever model you use).
    cache:
        A :class:`ResponseCache`; responses are looked up before any call.
    client:
        Optional pre-built client (anything with ``.messages.create(...)``);
        tests inject a fake. If omitted, the SDK is imported on first use and
        reads ``ANTHROPIC_API_KEY`` from the environment.
    offline:
        If True, a cache miss raises instead of calling the API.
    """

    def __init__(self, model: str, cache: ResponseCache | None = None, client: Any = None,
                 max_tokens: int = 512, offline: bool = False):
        self.model = model
        self.cache = cache if cache is not None else ResponseCache()   # (empty cache is falsy)
        self._client = client
        self.max_tokens = max_tokens
        self.offline = offline
        self.calls = 0                       # number of real (non-cached) calls

    def _get_client(self):
        if self._client is None:
            try:
                import anthropic  # lazy: optional dependency
            except ImportError as e:  # pragma: no cover - depends on environment
                raise ImportError("pip install anthropic to make real LLM calls") from e
            self._client = anthropic.Anthropic()
        return self._client

    def complete(self, prompt: str, *, step: str, variant: str, inputs: Mapping[str, Any],
                 seed: int = 0, temperature: float = 0.0, system: str | None = None,
                 model: str | None = None) -> str:
        """Return the response text for ``prompt``, cached on (step, variant, inputs, seed, ...).

        At temperature 0 the seed is dropped from the key: the call is treated
        as deterministic, so every world shares one cached response.
        """
        model = model or self.model
        key = ResponseCache.key(step=step, variant=variant, inputs=inputs, model=model,
                                prompt=prompt, system=system, temperature=temperature,
                                seed=seed if temperature > 0 else None)
        hit = self.cache.get(key)
        if hit is not None:
            return hit
        if self.offline:
            raise LookupError(f"cache miss for step '{step}' variant '{variant}' in offline mode")
        kwargs = dict(model=model, max_tokens=self.max_tokens, temperature=temperature,
                      messages=[{"role": "user", "content": prompt}])
        if system:
            kwargs["system"] = system
        resp = self._get_client().messages.create(**kwargs)
        self.calls += 1
        text = "".join(getattr(b, "text", "") for b in resp.content
                       if getattr(b, "type", "text") == "text")
        self.cache.put(key, text)
        return text


def _seed_from(noise, key: str = "llm_seed") -> int:
    """An integer seed drawn from the step's noise handle (identifies a sample slot)."""
    return int(noise.uniform(key) * 2**31)


def llm_text_step(name: str, llm: AnthropicLLM, prompt: str, deps: Sequence[str], *,
                  temperature: float = 0.0, system: str | None = None,
                  variants: Mapping[str, Mapping[str, Any]] | None = None,
                  parse: Callable[[str], Any] = lambda s: s, cost: float = 1.0) -> Step:
    """Build a :class:`Step` whose implementations are LLM calls.

    ``prompt`` is a ``str.format`` template over ``deps``. Each entry of
    ``variants`` overrides any of ``prompt``, ``system``, ``temperature``,
    ``model``, ``cost`` for an alternative implementation (e.g. a better
    prompt or stronger model). ``parse`` post-processes the text.
    """
    def make(vname: str, cfg: Mapping[str, Any]):
        p = cfg.get("prompt", prompt)
        t = cfg.get("temperature", temperature)
        sysm = cfg.get("system", system)
        mdl = cfg.get("model")

        def fn(noise, **inputs):
            seed = _seed_from(noise)
            text = llm.complete(p.format(**inputs), step=name, variant=vname, inputs=inputs,
                                seed=seed, temperature=t, system=sysm, model=mdl)
            return parse(text)
        return fn

    st = Step(make("as_run", {}), name=name, deps=deps, cost=cost)
    for vname, cfg in (variants or {}).items():
        st.variant(vname, cost=cfg.get("cost"))(make(vname, cfg))
    return st


_SCORE_INSTRUCTIONS = (
    "\n\nRate how appropriate each option is on a scale from 0 to 10. Reply with JSON only, "
    "mapping each option to its score, e.g. {{\"option_a\": 7, \"option_b\": 2}}. Options: {options}")


def _parse_scores(text: str, options: Sequence[str]) -> dict[str, float]:
    m = re.search(r"\{.*\}", text, re.S)
    raw = json.loads(m.group(0)) if m else {}
    return {o: float(raw.get(o, 0.0)) for o in options}


def llm_choice_step(name: str, llm: AnthropicLLM, prompt: str, options: Sequence[str],
                    deps: Sequence[str], *, sharpness: float = 1.0,
                    variants: Mapping[str, Mapping[str, Any]] | None = None,
                    cost: float = 1.0) -> Step:
    """A choice step whose randomness is client-side and therefore replayable.

    The model scores each option at temperature 0 (cached, deterministic);
    ``sharpness * score`` are used as logits and the option is sampled with
    ``noise.categorical`` (Gumbel-max). Counterfactual variants (other prompt
    or model) change the logits but reuse the same Gumbel noise per option.
    """
    def make(vname: str, cfg: Mapping[str, Any]):
        p = cfg.get("prompt", prompt)
        suffix = _SCORE_INSTRUCTIONS.format(options=", ".join(options))   # appended after .format
        mdl = cfg.get("model")
        k = cfg.get("sharpness", sharpness)

        def fn(noise, **inputs):
            text = llm.complete(p.format(**inputs) + suffix, step=name, variant=vname,
                                inputs=inputs, temperature=0.0, model=mdl)
            scores = _parse_scores(text, options)
            return noise.categorical({o: k * s for o, s in scores.items()})
        return fn

    st = Step(make("as_run", {}), name=name, deps=deps, cost=cost)
    for vname, cfg in (variants or {}).items():
        st.variant(vname, cost=cfg.get("cost"))(make(vname, cfg))
    return st