"""Authoritative identity boundary for a native GGUF Shard backend.

The native loader owns the facts about the mapped artifact.  This module turns
that immutable report and separately pinned deployment inputs into the one
``ShardIdentity`` DGR-003 permits a native worker to emit.  It is intentionally
not an adapter for the legacy Transformers backend: that backend has no
authoritative immutable GGUF artifact pin and must remain identity-free.
"""

from __future__ import annotations

import ctypes
import hashlib
import json
import re
from dataclasses import dataclass, field
from pathlib import Path

from .native_protocol import BUNDLE_VERSION, SCHEMA_VERSION, pb
from .runtime_pin import DEFAULT_LOCK_DIR, RuntimePin, load_runtime_pin
from .runtime_recipe import (
    ArtifactIdentity,
    DerivativeBinding,
    RecipeIdentityError,
    RuntimeRecipe,
    ShardIdentity,
    check_session_open,
    handshake_error,
)

_HEX40 = re.compile(r"^[0-9a-f]{40}$")
_HEX64 = re.compile(r"^[0-9a-f]{64}$")

# The executing-runtime attestation contract.
#
# The repository lock is world-readable, so a Python object holding
# lock-shaped values proves nothing about the runtime that will execute:
# copying `load_runtime_pin()` into a self-report is exactly the forgery
# DGR-025 forbids.  Attestation values are therefore accepted only when
# *extracted from the native artifact itself*, through two channels that must
# agree:
#
# 1. static — the artifact's bytes embed exactly one
#    ``MESHNET-RUNTIME-ATTESTATION.v1:<canonical json>`` marker (NUL
#    terminated).  The DGR-027 CMake ABI-marker lane is where the native
#    build bakes it in from the lock at configure time.
# 2. dynamic — the artifact must actually dlopen, and its exported
#    ``llama_meshnet_runtime_attestation`` symbol must return that same
#    marker.  A marker pasted into a plain file is not an executing runtime.
#
# What this cannot prove: a cross-compiler bit-reproducible binary SHA, or
# that an adversary did not *build* a native artifact that embeds lock-true
# values while lying about its source.  Manufacturing a lying native build is
# a categorically higher bar than authoring a Python dict, and real
# distributed certification (DGR-025's registered-but-dark ledger) remains
# the final backstop behind this boundary.
ATTESTATION_MARKER_PREFIX = b"MESHNET-RUNTIME-ATTESTATION.v1:"
ATTESTATION_SYMBOL = "llama_meshnet_runtime_attestation"

_ATTESTATION_STR_FIELDS = (
    "runtime_name",
    "upstream_commit",
    "patched_tree",
    "patch_stack_digest",
    "build_recipe_digest",
)
_ATTESTATION_INT_FIELDS = ("boundary_schema_version", "protocol_schema_version")

# Module-private capability: evidence can only be minted where an artifact
# was actually read, scanned, loaded, and queried.
_EVIDENCE_TOKEN = object()


def attestation_payload(
    *,
    runtime_name: str,
    upstream_commit: str,
    patched_tree: str,
    patch_stack_digest: str,
    build_recipe_digest: str,
    boundary_schema_version: int,
    protocol_schema_version: int,
) -> bytes:
    """The canonical marker payload for one exact runtime.

    This single encoding is shared by the build lane that embeds the marker,
    the extractor that parses it, and the binding check that ties attestation
    fields to the extracted evidence — so there is exactly one byte string a
    given runtime identity can legitimately embed.
    """
    return json.dumps(
        {
            "runtime_name": runtime_name,
            "upstream_commit": upstream_commit,
            "patched_tree": patched_tree,
            "patch_stack_digest": patch_stack_digest,
            "build_recipe_digest": build_recipe_digest,
            "boundary_schema_version": boundary_schema_version,
            "protocol_schema_version": protocol_schema_version,
        },
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    ).encode("utf-8")


def expected_attestation_payload(
    pin: RuntimePin,
    *,
    boundary_schema_version: int = BUNDLE_VERSION,
    protocol_schema_version: int = int(SCHEMA_VERSION),
) -> bytes:
    """The marker payload a native build of this lock workspace must embed."""
    return attestation_payload(
        runtime_name=pin.runtime_name,
        upstream_commit=pin.upstream_commit,
        patched_tree=pin.patched_tree,
        patch_stack_digest=pin.patch_stack_digest,
        build_recipe_digest=pin.build_recipe_digest,
        boundary_schema_version=boundary_schema_version,
        protocol_schema_version=protocol_schema_version,
    )


