"""The pinned GLM-5.2 Max target manifest and architecture snapshot.

This module is the *identity* half of the alpha target contract. It answers one
question: is the artifact on this disk the exact artifact alpha was locked
against?

Identity is pinned by repository revision **and** by every shard's LFS SHA-256.
A revision alone is not enough — a repository can be force-pushed, and a tag can
be moved — and a size alone is not enough, because two different quantizations of
the same model land within a rounding error of each other. The manifest therefore
carries both, plus the aggregate byte total, and cross-checks the aggregate
against the sum of the shards. A manifest whose declared total disagrees with its
own shards is rejected rather than trusted, because that is the exact shape a
hand-edited "it fits now" manifest takes.

Nothing here downloads a weight payload. Sizes and hashes come from the Hugging
Face metadata API (see ``scripts/refresh_glm_target_manifest.py``); verification
against a local file is DGR-018's job, using the digests locked here.
"""

from __future__ import annotations

import hashlib
import json
import re
from dataclasses import dataclass
from importlib.resources import files
from pathlib import Path
from typing import Any, Mapping

TARGET_MANIFEST_SCHEMA_VERSION = 1
ARCHITECTURE_SNAPSHOT_SCHEMA_VERSION = 1

ALPHA_QUANTIZATION = "UD-IQ1_S"
ALPHA_SHARD_COUNT = 6

_SHA256_RE = re.compile(r"\A[0-9a-f]{64}\Z")

_MANIFEST_RESOURCE = "target-manifest.json"
_ARCHITECTURE_RESOURCE = "architecture-snapshot.json"

GIB = 1024**3
GB = 1000**3


class GlmTargetError(ValueError):
    """Raised when the target manifest or architecture snapshot is not the pinned target."""


def canonical_sha256(value: Any) -> str:
    """SHA-256 over canonical JSON — the repository's digest convention."""
    payload = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def _require_mapping(value: Any, what: str) -> Mapping[str, Any]:
    if not isinstance(value, Mapping):
        raise GlmTargetError(f"{what} must be a JSON object, got {type(value).__name__}")
    return value


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


def _require_int(value: Any, what: str) -> int:
    if not isinstance(value, int) or isinstance(value, bool):
        raise GlmTargetError(f"{what} must be an integer, got {type(value).__name__}")
    return value


def _require_sha256(value: Any, what: str) -> str:
    text = _require_text(value, what)
    if not _SHA256_RE.match(text):
        raise GlmTargetError(
            f"{what} must be a lowercase 64-character hex SHA-256, got {text!r}"
        )
    return text


def _require_revision(value: Any, what: str) -> str:
    text = _require_text(value, what)
    if not re.fullmatch(r"[0-9a-f]{40}", text):
        raise GlmTargetError(
            f"{what} must be a full 40-character commit revision, got {text!r}; "
            "a branch name or short SHA is not an immutable pin"
        )
    return text


@dataclass(frozen=True)
class Shard:
    """One GGUF shard of the alpha artifact."""

    index: int
    path: str
    size_bytes: int
    sha256: str
    url: str

    def to_dict(self) -> dict:
        return {
            "index": self.index,
            "path": self.path,
            "size_bytes": self.size_bytes,
            "sha256": self.sha256,
            "url": self.url,
        }


@dataclass(frozen=True)
class TargetManifest:
    """The pinned, self-consistent GLM-5.2 ``UD-IQ1_S`` target."""

    schema_version: int
    manifest_version: int
    observed_at: str
    quantization: str
    source_repo_id: str
    source_revision: str
    source_license: str
    gguf_repo_id: str
    gguf_revision: str
    gguf_license: str
    total_bytes: int
    shards: tuple[Shard, ...]
    raw: Mapping[str, Any]
    source: str = "<memory>"

    @property
    def total_gib(self) -> float:
        return self.total_bytes / GIB

    @property
    def total_gb(self) -> float:
        return self.total_bytes / GB

    def shard(self, index: int) -> Shard:
        for shard in self.shards:
            if shard.index == index:
                return shard
        raise GlmTargetError(f"shard {index} is not in {self.source}")

    @property
    def digest(self) -> str:
        """Stable identity of this manifest, for the DGR-003 runtime recipe."""
        return canonical_sha256(self.raw)

    def to_dict(self) -> dict:
        return dict(self.raw)


