"""Shared tracing datatypes."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from torch import Tensor

from .probabilities import (
    gather_token_scores,
    logits_to_log_probs,
    logits_to_probs,
    select_final_token,
)
from .sites import SiteName


@dataclass(frozen=True)
class LayerSiteKey:
    """A stable key for a site capture."""

    site: SiteName
    layer: int | None = None


@dataclass(frozen=True)
class SiteRequest:
    """A normalized trace request for a site."""

    site: SiteName
    layers: tuple[int, ...] | None = None


@dataclass
class TraceResult:
    """Collected tensors from a traced forward pass."""

    activations: dict[SiteName, dict[int | None, Tensor]] = field(default_factory=dict)
    final_logits: Tensor | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    def add(self, site: SiteName, tensor: Tensor, layer: int | None = None) -> None:
        self.activations.setdefault(site, {})[layer] = tensor

    def get(self, site: SiteName, layer: int | None = None) -> Tensor:
        try:
            return self.activations[site][layer]
        except KeyError as exc:
            raise KeyError(f"No capture stored for site={site!r}, layer={layer!r}.") from exc

    def has(self, site: SiteName, layer: int | None = None) -> bool:
        return site in self.activations and layer in self.activations[site]

    @property
    def available_sites(self) -> tuple[SiteName, ...]:
        return tuple(self.activations)

    def final_log_probs(self) -> Tensor:
        """Return log-probabilities derived from the model's final logits."""

        if self.final_logits is None:
            raise ValueError("This trace does not contain final logits.")
        return logits_to_log_probs(self.final_logits)

    def final_probs(self) -> Tensor:
        """Return probabilities derived from the model's final logits."""

        if self.final_logits is None:
            raise ValueError("This trace does not contain final logits.")
        return logits_to_probs(self.final_logits)

    def layer_log_probs(self, layer: int) -> Tensor:
        """Return log-probabilities for a traced layer's logits."""

        return logits_to_log_probs(self.get("logits", layer))

    def layer_probs(self, layer: int) -> Tensor:
        """Return probabilities for a traced layer's logits."""

        return logits_to_probs(self.get("logits", layer))


@dataclass
class TokenScores:
    """Final-model token scores with logits-derived convenience accessors."""

    logits: Tensor
    input_ids: Tensor
    attention_mask: Tensor | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    temperature: float = 1.0
    _log_probs: Tensor | None = field(default=None, init=False, repr=False)
    _probs: Tensor | None = field(default=None, init=False, repr=False)

    def __post_init__(self) -> None:
        if self.logits.ndim != 3:
            raise ValueError("logits must have shape [batch, seq, vocab].")
        if self.input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, seq].")
        if self.logits.shape[:2] != self.input_ids.shape:
            raise ValueError("logits and input_ids must agree on [batch, seq].")
        if self.attention_mask is not None and self.attention_mask.shape != self.input_ids.shape:
            raise ValueError("attention_mask must match input_ids.")
        if self.temperature <= 0:
            raise ValueError("temperature must be positive.")

    @property
    def log_probs(self) -> Tensor:
        """Log-probabilities derived lazily from logits."""

        if self._log_probs is None:
            self._log_probs = logits_to_log_probs(self.logits)
        return self._log_probs

    @property
    def probs(self) -> Tensor:
        """Probabilities derived lazily from logits."""

        if self._probs is None:
            self._probs = self.log_probs.exp()
        return self._probs

    def final_token_logits(self) -> Tensor:
        """Return logits at the final valid token position for each sequence."""

        return select_final_token(self.logits, attention_mask=self.attention_mask)

    def final_token_log_probs(self) -> Tensor:
        """Return log-probabilities at the final valid token position."""

        return select_final_token(self.log_probs, attention_mask=self.attention_mask)

    def final_token_probs(self) -> Tensor:
        """Return probabilities at the final valid token position."""

        return select_final_token(self.probs, attention_mask=self.attention_mask)

    def gather_log_probs(self, token_ids: Tensor) -> Tensor:
        """Gather log-probabilities for selected token ids."""

        return gather_token_scores(self.log_probs, token_ids)
