"""Fail-closed admission: no routable registration without a fresh matching proof.

This module does not *produce* proof — `doctor` does that, by pushing a bounded
real forward through the selected shard (NCA-002). This module *decides whether a
proof covers what is about to be advertised*, and startup calls it immediately
before it registers with the tracker.

A capability report proves one combination: model artifact, shard range, recipe,
backend and device. Reusing it for anything else is the exact hole this closes —
a report that failed, aged out, or describes a different model, shard, recipe or
device is rejected here, and the node exits without ever registering an endpoint.

Nothing in here branches on a model, vendor or kernel name: identity fields are
opaque labels that are compared, never interpreted.
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any, Callable

from . import __version__ as _PACKAGE_VERSION
from .capability import CapabilityReport, config_fingerprint
from .doctor import DoctorSelection
from .recipe_manifest import Recipe, RecipeManifest
from .runtime_recipe import (
    build_artifact_identity,
    build_runtime_recipe_identity,
    compatibility_fingerprint,
    fingerprint_payload,
)
from .gguf_ownership import authoritative_dense_llama_ownership

# How long a passing report stays usable. Startup normally validates in-process
# (age ≈ 0); this bounds how far a report written by an earlier `doctor` run can
# be carried forward, after which the hardware, drivers or weights may have moved.
DEFAULT_MAX_REPORT_AGE_SECONDS = 900.0

# A report timestamped this far in the future is not fresh, it is wrong.
_MAX_CLOCK_SKEW_SECONDS = 60.0

REASON_NO_REPORT = "no-report"
REASON_NOT_PASSED = "not-passed"
REASON_STALE = "stale"
REASON_MODEL_MISMATCH = "model-mismatch"
REASON_SHARD_MISMATCH = "shard-mismatch"
REASON_RECIPE_MISMATCH = "recipe-mismatch"
REASON_BACKEND_MISMATCH = "backend-mismatch"
REASON_COMPATIBILITY_MISMATCH = "compatibility-mismatch"


class CapabilityAdmissionError(RuntimeError):
    """This node may not advertise the selection: the proof does not cover it."""

    def __init__(self, reason: str, message: str) -> None:
        super().__init__(message)
        self.reason = reason


@dataclass(frozen=True)
class CapabilityContext:
    """What is about to be advertised, and the loaded backend that would serve it."""

    backend: Any
    selection: DoctorSelection
    recipe: Recipe
    manifest: RecipeManifest
    device: str


# A validator turns the context into the report the gate then judges. Production
# uses `probe_capability`; tests pass an explicit test-safe one (see
# `meshnet_node.testing`) rather than switching this module into a lenient mode.
CapabilityValidator = Callable[[CapabilityContext], CapabilityReport]


@dataclass(frozen=True)
class AdmissionRequirement:
    """The one capability a report must prove for this node to register."""

    model_id: str
    shard_start: int
    shard_end: int
    recipe_id: str
    recipe_version: str
    backend_id: str
    device: str
    compatibility_fingerprint: str
    max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS

    @classmethod
    def for_context(
        cls,
        context: CapabilityContext,
        *,
        max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS,
    ) -> AdmissionRequirement:
        return cls(
            model_id=context.selection.model_id,
            shard_start=context.selection.shard_start,
            shard_end=context.selection.shard_end,
            recipe_id=context.recipe.id,
            recipe_version=context.recipe.version,
            backend_id=context.recipe.backend_id,
            device=context.device,
            compatibility_fingerprint=_compatibility_fingerprint_for_context(
                context
            ),
            max_age_seconds=max_age_seconds,
        )

    @property
    def shard_label(self) -> str:
        return f"layers {self.shard_start}–{self.shard_end}"


def admit(
    requirement: AdmissionRequirement,
    report: CapabilityReport | None,
    *,
    now: float | None = None,
) -> CapabilityReport:
    """Return `report` if it admits `requirement`; otherwise refuse to register.

    Checks run selection-first, so the operator is told the report is about the
    wrong thing before being told it is old.
    """
    if report is None:
        raise CapabilityAdmissionError(
            REASON_NO_REPORT,
            f"no capability report for {requirement.model_id} "
            f"{requirement.shard_label}: this node has not proven it can serve it",
        )

    if report.model.model_id != requirement.model_id:
        raise _mismatch(
            REASON_MODEL_MISMATCH,
            requirement,
            "model",
            report.model.model_id,
            requirement.model_id,
        )

    if (report.shard.start, report.shard.end) != (
        requirement.shard_start,
        requirement.shard_end,
    ):
        raise _mismatch(
            REASON_SHARD_MISMATCH,
            requirement,
            "shard",
            f"layers {report.shard.start}–{report.shard.end}",
            requirement.shard_label,
        )

    if (report.recipe.recipe_id, report.recipe.recipe_version) != (
        requirement.recipe_id,
        requirement.recipe_version,
    ):
        raise _mismatch(
            REASON_RECIPE_MISMATCH,
            requirement,
            "recipe",
            f"{report.recipe.recipe_id} (v{report.recipe.recipe_version})",
            f"{requirement.recipe_id} (v{requirement.recipe_version})",
        )

    if (report.backend.backend_id, report.backend.device) != (
        requirement.backend_id,
        requirement.device,
    ):
        raise _mismatch(
            REASON_BACKEND_MISMATCH,
            requirement,
            "backend",
            f"{report.backend.backend_id} on {report.backend.device}",
            f"{requirement.backend_id} on {requirement.device}",
        )

    if report.compatibility_fingerprint != requirement.compatibility_fingerprint:
        raise CapabilityAdmissionError(
            REASON_COMPATIBILITY_MISMATCH,
            f"capability proof fingerprint {report.compatibility_fingerprint!r} "
            f"does not match the expected compatibility fingerprint for "
            f"{requirement.model_id} {requirement.shard_label}; the artifact, "
            f"tokenizer, architecture, boundary schema, activation recipe or "
            f"cache layout differs",
        )

    if not report.passed:
        raise CapabilityAdmissionError(
            REASON_NOT_PASSED,
            f"capability validation {report.status} for {requirement.model_id} "
            f"{requirement.shard_label} with recipe {requirement.recipe_id}"
            + _diagnostics_suffix(report),
        )

    now = time.time() if now is None else now
    age = now - report.validated_at
    if age > requirement.max_age_seconds:
        raise CapabilityAdmissionError(
            REASON_STALE,
            f"capability report for {requirement.model_id} {requirement.shard_label} "
            f"is {age / 60:.0f} min old (limit "
            f"{requirement.max_age_seconds / 60:.0f} min); re-run `meshnet-node doctor`",
        )
    if age < -_MAX_CLOCK_SKEW_SECONDS:
        raise CapabilityAdmissionError(
            REASON_STALE,
            f"capability report for {requirement.model_id} {requirement.shard_label} "
            f"is timestamped {-age:.0f}s in the future; check this host's clock",
        )

    return report


def _mismatch(
    reason: str,
    requirement: AdmissionRequirement,
    field_name: str,
    reported: str,
    required: str,
) -> CapabilityAdmissionError:
    return CapabilityAdmissionError(
        reason,
        f"capability report proves a different {field_name}: it validated "
        f"{reported}, but this node would serve {required}. A report is only "
        "proof for the exact combination it ran.",
    )


def _diagnostics_suffix(report: CapabilityReport) -> str:
    if not report.diagnostics:
        return ""
    return " — " + " ".join(report.diagnostics)


def probe_capability(context: CapabilityContext) -> CapabilityReport:
    """Production validator: one bounded real forward through the loaded shard."""
    from .doctor import validate_loaded_backend

    return validate_loaded_backend(
        context.backend,
        context.selection,
        context.recipe,
        context.manifest,
    ).report


def _compatibility_fingerprint_for_context(context: CapabilityContext) -> str:
    backend = context.backend
    selection = context.selection
    recipe = context.recipe
    model_config = getattr(getattr(backend, "model", None), "config", None)
    model_config_payload = (
        model_config.to_dict() if hasattr(model_config, "to_dict") else model_config
    )
    runtime_versions = _runtime_versions()
    runtime_version = _PACKAGE_VERSION
    ownership = authoritative_dense_llama_ownership(backend, selection)
    artifact = build_artifact_identity(
        model_id=selection.model_id,
        revision=getattr(getattr(backend, "model", None), "revision", None),
        model_config=model_config_payload,
        shard_start=ownership.start_layer,
        shard_end=ownership.end_layer,
    )
    runtime_recipe = build_runtime_recipe_identity(
        model_id=selection.model_id,
        revision=getattr(getattr(backend, "model", None), "revision", None),
        model_config=model_config_payload,
        recipe_params=recipe.params,
        weight_quantization=selection.quantization,
        backend_id=recipe.backend_id,
        runtime_version=runtime_version,
        activation_dtype="bfloat16",
        compute_dtype=_backend_compute_dtype(backend),
        kv_dtype=_backend_kv_dtype(backend),
        kv_layout=_backend_kv_layout(backend),
        tokenizer_revision=_backend_tokenizer_revision(backend, selection),
        architecture_adapter=_backend_architecture_adapter(backend, recipe.backend_id),
        boundary_schema_version=1,
        cache_layout=_backend_cache_layout(backend, recipe.params),
    )
    return compatibility_fingerprint(
        fingerprint_payload(
            model={
                "model_id": selection.model_id,
                "revision": getattr(getattr(backend, "model", None), "revision", None),
                "config_fingerprint": config_fingerprint(model_config_payload),
            },
            shard={
                "start": ownership.start_layer,
                "end": ownership.end_layer,
                "owns_embedding": ownership.owns_embedding,
                "owns_final_head": ownership.owns_final_head,
            },
            recipe={
                "recipe_id": recipe.id,
                "recipe_version": recipe.version,
                "catalogue_version": context.manifest.catalogue_version,
            },
            backend={
                "backend_id": recipe.backend_id,
                "device": context.device,
                "device_name": _backend_device_name(context.device),
                "quantization": selection.quantization,
                "runtime": runtime_versions,
            },
            artifact=artifact.to_dict(),
            runtime_recipe=runtime_recipe.to_dict(),
        )
    )


def _runtime_versions() -> dict[str, str]:
    versions: dict[str, str] = {}
    for name in ("torch", "transformers"):
        try:
            module = __import__(name)
        except Exception:
            continue
        version = getattr(module, "__version__", None)
        if version:
            versions[name] = str(version)
    return versions


def _backend_compute_dtype(backend: Any) -> str:
    config = getattr(getattr(backend, "model", None), "config", None)
    for candidate in (config, getattr(config, "text_config", None)):
        if candidate is None:
            continue
        for attr in ("dtype", "torch_dtype"):
            value = getattr(candidate, attr, None)
            if value is None:
                continue
            return str(value).removeprefix("torch.")
    return "bfloat16"


def _backend_kv_dtype(backend: Any) -> str:
    return _backend_compute_dtype(backend)


def _backend_kv_layout(backend: Any) -> str:
    return "session-cache" if getattr(backend, "supports_kv_cache", False) else "stateless"


def _backend_tokenizer_revision(backend: Any, selection: DoctorSelection) -> str:
    model = getattr(backend, "model", None)
    revision = getattr(model, "revision", None)
    if isinstance(revision, str) and revision.strip():
        return revision
    tokenizer = getattr(backend, "tokenizer", None)
    for attr in ("revision", "model_id"):
        value = getattr(tokenizer, attr, None)
        if isinstance(value, str) and value.strip():
            return value
    return selection.model_id


def _backend_architecture_adapter(backend: Any, default: str) -> str:
    config = getattr(getattr(backend, "model", None), "config", None)
    for candidate in (config, getattr(config, "text_config", None)):
        if candidate is None:
            continue
        for attr in ("architecture_adapter", "model_type"):
            value = getattr(candidate, attr, None)
            if isinstance(value, str) and value.strip():
                return value
        architectures = getattr(candidate, "architectures", None)
        if isinstance(architectures, (list, tuple)) and architectures:
            first = architectures[0]
            if isinstance(first, str) and first.strip():
                return first
    return default


def _backend_device_name(device: str) -> str | None:
    if device != "cuda":
        return None
    from .hardware import detect_hardware

    try:
        return detect_hardware().get("gpu_name") or None
    except Exception:
        return None


def _backend_cache_layout(backend: Any, recipe_params: dict[str, Any] | None) -> str:
    if getattr(backend, "supports_kv_cache", False) is False:
        return "stateless"
    if recipe_params is None:
        return "local-hot-kv"
    if recipe_params.get("use_cache") is False:
        return "stateless"
    value = recipe_params.get("cache_layout")
    if isinstance(value, str) and value.strip():
        return value
    return "local-hot-kv"