"""Register a verified native Shard through the ordinary capability contract.

This is intentionally an adapter, not a second tracker protocol.  It converts
the native worker's immutable identity and enforced resource limits into the
same capability report every backend may submit.  The tracker remains the sole
owner of certification and decides whether the visible registration is dark.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

from .capability import ExecutionCapacity, RoutingMeasurements, build_capability_report
from .native_worker_supervisor import NativeWorkerProbe, NativeWorkerSpec, NativeWorkerSupervisor
from .runtime_recipe import ShardIdentity


class NativeRegistrationError(ValueError):
    """Native facts do not describe one coherent, registerable Shard."""


@dataclass(frozen=True)
class NativeShardRegistration:
    """One backend-neutral registration payload for a verified native Shard."""

    endpoint: str
    model_id: str
    identity: ShardIdentity
    worker: NativeWorkerSpec
    probe: NativeWorkerProbe
    device: str
    capacity: ExecutionCapacity
    duration_ms: int = 0
    routing: RoutingMeasurements | None = None

    def __post_init__(self) -> None:
        if not self.endpoint:
            raise NativeRegistrationError("native registration requires an endpoint")
        if not self.model_id:
            raise NativeRegistrationError("native registration requires a model id")
        if not self.device:
            raise NativeRegistrationError("native registration requires a device label")
        if self.identity.artifact.artifact_id != self.model_id:
            raise NativeRegistrationError("native registration model does not match its identity")
        if self.identity.fingerprint.model_artifact_digest != self.worker.artifact_digest:
            raise NativeRegistrationError("native worker artifact digest does not match its identity")
        if self.identity.fingerprint.runtime_recipe_digest != self.worker.recipe_digest:
            raise NativeRegistrationError("native worker recipe digest does not match its identity")
        expected = (
            self.worker.artifact_digest,
            self.worker.recipe_digest,
            self.worker.recipe_id,
            self.worker.recipe_version,
            self.worker.catalogue_version,
            self.worker.shard_start,
            self.worker.shard_end,
        )
        actual = (
            self.probe.artifact_digest,
            self.probe.recipe_digest,
            self.probe.recipe_id,
            self.probe.recipe_version,
            self.probe.catalogue_version,
            self.probe.shard_start,
            self.probe.shard_end,
        )
        if actual != expected:
            raise NativeRegistrationError("native worker probe differs from its startup identity/range")
        if not self.probe.serving:
            raise NativeRegistrationError("native worker is not serving; it cannot register a capability")
        if (
            self.identity.shard_start,
            self.identity.shard_end,
            self.identity.recipe.recipe_id,
            self.identity.recipe.recipe_version,
            self.identity.recipe.catalogue_version,
        ) != (
            self.worker.shard_start,
            self.worker.shard_end,
            self.worker.recipe_id,
            self.worker.recipe_version,
            self.worker.catalogue_version,
        ):
            raise NativeRegistrationError("native identity differs from worker range or recipe labels")
        if self.identity.recipe.axes["backend_id"] == "":
            raise NativeRegistrationError("native identity must name its backend")

    def payload(self) -> dict[str, Any]:
        """Return the existing tracker registration shape with no native branch."""
        report = build_capability_report(
            model_id=self.model_id,
            shard_start=self.identity.shard_start,
            shard_end=self.identity.shard_end - 1,
            recipe_id=self.identity.recipe.recipe_id,
            recipe_version=self.identity.recipe.recipe_version,
            catalogue_version=self.identity.recipe.catalogue_version,
            backend_id=self.identity.recipe.axes["backend_id"],
            device=self.device,
            quantization=self.identity.recipe.axes["weight_quantization"],
            model_config="sha256:" + self.identity.artifact.architecture_digest,
            revision=self.identity.artifact.revision,
            status="passed",
            duration_ms=self.duration_ms,
            identity=self.identity,
            capacity=self.capacity,
            routing=self.routing,
        )
        payload = {
            "endpoint": self.endpoint,
            "model": self.model_id.rsplit("/", 1)[-1],
            "hf_repo": self.model_id,
            "shard_start": self.identity.shard_start,
            "shard_end": self.identity.shard_end - 1,
            "recipe_id": self.identity.recipe.recipe_id,
            "recipe_version": self.identity.recipe.recipe_version,
            "capability_report": report.to_dict(),
            # Existing tracker capacity fields are retained for placement views.
            "ram_bytes": self.capacity.memory_capacity_bytes or 0,
            "max_loaded_shards": 1,
        }
        # These are the tracker’s established dynamic scoring inputs.  The
        # exact same optional report can be sent by any backend; no native
        # route or balancing branch is introduced here.
        if self.routing is not None:
            if self.routing.tokens_per_second is not None:
                payload["benchmark_tokens_per_sec"] = self.routing.tokens_per_second
            if self.routing.queue_depth is not None:
                payload["queue_depth"] = self.routing.queue_depth
        return payload


RegistrationSender = Callable[[dict[str, Any]], None]
WithdrawalSender = Callable[[str], None]


class NativeCapabilityRegistrar:
    """Publish/withdraw a native capability through caller-owned transport.

    The callbacks keep tracker HTTP, relay, billing, and provider mechanics out
    of the native worker.  A process supervisor calls ``withdraw`` on health
    loss; the caller supplies the existing tracker registration/withdrawal
    transport appropriate to its deployment.
    """

    def __init__(
        self,
        registration: NativeShardRegistration,
        *,
        register: RegistrationSender,
        withdraw: WithdrawalSender,
    ) -> None:
        self.registration = registration
        self._register = register
        self._withdraw = withdraw

    def publish(self) -> None:
        self._register(self.registration.payload())

    def unavailable(self, reason: str) -> None:
        self._withdraw(reason)

    def bind(self, supervisor: NativeWorkerSupervisor) -> None:
        """Publish only after DGR-040 verification; withdraw on health loss."""
        if supervisor.spec != self.registration.worker:
            raise NativeRegistrationError("registrar and supervisor must own the same native worker")
        supervisor.add_availability_callbacks(
            on_available=lambda _reason: self.publish(),
            on_unavailable=self.unavailable,
        )