"""Exact artifact and runtime-recipe identity helpers.

The runtime recipe is the compatibility contract for one routable shard. It is
kept separate from the user-facing recipe catalogue so the tracker can compare
the exact execution footprint that was validated, not just a named recipe.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from typing import Any, Mapping


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


def _optional_text(value: Any, field_name: str) -> str | None:
    if value is None:
        return None
    return _require_text(value, field_name)


def _sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _stable_json(data: Any) -> str:
    return json.dumps(
        data,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
        default=str,
    )


def _normalise_dtype(value: Any, default: str) -> str:
    if value is None:
        return default
    if isinstance(value, str):
        text = value.strip()
        if not text:
            return default
        return text.removeprefix("torch.")
    return str(value).removeprefix("torch.")


def _architecture_adapter_from_config(model_config: Any, default: str) -> str:
    if not isinstance(model_config, Mapping):
        return default
    for key in ("architecture_adapter", "model_type"):
        value = model_config.get(key)
        if isinstance(value, str) and value.strip():
            return value
    architectures = model_config.get("architectures")
    if isinstance(architectures, list) and architectures:
        first = architectures[0]
        if isinstance(first, str) and first.strip():
            return first
    text_config = model_config.get("text_config")
    if isinstance(text_config, Mapping):
        return _architecture_adapter_from_config(text_config, default)
    return default


def _tokenizer_revision_from_config(
    model_id: str,
    revision: str | None,
    model_config: Any,
) -> str:
    if isinstance(model_config, Mapping):
        for key in ("tokenizer_revision", "tokenizer_version", "_commit_hash"):
            value = model_config.get(key)
            if isinstance(value, str) and value.strip():
                return value
    if revision:
        return revision
    return model_id


def _cache_layout_from_recipe_params(recipe_params: Mapping[str, Any] | None) -> str:
    if not recipe_params:
        return "local-hot-kv"
    use_cache = recipe_params.get("use_cache")
    if use_cache is False:
        return "stateless"
    if "cache_layout" in recipe_params:
        value = recipe_params.get("cache_layout")
        if isinstance(value, str) and value.strip():
            return value
    return "local-hot-kv"


@dataclass(frozen=True)
class ArtifactIdentity:
    """Exact source artifact binding for a routable shard."""

    model_id: str
    revision: str | None = None
    artifact_hash: str | None = None
    shard_start: int | None = None
    shard_end: int | None = None

    def __post_init__(self) -> None:
        _require_text(self.model_id, "artifact.model_id")
        _optional_text(self.revision, "artifact.revision")
        _optional_text(self.artifact_hash, "artifact.artifact_hash")
        if self.shard_start is not None and self.shard_start < 0:
            raise ValueError("'artifact.shard_start' must be >= 0")
        if self.shard_end is not None and self.shard_end < 0:
            raise ValueError("'artifact.shard_end' must be >= 0")
        if (
            self.shard_start is not None
            and self.shard_end is not None
            and self.shard_end < self.shard_start
        ):
            raise ValueError("'artifact.shard_end' must be >= 'artifact.shard_start'")

    def to_dict(self) -> dict[str, Any]:
        return {
            "model_id": self.model_id,
            "revision": self.revision,
            "artifact_hash": self.artifact_hash,
            "shard_start": self.shard_start,
            "shard_end": self.shard_end,
        }

    @classmethod
    def from_dict(cls, data: Any) -> "ArtifactIdentity":
        if not isinstance(data, Mapping):
            raise ValueError(f"'artifact' must be a JSON object, got {type(data).__name__}")
        return cls(
            model_id=_require_text(data.get("model_id"), "artifact.model_id"),
            revision=_optional_text(data.get("revision"), "artifact.revision"),
            artifact_hash=_optional_text(
                data.get("artifact_hash"), "artifact.artifact_hash"
            ),
            shard_start=_optional_int(data.get("shard_start"), "artifact.shard_start"),
            shard_end=_optional_int(data.get("shard_end"), "artifact.shard_end"),
        )


@dataclass(frozen=True)
class RuntimeRecipeIdentity:
    """Exact runtime recipe used for admission and handshake compatibility."""

    weight_quantization: str
    activation_dtype: str
    compute_dtype: str
    kv_dtype: str
    kv_layout: str
    tokenizer_revision: str
    architecture_adapter: str
    backend_id: str
    runtime_version: str
    boundary_schema_version: int = 1
    cache_layout: str = "local-hot-kv"
    fingerprint: str | None = None

    def __post_init__(self) -> None:
        _require_text(self.weight_quantization, "runtime_recipe.weight_quantization")
        _require_text(self.activation_dtype, "runtime_recipe.activation_dtype")
        _require_text(self.compute_dtype, "runtime_recipe.compute_dtype")
        _require_text(self.kv_dtype, "runtime_recipe.kv_dtype")
        _require_text(self.kv_layout, "runtime_recipe.kv_layout")
        _require_text(self.tokenizer_revision, "runtime_recipe.tokenizer_revision")
        _require_text(self.architecture_adapter, "runtime_recipe.architecture_adapter")
        _require_text(self.backend_id, "runtime_recipe.backend_id")
        _require_text(self.runtime_version, "runtime_recipe.runtime_version")
        _require_text(self.cache_layout, "runtime_recipe.cache_layout")
        if self.boundary_schema_version < 1:
            raise ValueError("'runtime_recipe.boundary_schema_version' must be >= 1")
        expected = compatibility_fingerprint(self._fingerprint_payload())
        if not self.fingerprint:
            object.__setattr__(self, "fingerprint", expected)
        elif self.fingerprint != expected:
            raise ValueError(
                "'runtime_recipe.fingerprint' does not match the encoded fields"
            )

    def to_dict(self) -> dict[str, Any]:
        return {
            "weight_quantization": self.weight_quantization,
            "activation_dtype": self.activation_dtype,
            "compute_dtype": self.compute_dtype,
            "kv_dtype": self.kv_dtype,
            "kv_layout": self.kv_layout,
            "tokenizer_revision": self.tokenizer_revision,
            "architecture_adapter": self.architecture_adapter,
            "backend_id": self.backend_id,
            "runtime_version": self.runtime_version,
            "boundary_schema_version": self.boundary_schema_version,
            "cache_layout": self.cache_layout,
            "fingerprint": self.fingerprint,
        }

    @classmethod
    def from_dict(cls, data: Any) -> "RuntimeRecipeIdentity":
        if not isinstance(data, Mapping):
            raise ValueError(
                f"'runtime_recipe' must be a JSON object, got {type(data).__name__}"
            )
        boundary_schema_version = data.get("boundary_schema_version", 1)
        if isinstance(boundary_schema_version, bool) or not isinstance(
            boundary_schema_version, int
        ):
            raise ValueError(
                "'runtime_recipe.boundary_schema_version' must be an integer"
            )
        return cls(
            weight_quantization=_require_text(
                data.get("weight_quantization"), "runtime_recipe.weight_quantization"
            ),
            activation_dtype=_require_text(
                data.get("activation_dtype"), "runtime_recipe.activation_dtype"
            ),
            compute_dtype=_require_text(
                data.get("compute_dtype"), "runtime_recipe.compute_dtype"
            ),
            kv_dtype=_require_text(data.get("kv_dtype"), "runtime_recipe.kv_dtype"),
            kv_layout=_require_text(data.get("kv_layout"), "runtime_recipe.kv_layout"),
            tokenizer_revision=_require_text(
                data.get("tokenizer_revision"), "runtime_recipe.tokenizer_revision"
            ),
            architecture_adapter=_require_text(
                data.get("architecture_adapter"),
                "runtime_recipe.architecture_adapter",
            ),
            backend_id=_require_text(data.get("backend_id"), "runtime_recipe.backend_id"),
            runtime_version=_require_text(
                data.get("runtime_version"), "runtime_recipe.runtime_version"
            ),
            boundary_schema_version=boundary_schema_version,
            cache_layout=_require_text(data.get("cache_layout"), "runtime_recipe.cache_layout"),
            fingerprint=_optional_text(data.get("fingerprint"), "runtime_recipe.fingerprint"),
        )

    def _fingerprint_payload(self) -> dict[str, Any]:
        return {
            "weight_quantization": self.weight_quantization,
            "activation_dtype": self.activation_dtype,
            "compute_dtype": self.compute_dtype,
            "kv_dtype": self.kv_dtype,
            "kv_layout": self.kv_layout,
            "tokenizer_revision": self.tokenizer_revision,
            "architecture_adapter": self.architecture_adapter,
            "backend_id": self.backend_id,
            "runtime_version": self.runtime_version,
            "boundary_schema_version": self.boundary_schema_version,
            "cache_layout": self.cache_layout,
        }


def _optional_int(value: Any, field_name: str) -> int | None:
    if value is None:
        return None
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError(f"{field_name!r} must be an integer")
    if value < 0:
        raise ValueError(f"{field_name!r} must be >= 0")
    return value


def build_artifact_identity(
    *,
    model_id: str,
    revision: str | None = None,
    model_config: Any = None,
    artifact_hash: str | None = None,
    shard_start: int | None = None,
    shard_end: int | None = None,
) -> ArtifactIdentity:
    """Build a stable artifact binding from the locally loaded artifact."""
    resolved_hash = artifact_hash
    if resolved_hash is None:
        if isinstance(model_config, Mapping):
            resolved_hash = _hash_mapping(model_config)
        elif model_config is not None:
            resolved_hash = _sha256_text(_stable_json(model_config))
    if resolved_hash is None:
        resolved_hash = _sha256_text(
            _stable_json(
                {
                    "model_id": model_id,
                    "revision": revision,
                    "shard_start": shard_start,
                    "shard_end": shard_end,
                }
            )
        )
    return ArtifactIdentity(
        model_id=model_id,
        revision=revision,
        artifact_hash=resolved_hash,
        shard_start=shard_start,
        shard_end=shard_end,
    )


def build_runtime_recipe_identity(
    *,
    model_id: str,
    weight_quantization: str,
    backend_id: str,
    runtime_version: str,
    revision: str | None = None,
    model_config: Any = None,
    recipe_params: Mapping[str, Any] | None = None,
    activation_dtype: Any = None,
    compute_dtype: Any = None,
    kv_dtype: Any = None,
    kv_layout: str | None = None,
    tokenizer_revision: str | None = None,
    architecture_adapter: str | None = None,
    boundary_schema_version: int = 1,
    cache_layout: str | None = None,
) -> RuntimeRecipeIdentity:
    """Build the exact runtime recipe used for compatibility admission."""
    activation = _normalise_dtype(activation_dtype, "bfloat16")
    compute = _normalise_dtype(compute_dtype, activation)
    kv_dtype_text = _normalise_dtype(kv_dtype, compute)
    kv_layout_text = kv_layout or "session-cache"
    tokenizer = tokenizer_revision or _tokenizer_revision_from_config(
        model_id, revision, model_config
    )
    architecture = architecture_adapter or _architecture_adapter_from_config(
        model_config, backend_id
    )
    cache_layout_text = cache_layout or _cache_layout_from_recipe_params(recipe_params)
    return RuntimeRecipeIdentity(
        weight_quantization=weight_quantization,
        activation_dtype=activation,
        compute_dtype=compute,
        kv_dtype=kv_dtype_text,
        kv_layout=kv_layout_text,
        tokenizer_revision=tokenizer,
        architecture_adapter=architecture,
        backend_id=backend_id,
        runtime_version=runtime_version,
        boundary_schema_version=boundary_schema_version,
        cache_layout=cache_layout_text,
    )


def compatibility_fingerprint(data: Mapping[str, Any]) -> str:
    """Return a stable SHA256 compatibility fingerprint for an exact route."""
    return "sha256:" + _sha256_text(_stable_json(data))


def fingerprint_payload(
    *,
    model: Mapping[str, Any],
    shard: Mapping[str, Any],
    recipe: Mapping[str, Any],
    backend: Mapping[str, Any],
    artifact: Mapping[str, Any],
    runtime_recipe: Mapping[str, Any],
) -> dict[str, Any]:
    return {
        "model": dict(model),
        "shard": dict(shard),
        "recipe": dict(recipe),
        "backend": dict(backend),
        "artifact": dict(artifact),
        "runtime_recipe": dict(runtime_recipe),
    }


def _hash_mapping(data: Mapping[str, Any]) -> str:
    return "sha256:" + _sha256_text(_stable_json(data))