"""Tracker-side artifact and runtime recipe identity (DGR-003).

The node computes a compatibility fingerprint in `meshnet_node.runtime_recipe`.
This module recomputes it, and the recomputation is the entire point.

A fingerprint that arrives on the wire is a **claim**. If the tracker stored
what a node asserted, a node could assert the digest of a certified recipe while
running an uncertified one — and admission, route selection, and the gRPC
handshake would all wave it through, because they all compare the digest it
handed them. So the tracker derives the digest from the *axes* the node
declared, and refuses any report whose claim does not match the derivation. A
node can lie about its axes, and a real forward will catch that; it cannot lie
about the digest of the axes it declared.

This module deliberately does **not** import `meshnet_node`: the tracker package
does not depend on the node package, and an admission gate that shares an
implementation with the thing it admits is not an independent check. The two
implementations are pinned together by a committed conformance vector
(``tests/data/recipe_fingerprint_vectors.json``). If they drift, a test fails —
rather than a route quietly failing to form, or worse, quietly forming.

Recipes arrive **dark**: registered, visible to an operator, and not routable for
user traffic. Only a real distributed forward — at least two distinct nodes,
covering the whole model, emitting real tokens — takes a recipe out of the dark
(:class:`CertificationLedger`).
"""

from __future__ import annotations

import hashlib
import json
import re
import time
from dataclasses import dataclass, field
from typing import Any, Iterable, Mapping, Sequence

# Layout of the identity block this tracker reads (meshnet_node.runtime_recipe).
RECIPE_IDENTITY_SCHEMA_VERSION = 1

# Domain separation, byte-for-byte identical to the node's. These strings are
# part of the wire contract, not an implementation detail: changing one here
# without changing it there silently partitions every route.
ARTIFACT_DIGEST_DOMAIN = "meshnet.model-artifact.v1"
RECIPE_DIGEST_DOMAIN = "meshnet.runtime-recipe.v1"

# The axes a recipe digest commits to. Order is irrelevant (the canonical JSON
# sorts keys); membership is not — an axis missing here is an axis the tracker
# would let a node change without changing its identity.
RECIPE_AXES: tuple[str, ...] = (
    "weight_quantization",
    "activation_dtype",
    "compute_dtype",
    "kv_dtype",
    "kv_layout",
    "tokenizer_revision",
    "architecture_adapter",
    "backend_id",
    "runtime_version",
    "boundary_schema_version",
    "protocol_schema_version",
)

_INT_AXES = frozenset({"boundary_schema_version", "protocol_schema_version"})

MISMATCH_ARTIFACT = "artifact"
MISMATCH_TOKENIZER = "tokenizer"
MISMATCH_ARCHITECTURE = "architecture"
MISMATCH_RANGE = "range"
MISMATCH_BOUNDARY_SCHEMA = "boundary-schema"
MISMATCH_ACTIVATION_RECIPE = "activation-recipe"
MISMATCH_CACHE_LAYOUT = "cache-layout"
MISMATCH_FINGERPRINT = "fingerprint"
MISMATCH_UNCERTIFIED = "uncertified"

STATUS_UNKNOWN = "unknown"
STATUS_DARK = "dark"
STATUS_CERTIFIED = "certified"

MIN_CERTIFYING_NODES = 2

_HEX64 = re.compile(r"^[0-9a-f]{64}$")
_MOVING_REFS = frozenset({"main", "master", "head", "latest", "dev", "trunk"})


class RecipeIdentityError(ValueError):
    """A presented identity block is malformed or internally inconsistent."""


class FingerprintMismatch(RecipeIdentityError):
    """A supplied digest is inconsistent with its identity declaration.

    A digest is an integrity and compatibility claim, not an authenticated
    statement about what a node is actually executing. Distributed
    certification remains the trust boundary.
    """


def canonical_sha256(value: Any) -> str:
    payload = json.dumps(
        value, sort_keys=True, separators=(",", ":"), ensure_ascii=False
    )
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def _digest(domain: str, body: Mapping[str, Any]) -> str:
    return canonical_sha256({"domain": domain, "body": dict(body)})


