"""Adapter LossReport sidecar (ticket #000022).

Typed loss ledger for what bytes the adapter / canonicalizer dropped,
transformed, or normalized on the way from raw bytes to the prose the
model sees. SIDECAR ONLY: never enters ``cache_key``, ``document_root``,
``run_dag_root``, or the ``audit_events`` chain.

Loss policy lives in its own ``loss_report_policy_hash`` (computed
here) — separate from the QA ``governance_policy_hash``, so toggling
reporting does NOT invalidate prior cache entries. See ticket §3.7.

Public surface:

    LossEvent       — frozen dataclass, one per (loss_kind, chunk, stage).
    LossCollector   — accumulates per-kind aggregates during a normalizer
                      pass; merging duplicate kinds inside one chunk.
    record_losses   — batched, idempotent persistence.
    compute_loss_report_policy_hash — pure hash of the three policy fields.

The module has zero non-stdlib deps so callers in adapters / wikitext can
import it without dragging the wikitext / html extras.
"""

from __future__ import annotations

import hashlib
import json
import sqlite3
import time
from dataclasses import dataclass, field
from typing import Iterable, Optional


LOSS_REPORT_VERSION = "loss-report-v1"

DEFAULT_MAX_EXCERPT_BYTES = 200

VALID_LOSS_MODES = frozenset({"pure_drop", "transform", "quarantine", "normalize"})


@dataclass(frozen=True)
class LossEvent:
    """One aggregated loss event for a (chunk, stage, kind) tuple.

    ``bytes_dropped`` semantics depend on ``loss_mode``:

      pure_drop / quarantine — bytes removed from the prose stream;
                               participates in the byte-conservation
                               property test (§3.8).
      transform / normalize  — byte delta only; excluded from the
                               property test because content was
                               rewritten, not deleted.

    ``sample_excerpt`` is None when the policy disables excerpts;
    ``sample_hash`` stays populated either way so loss-signature
    recurrence remains detectable in PII-safe mode.
    """

    stage: str
    canonicalization_version: str
    loss_kind: str
    loss_mode: str
    bytes_dropped: int
    occurrence_count: int
    input_length_bytes: Optional[int] = None
    output_length_bytes: Optional[int] = None
    sample_excerpt: Optional[str] = None
    sample_hash: Optional[bytes] = None
    adapter_name: Optional[str] = None
    adapter_version: Optional[str] = None

    def __post_init__(self) -> None:
        if self.loss_mode not in VALID_LOSS_MODES:
            raise ValueError(
                f"loss_mode {self.loss_mode!r} not in {sorted(VALID_LOSS_MODES)}"
            )


def _truncate(s: str, limit: int) -> str:
    """UTF-8-aware byte truncation. Caller passes a byte limit; we return
    a string that encodes to <= limit bytes. Cuts on character boundary
    so the excerpt remains decodable."""
    encoded = s.encode("utf-8", errors="surrogatepass")
    if len(encoded) <= limit:
        return s
    cut = encoded[:limit]
    while cut:
        try:
            return cut.decode("utf-8")
        except UnicodeDecodeError:
            cut = cut[:-1]
    return ""


def _sha256_bytes(data: bytes) -> bytes:
    return hashlib.sha256(data).digest()