def _parse_shards(raw: Any, expected_count: int, expected_total: int) -> tuple[Shard, ...]:
    if not isinstance(raw, list):
        raise GlmTargetError("gguf_artifact.shards must be a JSON array")
    if len(raw) != expected_count:
        raise GlmTargetError(
            f"the alpha artifact has exactly {expected_count} shards, "
            f"but the manifest lists {len(raw)}"
        )

    shards: list[Shard] = []
    seen_index: set[int] = set()
    seen_sha: set[str] = set()
    for position, entry in enumerate(raw):
        item = _require_mapping(entry, f"shards[{position}]")
        index = _require_int(item.get("index"), f"shards[{position}].index")
        if index in seen_index:
            raise GlmTargetError(f"duplicate shard index {index} in the manifest")
        seen_index.add(index)
        size_bytes = _require_int(item.get("size_bytes"), f"shards[{index}].size_bytes")
        if size_bytes <= 0:
            raise GlmTargetError(f"shards[{index}].size_bytes must be positive")
        sha256 = _require_sha256(item.get("sha256"), f"shards[{index}].sha256")
        if sha256 in seen_sha:
            raise GlmTargetError(
                f"shard {index} repeats SHA-256 {sha256}; two distinct shards cannot "
                "have the same content digest"
            )
        seen_sha.add(sha256)
        shards.append(
            Shard(
                index=index,
                path=_require_text(item.get("path"), f"shards[{index}].path"),
                size_bytes=size_bytes,
                sha256=sha256,
                url=_require_text(item.get("url"), f"shards[{index}].url"),
            )
        )

    expected_indices = set(range(1, expected_count + 1))
    if seen_index != expected_indices:
        missing = sorted(expected_indices - seen_index)
        raise GlmTargetError(
            f"the manifest is missing shard(s) {missing}; all {expected_count} "
            "shards of the alpha artifact must be pinned"
        )

    summed = sum(shard.size_bytes for shard in shards)
    if summed != expected_total:
        raise GlmTargetError(
            f"declared total_bytes {expected_total} does not equal the sum of the "
            f"shard sizes {summed}; the manifest is not self-consistent"
        )

    return tuple(sorted(shards, key=lambda shard: shard.index))


def parse_target_manifest(data: Any, source: str = "<memory>") -> TargetManifest:
    """Validate an already-decoded target manifest, failing closed."""
    doc = _require_mapping(data, f"manifest root in {source}")

    schema_version = _require_int(doc.get("schema_version"), f"'schema_version' in {source}")
    if schema_version != TARGET_MANIFEST_SCHEMA_VERSION:
        raise GlmTargetError(
            f"{source} declares target-manifest schema version {schema_version}, "
            f"but this node reads version {TARGET_MANIFEST_SCHEMA_VERSION}"
        )

    quantization = _require_text(doc.get("alpha_quantization"), f"'alpha_quantization' in {source}")
    if quantization != ALPHA_QUANTIZATION:
        raise GlmTargetError(
            f"{source} pins quantization {quantization!r}, but the locked alpha "
            f"quantization is {ALPHA_QUANTIZATION!r}; a different quantization is a "
            "different target and requires a human contract change"
        )

    source_model = _require_mapping(doc.get("source_model"), f"'source_model' in {source}")
    gguf = _require_mapping(doc.get("gguf_artifact"), f"'gguf_artifact' in {source}")

    gguf_quant = _require_text(gguf.get("quantization"), f"gguf_artifact.quantization in {source}")
    if gguf_quant != quantization:
        raise GlmTargetError(
            f"{source} declares alpha_quantization {quantization!r} but the GGUF "
            f"artifact block says {gguf_quant!r}"
        )

    shard_count = _require_int(gguf.get("shard_count"), f"gguf_artifact.shard_count in {source}")
    if shard_count != ALPHA_SHARD_COUNT:
        raise GlmTargetError(
            f"{source} declares {shard_count} shards; the pinned alpha artifact has "
            f"exactly {ALPHA_SHARD_COUNT}"
        )

    total_bytes = _require_int(gguf.get("total_bytes"), f"gguf_artifact.total_bytes in {source}")
    shards = _parse_shards(gguf.get("shards"), shard_count, total_bytes)

    return TargetManifest(
        schema_version=schema_version,
        manifest_version=_require_int(doc.get("manifest_version"), f"'manifest_version' in {source}"),
        observed_at=_require_text(doc.get("observed_at"), f"'observed_at' in {source}"),
        quantization=quantization,
        source_repo_id=_require_text(source_model.get("repo_id"), "source_model.repo_id"),
        source_revision=_require_revision(source_model.get("revision"), "source_model.revision"),
        source_license=_require_text(source_model.get("weight_license"), "source_model.weight_license"),
        gguf_repo_id=_require_text(gguf.get("repo_id"), "gguf_artifact.repo_id"),
        gguf_revision=_require_revision(gguf.get("revision"), "gguf_artifact.revision"),
        gguf_license=_require_text(gguf.get("license"), "gguf_artifact.license"),
        total_bytes=total_bytes,
        shards=shards,
        raw=doc,
        source=source,
    )


