"""Immutable model/dataset artifact metadata and promotion gates."""

from __future__ import annotations

import hashlib
import json
from dataclasses import asdict, dataclass
from enum import Enum, StrEnum
from pathlib import Path


class ArtifactType(StrEnum):
    MODEL = "graph_reasoner_model"
    TOKENIZER = "graph_reasoner_tokenizer"
    CONFIG = "graph_reasoner_config"
    CALIBRATION = "graph_reasoner_calibration"
    EVALUATION = "graph_reasoner_evaluation"
    DATASET = "graph_reasoner_dataset"
    GENERATOR = "graph_reasoner_generator"


@dataclass(frozen=True)
class ArtifactMetadata:
    artifact_type: ArtifactType
    version: str
    path: str
    sha256: str
    created_at: str
    model_version: str | None = None
    graph_version: str | None = None
    parameter_count: int | None = None
    evaluation_passed: bool = False
    immutable: bool = True
    architecture_version: str | None = None
    dataset_version: str | None = None
    generator_version: str | None = None
    tokenizer_version: str | None = None
    ontology_version: str | None = None
    calibration_version: str | None = None
    evaluation_version: str | None = None
    config_hash: str | None = None
    git_commit: str | None = None
    pat_config_hash: str | None = None


class ArtifactRegistry:
    def __init__(self, root: str | Path) -> None:
        self.root = Path(root)
        self.root.mkdir(parents=True, exist_ok=True)
        self.index_path = self.root / "registry.json"
        self.active_path = self.root / "active.json"
        self._index: dict[str, ArtifactMetadata] = {}
        if self.index_path.exists():
            raw = json.loads(self.index_path.read_text(encoding="utf-8"))
            self._index = {
                key: ArtifactMetadata(
                    **{**value, "artifact_type": ArtifactType(value["artifact_type"])}
                )
                for key, value in raw.items()
            }

    def register(self, metadata: ArtifactMetadata) -> ArtifactMetadata:
        key = f"{metadata.artifact_type.value}:{metadata.version}"
        path = Path(metadata.path)
        if not path.exists():
            raise FileNotFoundError(path)
        if sha256_file(path) != metadata.sha256:
            raise ValueError(f"artifact digest does not match payload: {path}")
        if key in self._index:
            if self._index[key] != metadata:
                raise ValueError(
                    f"immutable artifact metadata already registered: {key}"
                )
            return self._index[key]
        self._index[key] = metadata
        self._write()
        return metadata

    def activate(self, artifact_type: ArtifactType, version: str) -> ArtifactMetadata:
        key = f"{artifact_type.value}:{version}"
        if key not in self._index:
            raise KeyError(key)
        metadata = self._index[key]
        if artifact_type == ArtifactType.MODEL and not metadata.evaluation_passed:
            raise ValueError("model cannot be activated before evaluation gates pass")
        # Registration hashed the payload; promotion must see the same bytes.
        if sha256_file(metadata.path) != metadata.sha256:
            raise ValueError(f"{key} changed after registration; refusing to activate")
        active = {}
        if self.active_path.exists():
            active = json.loads(self.active_path.read_text(encoding="utf-8"))
        active[artifact_type.value] = version
        self.active_path.write_text(
            json.dumps(active, indent=2, sort_keys=True), encoding="utf-8"
        )
        return metadata

    def get(self, artifact_type: ArtifactType, version: str) -> ArtifactMetadata:
        return self._index[f"{artifact_type.value}:{version}"]

    def _write(self) -> None:
        self.index_path.write_text(
            json.dumps(
                {key: asdict(value) for key, value in self._index.items()},
                indent=2,
                default=lambda value: value.value if isinstance(value, Enum) else value,
                sort_keys=True,
            ),
            encoding="utf-8",
        )


def sha256_file(path: str | Path) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()