@dataclass
class LossCollector:
    """Per-pass accumulator. One collector per (chunk, normalizer-pass).

    ``add()`` is the hot path; idempotent only at the (stage, kind) level
    — duplicate adds for the same kind merge by summing ``bytes_dropped``
    and ``occurrence_count``. The first sample seen wins for excerpt /
    hash (cheap; matches the §2.1 "representative excerpt" wording).

    ``events()`` returns the frozen `LossEvent` list once the pass
    completes. Subsequent ``add()`` after ``events()`` is allowed
    (collector stays mutable) — useful when one collector spans multiple
    micro-passes within a single canonicalizer.
    """

    excerpts_enabled: bool = True
    max_excerpt_bytes: int = DEFAULT_MAX_EXCERPT_BYTES
    adapter_name: Optional[str] = None
    adapter_version: Optional[str] = None
    input_length_bytes: Optional[int] = None
    output_length_bytes: Optional[int] = None
    _agg: dict[tuple[str, str, str], dict] = field(default_factory=dict)

    def add(
        self,
        *,
        stage: str,
        canonicalization_version: str,
        loss_kind: str,
        loss_mode: str,
        dropped: str | bytes,
        occurrences: int = 1,
    ) -> None:
        if loss_mode not in VALID_LOSS_MODES:
            raise ValueError(
                f"loss_mode {loss_mode!r} not in {sorted(VALID_LOSS_MODES)}"
            )
        if isinstance(dropped, str):
            dropped_bytes = dropped.encode("utf-8", errors="surrogatepass")
        else:
            dropped_bytes = dropped
        n_bytes = len(dropped_bytes)

        key = (stage, canonicalization_version, loss_kind)
        slot = self._agg.get(key)
        if slot is None:
            sample_text: Optional[str]
            try:
                decoded = dropped_bytes.decode("utf-8", errors="replace")
            except Exception:
                decoded = ""
            if self.excerpts_enabled:
                sample_text = _truncate(decoded, self.max_excerpt_bytes)
            else:
                sample_text = None
            slot = {
                "stage": stage,
                "canonicalization_version": canonicalization_version,
                "loss_kind": loss_kind,
                "loss_mode": loss_mode,
                "bytes_dropped": 0,
                "occurrence_count": 0,
                "sample_excerpt": sample_text,
                "sample_hash": _sha256_bytes(
                    dropped_bytes[: self.max_excerpt_bytes]
                ),
            }
            self._agg[key] = slot
        elif slot["loss_mode"] != loss_mode:
            raise ValueError(
                f"loss_mode mismatch on duplicate add for {key!r}: "
                f"existing={slot['loss_mode']!r} new={loss_mode!r}"
            )
        slot["bytes_dropped"] += n_bytes
        slot["occurrence_count"] += occurrences

    def record_delta(
        self,
        *,
        stage: str,
        canonicalization_version: str,
        loss_kind: str,
        loss_mode: str,
        bytes_delta: int,
        occurrences: int = 1,
    ) -> None:
        """Record a byte-length delta with no associated content sample.

        Use for ``transform`` or ``normalize`` events where bytes were
        rewritten in place (e.g. ``strip_code``, whitespace collapse) and
        there is no single span to capture. ``bytes_delta`` aggregates
        across multiple ``record_delta`` calls for the same kind.
        """
        if loss_mode not in VALID_LOSS_MODES:
            raise ValueError(
                f"loss_mode {loss_mode!r} not in {sorted(VALID_LOSS_MODES)}"
            )
        if bytes_delta <= 0:
            return
        key = (stage, canonicalization_version, loss_kind)
        slot = self._agg.get(key)
        if slot is None:
            slot = {
                "stage": stage,
                "canonicalization_version": canonicalization_version,
                "loss_kind": loss_kind,
                "loss_mode": loss_mode,
                "bytes_dropped": 0,
                "occurrence_count": 0,
                "sample_excerpt": None,
                "sample_hash": None,
            }
            self._agg[key] = slot
        elif slot["loss_mode"] != loss_mode:
            raise ValueError(
                f"loss_mode mismatch on duplicate add for {key!r}: "
                f"existing={slot['loss_mode']!r} new={loss_mode!r}"
            )
        slot["bytes_dropped"] += int(bytes_delta)
        slot["occurrence_count"] += occurrences

    def set_lengths(self, *, input_bytes: int, output_bytes: int) -> None:
        """Record the chunk-level input/output sizes for the byte-
        conservation property test. Called once per pass after
        normalization completes."""
        self.input_length_bytes = input_bytes
        self.output_length_bytes = output_bytes

    def events(self) -> list[LossEvent]:
        """Snapshot the accumulator as a list of frozen `LossEvent`s."""
        out: list[LossEvent] = []
        for slot in self._agg.values():
            out.append(
                LossEvent(
                    stage=slot["stage"],
                    canonicalization_version=slot["canonicalization_version"],
                    loss_kind=slot["loss_kind"],
                    loss_mode=slot["loss_mode"],
                    bytes_dropped=slot["bytes_dropped"],
                    occurrence_count=slot["occurrence_count"],
                    input_length_bytes=self.input_length_bytes,
                    output_length_bytes=self.output_length_bytes,
                    sample_excerpt=slot["sample_excerpt"],
                    sample_hash=slot["sample_hash"],
                    adapter_name=self.adapter_name,
                    adapter_version=self.adapter_version,
                )
            )
        return out


def compute_loss_report_policy_hash(
    *,
    enabled: bool,
    excerpts: bool,
    max_excerpt_bytes: int,
    version: str = LOSS_REPORT_VERSION,
) -> str:
    """Hash the three loss-reporting policy fields. Pure function.

    Returns 64-char hex sha256 over a canonical JSON. Stored on each
    ``adapter_loss_reports`` row so an auditor can correlate policy
    changes with loss data. NOT folded into ``governance_policy_hash``.
    """
    body = {
        "version": version,
        "enabled": bool(enabled),
        "excerpts": bool(excerpts),
        "max_excerpt_bytes": int(max_excerpt_bytes),
    }
    canonical = json.dumps(
        body, sort_keys=True, separators=(",", ":"), ensure_ascii=False
    )
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def record_losses(
    conn: sqlite3.Connection,
    *,
    chunk_id: int,
    document_root: str,
    events: Iterable[LossEvent],
    loss_report_policy_hash: str,
    ts: int | None = None,
) -> int:
    """Batched, idempotent insert of `LossEvent`s for one chunk.

    Idempotency: PK is (chunk_id, stage, canonicalization_version,
    loss_kind). Re-running the same normalizer over the same chunk
    is a no-op via INSERT OR REPLACE.

    Caller wraps in a transaction (or this runs inside the ingest
    batch's transaction). Returns the row count touched.
    """
    if ts is None:
        ts = int(time.time())
    rows: list[tuple] = []
    for ev in events:
        rows.append(
            (
                chunk_id,
                document_root,
                ev.stage,
                ev.canonicalization_version,
                ev.loss_kind,
                ev.loss_mode,
                ev.bytes_dropped,
                ev.occurrence_count,
                ev.input_length_bytes,
                ev.output_length_bytes,
                ev.sample_excerpt,
                ev.sample_hash.hex() if ev.sample_hash is not None else None,
                ev.adapter_name,
                ev.adapter_version,
                loss_report_policy_hash,
                ts,
            )
        )
    if not rows:
        return 0
    conn.executemany(
        "INSERT OR REPLACE INTO adapter_loss_reports "
        "(chunk_id, document_root, stage, canonicalization_version, "
        " loss_kind, loss_mode, bytes_dropped, occurrence_count, "
        " input_length_bytes, output_length_bytes, sample_excerpt, "
        " sample_hash, adapter_name, adapter_version, "
        " loss_report_policy_hash, created_at) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        rows,
    )
    return len(rows)