def _text(value: Any, what: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise RecipeIdentityError(f"{what!r} must be a non-empty string")
    return value


def _hex64(value: Any, what: str) -> str:
    text = _text(value, what)
    if not _HEX64.match(text):
        raise RecipeIdentityError(f"{what!r} must be a SHA-256 hex digest")
    return text


def _integer(value: Any, what: str, minimum: int) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise RecipeIdentityError(f"{what!r} must be an integer")
    if value < minimum:
        raise RecipeIdentityError(f"{what!r} must be >= {minimum}")
    return value


def _pin(value: Any, what: str) -> str:
    text = _text(value, what)
    lowered = text.strip().lower()
    if lowered in _MOVING_REFS or lowered.startswith("refs/"):
        raise RecipeIdentityError(
            f"{what!r} is a moving reference, not an exact revision pin"
        )
    return text


def _mapping(value: Any, what: str) -> Mapping[str, Any]:
    if not isinstance(value, Mapping):
        raise RecipeIdentityError(f"{what!r} must be a JSON object")
    return value


def artifact_digest(
    *,
    source_digest: str,
    architecture: str,
    architecture_digest: str,
    layer_count: int,
) -> str:
    """The artifact half of the fingerprint, derived exactly as the node derives it."""
    return _digest(
        ARTIFACT_DIGEST_DOMAIN,
        {
            "source_digest": source_digest,
            "architecture": architecture,
            "architecture_digest": architecture_digest,
            "layer_count": layer_count,
        },
    )


def recipe_digest(axes: Mapping[str, Any]) -> str:
    """The recipe half of the fingerprint, derived exactly as the node derives it."""
    missing = [axis for axis in RECIPE_AXES if axis not in axes]
    if missing:
        raise RecipeIdentityError(
            "recipe is missing required axes: " + ", ".join(missing)
        )
    return _digest(RECIPE_DIGEST_DOMAIN, {axis: axes[axis] for axis in RECIPE_AXES})


@dataclass(frozen=True)
class PresentedIdentity:
    """One node's declared artifact/recipe identity, with digests re-derived here.

    `model_artifact_digest` and `runtime_recipe_digest` are computed by this
    tracker from the declared axes. They are never copied from the report.
    """

    artifact_id: str
    revision: str
    source_digest: str
    architecture: str
    architecture_digest: str
    layer_count: int
    is_derivative: bool
    shard_start: int
    shard_end: int
    axes: Mapping[str, Any]
    recipe_id: str
    recipe_version: str
    catalogue_version: str

    @property
    def model_artifact_digest(self) -> str:
        return artifact_digest(
            source_digest=self.source_digest,
            architecture=self.architecture,
            architecture_digest=self.architecture_digest,
            layer_count=self.layer_count,
        )

    @property
    def runtime_recipe_digest(self) -> str:
        return recipe_digest(self.axes)

    @property
    def key(self) -> tuple[str, str]:
        return (self.model_artifact_digest, self.runtime_recipe_digest)

    @property
    def tokenizer_revision(self) -> str:
        return str(self.axes["tokenizer_revision"])

    @property
    def architecture_adapter(self) -> str:
        return str(self.axes["architecture_adapter"])

    def fingerprint_dict(self) -> dict:
        return {
            "model_artifact_digest": self.model_artifact_digest,
            "runtime_recipe_digest": self.runtime_recipe_digest,
            "recipe_id": self.recipe_id,
            "recipe_version": self.recipe_version,
            "catalogue_version": self.catalogue_version,
        }


def parse_identity(data: Any) -> PresentedIdentity:
    """Parse and *verify* a node's identity block.

    Raises when the block is malformed, when a split artifact does not name the
    exact source it was cut from, when a Shard advertises layers its artifact
    does not contain, or when the declared fingerprint does not match the digest
    of the axes it was declared with.
    """
    doc = _mapping(data, "identity")

    schema_version = doc.get("schema_version")
    if schema_version != RECIPE_IDENTITY_SCHEMA_VERSION:
        raise RecipeIdentityError(
            f"identity block declares schema version {schema_version!r}; this tracker "
            f"reads version {RECIPE_IDENTITY_SCHEMA_VERSION}"
        )

    artifact = _mapping(doc.get("artifact"), "identity.artifact")
    recipe = _mapping(doc.get("recipe"), "identity.recipe")

    layer_count = _integer(artifact.get("layer_count"), "artifact.layer_count", 1)
    content_digest = _hex64(artifact.get("content_digest"), "artifact.content_digest")

    raw_binding = artifact.get("derived_from")
    if raw_binding is None:
        source_digest = content_digest
        binding: tuple[int, int] | None = None
    else:
        derived = _mapping(raw_binding, "artifact.derived_from")
        source_digest = _hex64(
            derived.get("source_artifact_digest"),
            "artifact.derived_from.source_artifact_digest",
        )
        binding = (
            _integer(derived.get("shard_start"), "derived_from.shard_start", 0),
            _integer(derived.get("shard_end"), "derived_from.shard_end", 1),
        )
        if binding[1] <= binding[0]:
            raise RecipeIdentityError(
                "'derived_from' covers no layers; an empty split proves nothing"
            )
        if binding[1] > layer_count:
            raise RecipeIdentityError(
                f"'derived_from' claims layers {binding[0]}–{binding[1]}, but the "
                f"source model has only {layer_count} layers"
            )
        if source_digest == content_digest:
            raise RecipeIdentityError(
                "a split artifact's source digest equals its own content digest; "
                "an artifact is not a split of itself"
            )

    shard_start = _integer(doc.get("shard_start"), "shard_start", 0)
    shard_end = _integer(doc.get("shard_end"), "shard_end", 1)
    if shard_end <= shard_start:
        raise RecipeIdentityError("a Shard owning no layer computes nothing")
    if shard_end > layer_count:
        raise RecipeIdentityError(
            f"Shard owns layers {shard_start}–{shard_end}, but the artifact has only "
            f"{layer_count} layers"
        )
    if binding is not None and not (
        binding[0] <= shard_start and shard_end <= binding[1]
    ):
        raise RecipeIdentityError(
            f"Shard advertises layers {shard_start}–{shard_end}, but its split "
            f"artifact only contains layers {binding[0]}–{binding[1]}"
        )

    axes: dict[str, Any] = {}
    for axis in RECIPE_AXES:
        if axis not in recipe:
            raise RecipeIdentityError(
                f"recipe is missing axis {axis!r}; an unstated axis cannot default"
            )
        value = recipe[axis]
        if axis in _INT_AXES:
            axes[axis] = _integer(value, f"recipe.{axis}", 1)
        else:
            axes[axis] = _text(value, f"recipe.{axis}")
    _pin(axes["tokenizer_revision"], "recipe.tokenizer_revision")

    identity = PresentedIdentity(
        artifact_id=_text(artifact.get("artifact_id"), "artifact.artifact_id"),
        revision=_pin(artifact.get("revision"), "artifact.revision"),
        source_digest=source_digest,
        architecture=_text(artifact.get("architecture"), "artifact.architecture"),
        architecture_digest=_hex64(
            artifact.get("architecture_digest"), "artifact.architecture_digest"
        ),
        layer_count=layer_count,
        is_derivative=binding is not None,
        shard_start=shard_start,
        shard_end=shard_end,
        axes=axes,
        recipe_id=_text(recipe.get("recipe_id"), "recipe.recipe_id"),
        recipe_version=_text(recipe.get("recipe_version"), "recipe.recipe_version"),
        catalogue_version=_text(
            recipe.get("catalogue_version"), "recipe.catalogue_version"
        ),
    )

    declared = doc.get("fingerprint")
    if declared is not None:
        claim = _mapping(declared, "identity.fingerprint")
        claimed_artifact = _hex64(
            claim.get("model_artifact_digest"), "fingerprint.model_artifact_digest"
        )
        claimed_recipe = _hex64(
            claim.get("runtime_recipe_digest"), "fingerprint.runtime_recipe_digest"
        )
        if (claimed_artifact, claimed_recipe) != identity.key:
            raise FingerprintMismatch(
                "declared fingerprint is inconsistent with the artifact and recipe "
                "claim; the tracker recomputes compatibility digests"
            )

    return identity


@dataclass(frozen=True)
class RouteMismatch:
    reason: str
    detail: str

    def describe(self) -> str:
        return f"{self.reason}: {self.detail}"


def compare(
    route: PresentedIdentity, candidate: PresentedIdentity
) -> tuple[RouteMismatch, ...]:
    """Why `candidate` may not join a route already running `route`."""
    if route.key == candidate.key:
        return ()

    reasons: list[RouteMismatch] = []
    if route.source_digest != candidate.source_digest:
        reasons.append(
            RouteMismatch(
                MISMATCH_ARTIFACT,
                f"serves Model Artifact {candidate.source_digest[:12]}…, the route "
                f"runs {route.source_digest[:12]}…",
            )
        )
    if route.architecture_digest != candidate.architecture_digest:
        reasons.append(
            RouteMismatch(
                MISMATCH_ARCHITECTURE,
                "architecture/config snapshot differs from the route's",
            )
        )
    if route.layer_count != candidate.layer_count:
        reasons.append(
            RouteMismatch(
                MISMATCH_ARTIFACT,
                f"artifact has {candidate.layer_count} layers, the route's has "
                f"{route.layer_count}",
            )
        )

    axis_reason = {
        "tokenizer_revision": MISMATCH_TOKENIZER,
        "architecture_adapter": MISMATCH_ARCHITECTURE,
        "activation_dtype": MISMATCH_ACTIVATION_RECIPE,
        "compute_dtype": MISMATCH_ACTIVATION_RECIPE,
        "kv_dtype": MISMATCH_CACHE_LAYOUT,
        "kv_layout": MISMATCH_CACHE_LAYOUT,
        "boundary_schema_version": MISMATCH_BOUNDARY_SCHEMA,
        "protocol_schema_version": MISMATCH_BOUNDARY_SCHEMA,
    }
    for axis in RECIPE_AXES:
        mine = route.axes.get(axis)
        theirs = candidate.axes.get(axis)
        if mine != theirs:
            reasons.append(
                RouteMismatch(
                    axis_reason.get(axis, MISMATCH_FINGERPRINT),
                    f"{axis} is {theirs!r}, the route runs {mine!r}",
                )
            )

    if not reasons:
        reasons.append(
            RouteMismatch(
                MISMATCH_FINGERPRINT,
                "fingerprints differ on a field this comparison does not cover",
            )
        )
    return tuple(reasons)


def coverage_gap(
    ranges: Iterable[tuple[int, int]], layer_count: int
) -> str | None:
    """Why `ranges` fail to tile ``[0, layer_count)``, or None when they do.

    Overlaps are legal — ADR-0012 lets the Tracker resolve one by telling a hop
    where the previous hop stopped. A hole is not: its layers are never computed,
    and the route emits fluent tokens from a truncated model.
    """
    ordered = sorted(ranges)
    if not ordered:
        return "the route owns no layers"
    if ordered[0][0] != 0:
        return f"layers 0–{ordered[0][0]} are owned by no Shard on the route"

    covered = 0
    for start, end in ordered:
        if start > covered:
            return f"layers {covered}–{start} are owned by no Shard on the route"
        covered = max(covered, end)

    if covered < layer_count:
        return f"layers {covered}–{layer_count} are owned by no Shard on the route"
    if covered > layer_count:
        return (
            f"the route claims {covered} layers, but the model has only {layer_count}"
        )
    return None


@dataclass(frozen=True)
class DistributedForwardEvidence:
    """A real distributed forward — the only thing that certifies a recipe."""

    route_session_id: str
    route_epoch: int
    node_ids: tuple[str, ...]
    shard_ranges: tuple[tuple[int, int], ...]
    tokens_generated: int
    layer_count: int
    fingerprint: tuple[str, str]
    synthetic: bool = False
    certified_at: float = field(default_factory=time.time)

    def rejection(self) -> str | None:
        if self.synthetic:
            return (
                "evidence comes from a synthetic worker; only a real distributed "
                "forward certifies a recipe"
            )
        distinct = set(self.node_ids)
        if len(distinct) < MIN_CERTIFYING_NODES:
            return (
                f"evidence covers {len(distinct)} distinct node(s); a distributed "
                f"forward requires at least {MIN_CERTIFYING_NODES}"
            )
        if len(distinct) != len(self.node_ids):
            return "the same node is counted more than once in the certifying route"
        if self.tokens_generated < 1:
            return "the certifying forward generated no tokens"
        return coverage_gap(self.shard_ranges, self.layer_count)


@dataclass(frozen=True)
class RecipeStatus:
    status: str
    detail: str = ""
    certified_at: float | None = None

    @property
    def may_serve(self) -> bool:
        return self.status == STATUS_CERTIFIED

    @property
    def may_certify(self) -> bool:
        """A dark recipe may be routed to *certify* it, and for nothing else.

        Without this, certification is unreachable by construction: serving needs
        certification, certification needs a real distributed forward, and a real
        distributed forward needs a route.
        """
        return self.status in (STATUS_DARK, STATUS_CERTIFIED)

    def to_dict(self) -> dict:
        return {
            "status": self.status,
            "detail": self.detail,
            "certified_at": self.certified_at,
        }


_DARK_DETAIL = "registered; dark until a real distributed forward certifies it"
_UNKNOWN_DETAIL = "this recipe has never been registered with the tracker"


class CertificationLedger:
    """Which recipes the tracker has seen, and which a real forward has proven."""

    def __init__(self) -> None:
        self._dark: set[tuple[str, str]] = set()
        self._certified: dict[tuple[str, str], RecipeStatus] = {}

    def register(self, identity: PresentedIdentity) -> RecipeStatus:
        key = identity.key
        if key in self._certified:
            return self._certified[key]
        self._dark.add(key)
        return RecipeStatus(STATUS_DARK, _DARK_DETAIL)

    def status(self, identity: PresentedIdentity | tuple[str, str]) -> RecipeStatus:
        key = identity if isinstance(identity, tuple) else identity.key
        if key in self._certified:
            return self._certified[key]
        if key in self._dark:
            return RecipeStatus(STATUS_DARK, _DARK_DETAIL)
        return RecipeStatus(STATUS_UNKNOWN, _UNKNOWN_DETAIL)

    def certify(
        self,
        identity: PresentedIdentity | tuple[str, str],
        evidence: DistributedForwardEvidence,
    ) -> RecipeStatus:
        """Promote a recipe out of the dark, or raise saying why the evidence is short."""
        key = identity if isinstance(identity, tuple) else identity.key
        if key not in self._dark and key not in self._certified:
            raise RecipeIdentityError(
                "this fingerprint is not registered; unknown recipes cannot be certified"
            )
        if evidence.fingerprint != key:
            raise RecipeIdentityError(
                "certification evidence fingerprint does not match the recipe being promoted"
            )
        rejection = evidence.rejection()
        if rejection is not None:
            raise RecipeIdentityError(f"this evidence does not certify: {rejection}")
        status = RecipeStatus(
            STATUS_CERTIFIED,
            (
                f"certified by route session {evidence.route_session_id} across "
                f"{len(set(evidence.node_ids))} nodes"
            ),
            certified_at=evidence.certified_at,
        )
        self._certified[key] = status
        self._dark.discard(key)
        return status

    def to_dict(self) -> dict:
        return {
            "dark": [list(key) for key in sorted(self._dark)],
            "certified": {
                "/".join(key): status.to_dict()
                for key, status in sorted(self._certified.items())
            },
        }


def admit_route(
    identities: Sequence[PresentedIdentity],
    *,
    ledger: CertificationLedger | None = None,
    for_certification: bool = False,
) -> tuple[RouteMismatch, ...]:
    """Every reason these Shards may not form one Inference Route; empty when they may.

    Fail-closed by construction: an empty route, a hole in the coverage, one
    differing dtype, and an uncertified recipe each produce a reason, so a caller
    that proceeds only on an empty result cannot admit any of them.
    """
    if not identities:
        return (RouteMismatch(MISMATCH_RANGE, "a route needs at least one Shard"),)

    head = identities[0]
    reasons: list[RouteMismatch] = []
    for candidate in identities[1:]:
        reasons.extend(compare(head, candidate))

    gap = coverage_gap(
        [(i.shard_start, i.shard_end) for i in identities], head.layer_count
    )
    if gap is not None:
        reasons.append(RouteMismatch(MISMATCH_RANGE, gap))

    if ledger is not None:
        status = ledger.status(head)
        allowed = status.may_certify if for_certification else status.may_serve
        if not allowed:
            reasons.append(
                RouteMismatch(MISMATCH_UNCERTIFIED, f"{status.status}: {status.detail}")
            )

    return tuple(reasons)