@dataclass(frozen=True)
class ArchitectureSnapshot:
    """Architecture-critical metadata derived from the pinned ``config.json``."""

    schema_version: int
    source_repo_id: str
    source_revision: str
    architecture: Mapping[str, Any]
    reasoning_effort: Mapping[str, Any]
    source_files: Mapping[str, Mapping[str, Any]]
    raw: Mapping[str, Any]
    source: str = "<memory>"

    def __getitem__(self, key: str) -> Any:
        if key not in self.architecture:
            raise GlmTargetError(f"architecture field {key!r} is missing from {self.source}")
        return self.architecture[key]

    @property
    def digest(self) -> str:
        return canonical_sha256(self.raw)

    def file_sha256(self, path: str) -> str:
        entry = self.source_files.get(path)
        if entry is None:
            raise GlmTargetError(f"{path!r} is not pinned in {self.source}")
        return str(entry["sha256"])

    def to_dict(self) -> dict:
        return dict(self.raw)


# Fields the distributed runtime cannot plan or shard without. Absent or
# contradictory values fail closed rather than defaulting.
REQUIRED_ARCHITECTURE_FIELDS: tuple[str, ...] = (
    "model_type",
    "num_hidden_layers",
    "num_nextn_predict_layers",
    "total_artifact_layers",
    "first_k_dense_replace",
    "dense_layers",
    "sparse_moe_layers",
    "hidden_size",
    "n_routed_experts",
    "num_experts_per_tok",
    "n_shared_experts",
    "kv_lora_rank",
    "qk_rope_head_dim",
    "mla_cached_values_per_token_per_layer",
    "index_topk",
    "index_head_dim",
    "indexer_full_layers",
    "indexer_shared_layers",
    "max_position_embeddings",
    "vocab_size",
)


def parse_architecture_snapshot(data: Any, source: str = "<memory>") -> ArchitectureSnapshot:
    """Validate an architecture snapshot and its internal arithmetic."""
    doc = _require_mapping(data, f"snapshot root in {source}")

    schema_version = _require_int(doc.get("schema_version"), f"'schema_version' in {source}")
    if schema_version != ARCHITECTURE_SNAPSHOT_SCHEMA_VERSION:
        raise GlmTargetError(
            f"{source} declares architecture-snapshot schema version {schema_version}, "
            f"but this node reads version {ARCHITECTURE_SNAPSHOT_SCHEMA_VERSION}"
        )

    arch = _require_mapping(doc.get("architecture"), f"'architecture' in {source}")
    missing = [field for field in REQUIRED_ARCHITECTURE_FIELDS if field not in arch]
    if missing:
        raise GlmTargetError(
            f"{source} is missing architecture-critical field(s) {missing}; the "
            "runtime cannot shard or plan an architecture it cannot fully describe"
        )

    layers = _require_int(arch["num_hidden_layers"], "num_hidden_layers")
    nextn = _require_int(arch["num_nextn_predict_layers"], "num_nextn_predict_layers")
    total_layers = _require_int(arch["total_artifact_layers"], "total_artifact_layers")
    if total_layers != layers + nextn:
        raise GlmTargetError(
            f"total_artifact_layers {total_layers} != num_hidden_layers {layers} + "
            f"num_nextn_predict_layers {nextn}; the NextN layer must be counted "
            "explicitly, never folded into the backbone"
        )

    dense = _require_int(arch["dense_layers"], "dense_layers")
    sparse = _require_int(arch["sparse_moe_layers"], "sparse_moe_layers")
    if dense != _require_int(arch["first_k_dense_replace"], "first_k_dense_replace"):
        raise GlmTargetError("dense_layers must equal first_k_dense_replace")
    if dense + sparse != layers:
        raise GlmTargetError(
            f"dense_layers {dense} + sparse_moe_layers {sparse} != num_hidden_layers {layers}"
        )

    full = _require_int(arch["indexer_full_layers"], "indexer_full_layers")
    shared = _require_int(arch["indexer_shared_layers"], "indexer_shared_layers")
    if full + shared != layers:
        raise GlmTargetError(
            f"indexer_full_layers {full} + indexer_shared_layers {shared} != "
            f"num_hidden_layers {layers}; every layer holds exactly one IndexShare role"
        )
    if full <= 0:
        raise GlmTargetError(
            "indexer_full_layers must be positive; a route with no Full producer layer "
            "has no index for its Shared consumers to reuse"
        )

    mla = _require_int(
        arch["mla_cached_values_per_token_per_layer"], "mla_cached_values_per_token_per_layer"
    )
    expected_mla = _require_int(arch["kv_lora_rank"], "kv_lora_rank") + _require_int(
        arch["qk_rope_head_dim"], "qk_rope_head_dim"
    )
    if mla != expected_mla:
        raise GlmTargetError(
            f"mla_cached_values_per_token_per_layer {mla} != kv_lora_rank + "
            f"qk_rope_head_dim ({expected_mla})"
        )

    reasoning = _require_mapping(doc.get("reasoning_effort"), f"'reasoning_effort' in {source}")
    if reasoning.get("alpha_mode") != "max":
        raise GlmTargetError(
            f"{source} does not lock reasoning_effort=max; alpha is defined as the "
            "Max reasoning mode of this exact checkpoint"
        )
    _require_text(reasoning.get("rendered_marker"), "reasoning_effort.rendered_marker")

    files_raw = doc.get("source_files")
    if not isinstance(files_raw, list) or not files_raw:
        raise GlmTargetError(f"'source_files' in {source} must be a non-empty JSON array")
    source_files: dict[str, Mapping[str, Any]] = {}
    for position, entry in enumerate(files_raw):
        item = _require_mapping(entry, f"source_files[{position}]")
        path = _require_text(item.get("path"), f"source_files[{position}].path")
        _require_sha256(item.get("sha256"), f"source_files[{path}].sha256")
        source_files[path] = item
    for required in ("config.json", "chat_template.jinja"):
        if required not in source_files:
            raise GlmTargetError(
                f"{source} does not pin {required!r}; config and chat-template drift "
                "silently changes runtime semantics"
            )

    return ArchitectureSnapshot(
        schema_version=schema_version,
        source_repo_id=_require_text(doc.get("source_repo_id"), "source_repo_id"),
        source_revision=_require_revision(doc.get("source_revision"), "source_revision"),
        architecture=arch,
        reasoning_effort=reasoning,
        source_files=source_files,
        raw=doc,
        source=source,
    )