@dataclass(frozen=True)
class NativeArtifactEvidence:
    """Proof that attestation values came out of a loadable native artifact.

    ``binary_digest`` pins *which* artifact bytes were attested;
    ``payload_digest`` pins *what* those bytes attested, and is re-derived
    from the attestation's own fields on construction so the values cannot be
    edited after extraction (``dataclasses.replace`` laundering fails).
    Only :func:`attest_loaded_runtime` can mint this object.
    """

    artifact_path: str
    binary_digest: str
    payload_digest: str
    _token: object = field(default=None, repr=False, compare=False)

    def __post_init__(self) -> None:
        if self._token is not _EVIDENCE_TOKEN:
            raise RecipeIdentityError(
                "native artifact evidence can only be minted by "
                "attest_loaded_runtime() from an actually loaded native "
                "artifact; it cannot be authored from repository lock values"
            )
        if not isinstance(self.artifact_path, str) or not self.artifact_path:
            raise RecipeIdentityError("native artifact evidence must name the artifact")
        for field_name in ("binary_digest", "payload_digest"):
            value = getattr(self, field_name)
            if not isinstance(value, str) or not _HEX64.fullmatch(value):
                raise RecipeIdentityError(
                    f"native artifact evidence {field_name!r} must be a 64-hex sha256"
                )


@dataclass(frozen=True)
class NativeRuntimeAttestation:
    """What the *executing* runtime reports about itself, at load time.

    The repository lock says what the runtime is supposed to be; this says
    what the loaded runtime *is* — the source tree it was built from, the
    patch stack compiled into it, the numerically relevant build recipe, and
    the boundary/protocol schema (ABI) it speaks.  The values are never
    accepted from a caller: they must arrive bound to
    :class:`NativeArtifactEvidence`, which only
    :func:`attest_loaded_runtime` can produce by reading, loading, and
    querying the native artifact itself.  Identity construction then compares
    them to the lock/build-derived expectation and refuses on any difference,
    so a worker cannot serve a lock it is not actually running — and cannot
    fake one by copying the lock into a Python self-report.

    Deliberately *not* attested: a compiler-specific binary SHA. Binding the
    recorded build recipe is honest about what the manifest can prove;
    bit-reproducible binary attestation is not claimed.
    """

    runtime_name: str
    upstream_commit: str
    patched_tree: str
    patch_stack_digest: str
    build_recipe_digest: str
    boundary_schema_version: int
    protocol_schema_version: int
    evidence: NativeArtifactEvidence

    def __post_init__(self) -> None:
        if not isinstance(self.evidence, NativeArtifactEvidence):
            raise RecipeIdentityError(
                "runtime attestation values must be extracted from the loaded "
                "native artifact via attest_loaded_runtime(); a Python "
                "self-report carrying copied lock values is not an attestation"
            )
        if not isinstance(self.runtime_name, str) or not self.runtime_name.strip():
            raise RecipeIdentityError(
                "runtime attestation must name the executing runtime"
            )
        for field_name, pattern, what in (
            ("upstream_commit", _HEX40, "40-hex upstream commit"),
            ("patched_tree", _HEX40, "40-hex patched source tree id"),
            ("patch_stack_digest", _HEX64, "64-hex patch-stack digest"),
            ("build_recipe_digest", _HEX64, "64-hex build-recipe digest"),
        ):
            value = getattr(self, field_name)
            if not isinstance(value, str) or not pattern.fullmatch(value):
                raise RecipeIdentityError(
                    f"runtime attestation {field_name!r} must be an exact {what}"
                )
        for field_name in _ATTESTATION_INT_FIELDS:
            value = getattr(self, field_name)
            if isinstance(value, bool) or not isinstance(value, int) or value < 1:
                raise RecipeIdentityError(
                    f"runtime attestation {field_name!r} must be a positive integer"
                )
        payload = attestation_payload(
            **{name: getattr(self, name) for name in _ATTESTATION_STR_FIELDS},
            **{name: getattr(self, name) for name in _ATTESTATION_INT_FIELDS},
        )
        if hashlib.sha256(payload).hexdigest() != self.evidence.payload_digest:
            raise RecipeIdentityError(
                "runtime attestation fields do not match the attestation "
                "extracted from the native artifact; refusing values edited "
                "after extraction"
            )


