"""Immutable canonical snapshot and active check-in lineage resolution."""

from __future__ import annotations

import hashlib
import json
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import date, datetime
from typing import Final, TypeVar

from .models import ContractStatus, Event, EventType
from .store import CanonicalEventSnapshot
from .weekly_operations import CanonicalPin

_ROOT_TYPES: Final = frozenset((EventType.NUTRITION_CHECKIN, EventType.MORNING_CHECKIN))
_RowValue = TypeVar("_RowValue")


@dataclass(frozen=True, slots=True)
class InvalidCanonical:
    pass


@dataclass(frozen=True, slots=True)
class InvalidLineage:
    pass


@dataclass(frozen=True, slots=True)
class MissingLineage:
    pin: CanonicalPin


@dataclass(frozen=True, slots=True)
class SelectedLineage:
    pin: CanonicalPin
    root: Event
    source: Event
    source_digest: str


LineageResolution = InvalidCanonical | InvalidLineage | MissingLineage | SelectedLineage


def _typed_row(row: Mapping[str, _RowValue]) -> dict[str, str | int] | None:
    result: dict[str, str | int] = {}
    for key, value in row.items():
        if not isinstance(value, (str, int)):
            return None
        result[key] = value
    return result


def canonical_prefix_pin(snapshot: CanonicalEventSnapshot, sequence: int) -> CanonicalPin:
    """Digest one already-validated canonical sequence prefix."""
    rows = tuple(_typed_row(row) for row in snapshot.sequence_rows[:sequence])
    if any(row is None for row in rows):
        raise AssertionError("validated canonical row became invalid")
    payload = {"sequence": sequence, "rows": [row for row in rows if row is not None]}
    encoded = json.dumps(
        payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")
    ).encode()
    return CanonicalPin(sequence, hashlib.sha256(encoded).hexdigest())


def _snapshot_pin(snapshot: CanonicalEventSnapshot) -> CanonicalPin | None:
    if len(snapshot.events) != len(snapshot.sequence_rows):
        return None
    for expected, (event, row) in enumerate(
        zip(snapshot.events, snapshot.sequence_rows, strict=True), start=1
    ):
        typed_row = _typed_row(row)
        if typed_row is None:
            return None
        body = dict(typed_row)
        supplied = body.pop("row_digest", None)
        encoded = json.dumps(
            body, ensure_ascii=False, sort_keys=True, separators=(",", ":")
        ).encode()
        event_bytes = (event.model_dump_json(exclude_none=True) + "\n").encode()
        valid = (
            row.get("schema_version") == "canonical_sequence_v1"
            and row.get("sequence") == expected
            and row.get("event_id") == event.event_id
            and row.get("event_digest") == hashlib.sha256(event_bytes).hexdigest()
            and supplied == hashlib.sha256(encoded).hexdigest()
        )
        if not valid:
            return None
    return canonical_prefix_pin(snapshot, len(snapshot.sequence_rows))


def _trace(tip: Event, by_id: Mapping[str, Event], kst_day: date) -> tuple[Event, ...] | None:
    lineage: list[Event] = []
    seen: set[str] = set()
    current = tip
    while True:
        if current.event_id in seen or current.status is not ContractStatus.ACCEPTED:
            return None
        if current.safety is not None or current.check_in is None:
            return None
        seen.add(current.event_id)
        lineage.append(current)
        if current.supersedes is None:
            break
        if current.event_type is not EventType.CORRECTION:
            return None
        parent = by_id.get(current.supersedes)
        if parent is None:
            return None
        current = parent
    root = lineage[-1]
    if root.event_type not in _ROOT_TYPES:
        return None
    try:
        occurred = datetime.fromisoformat(root.occurred_at_kst)
    except ValueError:
        return None
    return tuple(lineage) if occurred.date() == kst_day else None


def resolve_canonical_lineage(
    snapshot: CanonicalEventSnapshot, kst_day: date
) -> LineageResolution:
    """Resolve exactly one active accepted nutrition-or-morning lineage."""
    pin = _snapshot_pin(snapshot)
    if pin is None:
        return InvalidCanonical()
    by_id = {event.event_id: event for event in snapshot.events}
    superseded = {event.supersedes for event in snapshot.events if event.supersedes is not None}
    lineages = tuple(
        lineage
        for event in snapshot.events
        if event.event_id not in superseded
        and event.event_type in (*_ROOT_TYPES, EventType.CORRECTION)
        and (lineage := _trace(event, by_id, kst_day)) is not None
    )
    if not lineages:
        return MissingLineage(pin)
    if len(lineages) > 1:
        return InvalidLineage()
    lineage = next(iter(lineages))
    source = lineage[0]
    source_row = next(
        row for row in snapshot.sequence_rows if row.get("event_id") == source.event_id
    )
    digest = source_row.get("event_digest")
    return (
        SelectedLineage(pin, lineage[-1], source, digest)
        if isinstance(digest, str)
        else InvalidLineage()
    )