def _read_resource(resource: str, path: Path | None) -> tuple[str, str]:
    if path is not None:
        try:
            return str(path), path.read_text(encoding="utf-8")
        except OSError as exc:
            raise GlmTargetError(f"cannot read {path}: {exc.strerror or exc}") from exc
    source = f"packaged {resource}"
    try:
        raw = files("meshnet_node.glm_alpha").joinpath("data", resource).read_text(encoding="utf-8")
    except (OSError, FileNotFoundError, ModuleNotFoundError) as exc:
        raise GlmTargetError(
            f"{source} is missing from this node installation ({type(exc).__name__})"
        ) from exc
    return source, raw


def _load_json(resource: str, path: Path | None) -> tuple[str, Any]:
    source, raw = _read_resource(resource, path)
    try:
        return source, json.loads(raw)
    except json.JSONDecodeError as exc:
        raise GlmTargetError(
            f"{source} is not valid JSON: {exc.msg} at line {exc.lineno} column {exc.colno}"
        ) from exc


def load_target_manifest(path: Path | None = None) -> TargetManifest:
    """Load the packaged target manifest, or one at ``path``."""
    source, data = _load_json(_MANIFEST_RESOURCE, path)
    return parse_target_manifest(data, source=source)


def load_architecture_snapshot(path: Path | None = None) -> ArchitectureSnapshot:
    """Load the packaged architecture snapshot, or one at ``path``."""
    source, data = _load_json(_ARCHITECTURE_RESOURCE, path)
    return parse_architecture_snapshot(data, source=source)


def require_pinned_target(
    manifest: TargetManifest,
    snapshot: ArchitectureSnapshot,
    *,
    expected_source_revision: str,
    expected_gguf_revision: str,
) -> None:
    """Reject any target whose revisions are not the ones alpha was locked against.

    Callers pass the revisions from the locked alpha contract, so a swapped
    manifest cannot quietly re-point the target at a different upstream commit.
    """
    if manifest.source_revision != expected_source_revision:
        raise GlmTargetError(
            f"source revision {manifest.source_revision} does not match the locked "
            f"alpha revision {expected_source_revision}"
        )
    if manifest.gguf_revision != expected_gguf_revision:
        raise GlmTargetError(
            f"GGUF revision {manifest.gguf_revision} does not match the locked alpha "
            f"revision {expected_gguf_revision}"
        )
    if snapshot.source_revision != manifest.source_revision:
        raise GlmTargetError(
            f"the architecture snapshot was taken at {snapshot.source_revision} but the "
            f"manifest pins {manifest.source_revision}; config metadata and weights must "
            "come from one revision"
        )