def _parse_attestation_payload(payload: bytes) -> dict[str, object]:
    """Strictly parse one embedded marker payload, or refuse."""
    try:
        doc = json.loads(payload.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise RecipeIdentityError(
            f"embedded runtime attestation marker is not valid JSON: {exc}"
        ) from exc
    expected_keys = set(_ATTESTATION_STR_FIELDS) | set(_ATTESTATION_INT_FIELDS)
    if not isinstance(doc, dict) or set(doc) != expected_keys:
        raise RecipeIdentityError(
            "embedded runtime attestation marker must record exactly the "
            "attestation fields"
        )
    for name in _ATTESTATION_STR_FIELDS:
        if not isinstance(doc[name], str):
            raise RecipeIdentityError(
                f"embedded runtime attestation field {name!r} must be a string"
            )
    for name in _ATTESTATION_INT_FIELDS:
        if isinstance(doc[name], bool) or not isinstance(doc[name], int):
            raise RecipeIdentityError(
                f"embedded runtime attestation field {name!r} must be an integer"
            )
    if attestation_payload(**doc) != payload:  # type: ignore[arg-type]
        raise RecipeIdentityError(
            "embedded runtime attestation marker is not in canonical form"
        )
    return doc


def attest_loaded_runtime(artifact_path: Path | str) -> NativeRuntimeAttestation:
    """Extract the executing runtime's attestation from its native artifact.

    Fails closed when the artifact is missing or empty, embeds no attestation
    marker (a runtime built without the attestation lane cannot prove what it
    is), embeds conflicting markers, is not a loadable shared object, does
    not export :data:`ATTESTATION_SYMBOL`, or reports through that symbol
    anything other than the embedded marker.

    The returned attestation is bound to the artifact by its byte digest and
    to the extracted values by the payload digest.  Loading the artifact does
    execute its initializers — this is the same artifact the worker is about
    to run inference with, so that adds no new execution.  An OS-level swap
    of the file between the byte read and the dlopen is a documented
    residual race; distributed certification remains the final backstop.
    """
    path = Path(artifact_path)
    try:
        data = path.read_bytes()
    except FileNotFoundError:
        raise RecipeIdentityError(
            f"native runtime artifact not found at {path}; without the built "
            "native runtime there is no executing identity to attest"
        ) from None
    except OSError as exc:
        raise RecipeIdentityError(
            f"native runtime artifact at {path} is unreadable: {exc}"
        ) from exc
    if not data:
        raise RecipeIdentityError(f"native runtime artifact at {path} is empty")

    payloads: list[bytes] = []
    cursor = 0
    while (start := data.find(ATTESTATION_MARKER_PREFIX, cursor)) >= 0:
        end = data.find(b"\x00", start)
        if end < 0:
            raise RecipeIdentityError(
                "embedded runtime attestation marker is not NUL-terminated"
            )
        payloads.append(data[start + len(ATTESTATION_MARKER_PREFIX) : end])
        cursor = end
    if not payloads:
        raise RecipeIdentityError(
            f"native runtime artifact at {path} embeds no runtime attestation "
            "marker; a runtime built without the attestation lane cannot "
            "prove what it is"
        )
    if len(set(payloads)) != 1:
        raise RecipeIdentityError(
            f"native runtime artifact at {path} embeds conflicting runtime "
            "attestation markers"
        )
    payload = payloads[0]
    doc = _parse_attestation_payload(payload)

    try:
        library = ctypes.CDLL(str(path), mode=ctypes.RTLD_LOCAL)
    except OSError as exc:
        raise RecipeIdentityError(
            f"native runtime artifact at {path} is not a loadable native "
            "artifact; an attestation marker copied into a plain file is not "
            "an executing runtime"
        ) from exc
    try:
        symbol = getattr(library, ATTESTATION_SYMBOL)
    except AttributeError:
        raise RecipeIdentityError(
            f"native runtime artifact at {path} does not export "
            f"{ATTESTATION_SYMBOL}; the loaded runtime itself must report "
            "its attestation"
        ) from None
    symbol.restype = ctypes.c_char_p
    symbol.argtypes = []
    reported = symbol()
    if reported != ATTESTATION_MARKER_PREFIX + payload:
        raise RecipeIdentityError(
            f"the runtime loaded from {path} reports a different attestation "
            "than its artifact embeds; refusing an artifact that disagrees "
            "with itself"
        )

    evidence = NativeArtifactEvidence(
        artifact_path=str(path),
        binary_digest=hashlib.sha256(data).hexdigest(),
        payload_digest=hashlib.sha256(payload).hexdigest(),
        _token=_EVIDENCE_TOKEN,
    )
    return NativeRuntimeAttestation(evidence=evidence, **doc)  # type: ignore[arg-type]


@dataclass(frozen=True)
class NativeLoadedArtifactReport:
    """Immutable GGUF facts returned by the loaded native model.

    The report is copied from ``llama_model_meshnet_range_report`` plus the
    parsed GGUF metadata while the model is live.  Byte counts are operational
    evidence rather than compatibility axes, but keeping them beside the range
    prevents a caller from substituting an unverified range declaration.
    ``runtime_attestation`` must be the evidence-bound attestation extracted
    from the loaded native artifact by :func:`attest_loaded_runtime`; a
    report without one cannot be turned into an identity at all, and one
    cannot exist without an actual native artifact to extract it from.
    """

    owned_start_layer: int
    owned_end_layer: int
    mapped_bytes: int
    resident_bytes: int
    registered_bytes: int
    architecture: str
    architecture_digest: str
    layer_count: int
    runtime_attestation: NativeRuntimeAttestation

    def __post_init__(self) -> None:
        if self.owned_start_layer < 0 or self.owned_end_layer <= self.owned_start_layer:
            raise RecipeIdentityError("native report has an invalid owned layer range")
        if self.layer_count < 1 or self.owned_end_layer > self.layer_count:
            raise RecipeIdentityError("native report range is outside GGUF layer metadata")
        if min(self.mapped_bytes, self.resident_bytes, self.registered_bytes) < 0:
            raise RecipeIdentityError("native report byte counts must be non-negative")
        if not isinstance(self.runtime_attestation, NativeRuntimeAttestation):
            raise RecipeIdentityError(
                "native report must carry the executing runtime's attestation"
            )


@dataclass(frozen=True)
class ImmutableArtifactPin:
    """Deployment-supplied immutable bytes pin, never inferred from a model name."""

    artifact_id: str
    revision: str
    content_digest: str
    derived_from: DerivativeBinding | None = None


@dataclass(frozen=True)
class NativeNumericalRecipe:
    """Immutable numerical inputs selected for this native worker instance."""

    weight_quantization: str
    activation_dtype: str
    compute_dtype: str
    kv_dtype: str
    kv_layout: str
    architecture_adapter: str
    backend_id: str
    recipe_id: str
    recipe_version: str
    catalogue_version: str
    boundary_schema_version: int = BUNDLE_VERSION
    protocol_schema_version: int = int(SCHEMA_VERSION)


@dataclass(frozen=True)
class NativeIdentityInputs:
    """Everything a native backend needs to emit one exact identity.

    ``tokenizer_revision`` must be the content-addressed identity computed by
    :func:`meshnet_node.runtime_recipe.tokenizer_identity` over the loaded
    tokenizer/config bytes; identity construction rejects anything else.
    """

    loaded_artifact: NativeLoadedArtifactReport
    artifact_pin: ImmutableArtifactPin
    tokenizer_revision: str
    numerical_recipe: NativeNumericalRecipe


def _require_attested_runtime(
    attested: NativeRuntimeAttestation,
    expected: RuntimePin,
    recipe: NativeNumericalRecipe,
) -> None:
    """Fail closed unless the executing runtime is the locked, built runtime.

    Every comparison is exact and every difference is separately fatal: an
    attestation that agrees on the commit but not the patch stack (or the
    patched tree, or the build recipe, or the ABI) is a different runtime
    wearing the lock's name, and letting it emit the lock's identity is
    exactly the substitution DGR-025 exists to prevent.
    """
    for what, got, want in (
        ("runtime name", attested.runtime_name, expected.runtime_name),
        ("upstream commit", attested.upstream_commit, expected.upstream_commit),
        ("patched source tree", attested.patched_tree, expected.patched_tree),
        (
            "ordered patch stack",
            attested.patch_stack_digest,
            expected.patch_stack_digest,
        ),
        (
            "build recipe",
            attested.build_recipe_digest,
            expected.build_recipe_digest,
        ),
    ):
        if got != want:
            raise RecipeIdentityError(
                f"the executing runtime's attested {what} does not match the "
                "lock/build-derived expectation; refusing to emit an identity "
                "for a runtime this node is not provably running"
            )
    for what, got, want in (
        (
            "boundary schema",
            attested.boundary_schema_version,
            recipe.boundary_schema_version,
        ),
        (
            "protocol schema",
            attested.protocol_schema_version,
            recipe.protocol_schema_version,
        ),
    ):
        if got != want:
            raise RecipeIdentityError(
                f"the executing runtime's attested {what} version ({got}) does "
                f"not match the recipe's ({want}); an ABI the runtime does not "
                "actually speak cannot be part of its identity"
            )


def shard_identity_from_native_report(
    inputs: NativeIdentityInputs,
    *,
    lock_dir: Path = DEFAULT_LOCK_DIR,
) -> ShardIdentity:
    """Derive identity only from the native report and immutable pinned inputs.

    The ``runtime_version`` axis is never accepted from a caller: it is derived
    from the committed lock workspace, and the loaded runtime's attestation
    must match that lock/build-derived expectation exactly — otherwise this
    raises and no identity exists to register, admit, or certify.
    """
    report = inputs.loaded_artifact
    pin = inputs.artifact_pin
    recipe = inputs.numerical_recipe
    if recipe.backend_id.strip().lower() not in {"llama.cpp", "llama-cpp"}:
        raise RecipeIdentityError(
            "native llama.cpp identity requires backend_id 'llama.cpp' or 'llama-cpp'"
        )
    runtime_pin = load_runtime_pin(lock_dir)
    _require_attested_runtime(report.runtime_attestation, runtime_pin, recipe)
    # The lock-derived prefix identifies the intended source/patch/build
    # recipe.  The executing artifact digest identifies the bytes that actually
    # supplied the attestation.  Without this suffix, any independently built
    # shared object could copy the public lock values into its marker and claim
    # the exact same compatibility identity as the certified artifact.
    runtime_version = (
        f"{runtime_pin.runtime_version}"
        f"+artifact.{report.runtime_attestation.evidence.binary_digest}"
    )
    artifact = ArtifactIdentity(
        artifact_id=pin.artifact_id,
        revision=pin.revision,
        content_digest=pin.content_digest,
        architecture=report.architecture,
        architecture_digest=report.architecture_digest,
        layer_count=report.layer_count,
        derived_from=pin.derived_from,
    )
    return ShardIdentity(
        artifact=artifact,
        recipe=RuntimeRecipe(
            weight_quantization=recipe.weight_quantization,
            activation_dtype=recipe.activation_dtype,
            compute_dtype=recipe.compute_dtype,
            kv_dtype=recipe.kv_dtype,
            kv_layout=recipe.kv_layout,
            tokenizer_revision=inputs.tokenizer_revision,
            architecture_adapter=recipe.architecture_adapter,
            backend_id=recipe.backend_id,
            runtime_version=runtime_version,
            boundary_schema_version=recipe.boundary_schema_version,
            protocol_schema_version=recipe.protocol_schema_version,
            recipe_id=recipe.recipe_id,
            recipe_version=recipe.recipe_version,
            catalogue_version=recipe.catalogue_version,
        ),
        shard_start=report.owned_start_layer,
        shard_end=report.owned_end_layer,
    )


class NativeSessionRejected(RecipeIdentityError):
    """A native worker refused a ``SessionOpen`` before allocating session state."""

    def __init__(self, error: "pb.ShardError") -> None:
        super().__init__(f"native SessionOpen rejected: {error.code}")
        self.error = error


class NativeWorkerBackendAdapter:
    """Small backend-facing adapter around the native loaded-artifact seam."""

    def __init__(self, identity_inputs: NativeIdentityInputs) -> None:
        self.identity_inputs = identity_inputs
        self.identity = shard_identity_from_native_report(identity_inputs)

    @property
    def loaded_artifact_report(self) -> NativeLoadedArtifactReport:
        return self.identity_inputs.loaded_artifact

    def check_session_open(
        self,
        opened: "pb.SessionOpen",
        *,
        expected_route_session_id: str | None = None,
        expected_route_epoch: int | None = None,
    ) -> None:
        """Reject incompatible/stale opens at the native worker boundary."""
        mismatches = check_session_open(
            self.identity,
            opened,
            expected_route_session_id=expected_route_session_id,
            expected_route_epoch=expected_route_epoch,
        )
        error = handshake_error(mismatches)
        if error is not None:
            raise NativeSessionRejected(error)

    def on_session_open(
        self,
        opened: "pb.SessionOpen",
        *,
        expected_route_session_id: str | None = None,
        expected_route_epoch: int | None = None,
    ) -> "pb.SessionAccepted":
        """The native worker's SessionOpen boundary, before session allocation."""
        self.check_session_open(
            opened,
            expected_route_session_id=expected_route_session_id,
            expected_route_epoch=expected_route_epoch,
        )
        return pb.SessionAccepted(
            schema_version=SCHEMA_VERSION,
            route_session_id=opened.route_session_id,
            route_epoch=opened.route_epoch,
            fingerprint=self.identity.fingerprint.to_proto(),
        )