"""
HuggingFace Inference Endpoint handler — capture test (T2).

Commit this as `handler.py` with `capture.py` and `requirements.txt` beside it in a model
repo. HuggingFace's Default container detects the handler and runs this instead of the stock
pipeline.

The signals are returned INSIDE the HTTP response under "_capture" rather than posted
anywhere. That is deliberate:

  * it proves capture works on this surface even if outbound network is blocked, so the
    capture question and the egress question stay independent;
  * it needs no API key, no database, and no platform to be up;
  * and it happens to be the fallback transport design we would need anyway if T1 fails.

Env vars (all optional):
    CAPTURE_ATTENTION=1     load with eager attention and capture attention metrics
                            (disables SDPA/FlashAttention — measure the cost, do not
                             leave this on)
    CAPTURE_TASK=causal     for decoder LMs; default "classification"
    CAPTURE_MAX_LAYERS=4    how many transformer blocks to sample
    CAPTURE_EXPORT_DIMS=256 how much of the embedding to return
"""

from __future__ import annotations

import os
import time
import traceback
from typing import Any

import torch

from capture import CaptureProbe, load_model

ATTENTION = os.environ.get("CAPTURE_ATTENTION", "0") == "1"
TASK = os.environ.get("CAPTURE_TASK", "classification")
MAX_LAYERS = int(os.environ.get("CAPTURE_MAX_LAYERS", "4"))
EXPORT_DIMS = int(os.environ.get("CAPTURE_EXPORT_DIMS", "256"))


class EndpointHandler:
    def __init__(self, path: str = ""):
        # HF mounts the selected model repo at `path` (/repository in the container).
        self.model, self.tokenizer = load_model(
            path or ".", task=TASK, capture_attention=ATTENTION
        )
        self.probe = CaptureProbe(
            self.model, capture_attention=ATTENTION, max_layers=MAX_LAYERS
        )
        self.env = {
            "capture_attention": ATTENTION,
            "task": TASK,
            "attn_implementation": getattr(self.model.config, "_attn_implementation", "unknown"),
            "torch": torch.__version__,
        }

    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
        text = data.pop("inputs", data)
        started = time.perf_counter()

        enc = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=256)

        try:
            with self.probe.capture() as cap:
                with torch.no_grad():
                    out = self.model(
                        **enc,
                        output_hidden_states=True,
                        output_attentions=ATTENTION,
                    )
            signals = self.probe.summarise(
                cap, inputs=enc["input_ids"], outputs=out, export_dims=EXPORT_DIMS
            )
            capture_error = None
        except Exception:                                   # noqa: BLE001
            # Capture must never take the endpoint down — that is the whole point of
            # testing it before anyone depends on it.
            signals, capture_error = {}, traceback.format_exc(limit=4)
            with torch.no_grad():
                out = self.model(**enc)

        total_ms = round((time.perf_counter() - started) * 1000, 2)

        prediction = self._prediction(out)
        return {
            "prediction": prediction,
            "_capture": signals,
            "_meta": {
                "total_ms": total_ms,
                "env": self.env,
                "capture_error": capture_error,
            },
        }

    def _prediction(self, out: Any) -> Any:
        logits = getattr(out, "logits", None)
        if logits is None:
            return None
        if logits.dim() == 3:                                # causal LM
            return {"next_token_id": int(logits[0, -1].argmax())}
        probs = torch.softmax(logits[0], dim=-1)
        id2label = getattr(self.model.config, "id2label", {})
        return [
            {"label": id2label.get(i, str(i)), "score": round(float(p), 6)}
            for i, p in enumerate(probs.tolist())
        ]