"""Model checkpoint — serialize and restore WorldDAG state.

Provides lossless round-trip serialization of WorldDAG instances
(entities + edges) to JSON-compatible dicts. Includes metadata
(version, timestamp, content hash) for audit trail compliance.

All implementations are stdlib-only (ADR-007 Tier-0 compliance).
"""

from __future__ import annotations

import hashlib
import json
import time
from dataclasses import dataclass, field
from pathlib import Path

from orxaq.kernel.types import (
    CausalEdge,
    CausalRole,
    Entity,
    SensitivityLevel,
    VariableType,
)
from orxaq.kernel.world_dag import WorldDAG


@dataclass(frozen=True)
class Checkpoint:
    """Immutable snapshot of a WorldDAG with audit metadata."""

    name: str
    version: str
    timestamp: float
    content_hash: str
    entities: tuple[dict, ...]
    edges: tuple[dict, ...]
    metadata: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        """Serialize checkpoint to a JSON-compatible dict."""
        return {
            "name": self.name,
            "version": self.version,
            "timestamp": self.timestamp,
            "content_hash": self.content_hash,
            "entities": list(self.entities),
            "edges": list(self.edges),
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict) -> Checkpoint:
        """Reconstruct a Checkpoint from a dict."""
        return cls(
            name=data["name"],
            version=data["version"],
            timestamp=data["timestamp"],
            content_hash=data["content_hash"],
            entities=tuple(data["entities"]),
            edges=tuple(data["edges"]),
            metadata=data.get("metadata", {}),
        )


def _entity_to_dict(entity: Entity) -> dict:
    """Serialize an Entity to a JSON-compatible dict."""
    return {
        "entity_id": entity.entity_id,
        "name": entity.name,
        "description": entity.description,
        "causal_role": entity.causal_role.value,
        "variable_type": entity.variable_type.value,
        "sensitivity": entity.sensitivity.value,
        "domain": entity.domain,
        "units": entity.units,
        "bounds": list(entity.bounds),
        "regulatory_tags": list(entity.regulatory_tags),
        "metadata": list(entity.metadata),
    }


def _entity_from_dict(data: dict) -> Entity:
    """Reconstruct an Entity from a dict."""
    return Entity(
        entity_id=data["entity_id"],
        name=data["name"],
        description=data.get("description", ""),
        causal_role=CausalRole(data.get("causal_role", "exogenous")),
        variable_type=VariableType(data.get("variable_type", "static")),
        sensitivity=SensitivityLevel(data.get("sensitivity", "internal")),
        domain=data.get("domain", ""),
        units=data.get("units", ""),
        bounds=tuple(data.get("bounds", [None, None])),
        regulatory_tags=tuple(data.get("regulatory_tags", ())),
        metadata=tuple(tuple(m) for m in data.get("metadata", ())),
    )


def _edge_to_dict(edge: CausalEdge) -> dict:
    """Serialize a CausalEdge to a JSON-compatible dict."""
    return {
        "source": edge.source,
        "target": edge.target,
        "mechanism": edge.mechanism,
        "confidence": edge.confidence,
        "soliton_score": edge.soliton_score,
        "provenance": edge.provenance,
        "is_immutable": edge.is_immutable,
        "edge_type": edge.edge_type,
        "temporal_lag": edge.temporal_lag,
    }


def _edge_from_dict(data: dict) -> CausalEdge:
    """Reconstruct a CausalEdge from a dict."""
    return CausalEdge(
        source=data["source"],
        target=data["target"],
        mechanism=data.get("mechanism", ""),
        confidence=data.get("confidence", 0.5),
        soliton_score=data.get("soliton_score", 0.5),
        provenance=data.get("provenance", ""),
        is_immutable=data.get("is_immutable", False),
        edge_type=data.get("edge_type", "causal"),
        temporal_lag=data.get("temporal_lag", 0.0),
    )


def _content_hash(entities: list[dict], edges: list[dict]) -> str:
    """Compute a deterministic SHA-256 hash of DAG content."""
    content = json.dumps(
        {
            "entities": sorted(entities, key=lambda e: e["entity_id"]),
            "edges": sorted(edges, key=lambda e: (e["source"], e["target"])),
        },
        sort_keys=True,
    )
    return hashlib.sha256(content.encode()).hexdigest()


def serialize_dag(dag: WorldDAG) -> tuple[list[dict], list[dict]]:
    """Serialize a WorldDAG to lists of entity and edge dicts."""
    entities = [_entity_to_dict(e) for e in dag.entities.values()]
    edges = [_edge_to_dict(e) for e in dag.edges.values()]
    return entities, edges


def deserialize_dag(entities: list[dict], edges: list[dict]) -> WorldDAG:
    """Reconstruct a WorldDAG from serialized entity and edge dicts."""
    dag = WorldDAG()
    for e_data in entities:
        dag.add_entity(_entity_from_dict(e_data))
    for edge_data in edges:
        dag.add_edge(_edge_from_dict(edge_data))
    return dag


def create_checkpoint(
    dag: WorldDAG,
    name: str,
    version: str = "1.0.0",
    metadata: dict | None = None,
) -> Checkpoint:
    """Create a checkpoint from a WorldDAG."""
    entities, edges = serialize_dag(dag)
    return Checkpoint(
        name=name,
        version=version,
        timestamp=time.time(),
        content_hash=_content_hash(entities, edges),
        entities=tuple(entities),
        edges=tuple(edges),
        metadata=metadata or {},
    )


def restore_from_checkpoint(checkpoint: Checkpoint) -> WorldDAG:
    """Restore a WorldDAG from a checkpoint."""
    return deserialize_dag(list(checkpoint.entities), list(checkpoint.edges))


class CheckpointStore:
    """Filesystem-backed checkpoint storage with versioning.

    Stores checkpoints as JSON files in a directory structure:
    ``<base_dir>/<model_name>/<version>.json``
    """

    def __init__(self, base_dir: str | Path) -> None:
        self._base_dir = Path(base_dir)

    def save(self, checkpoint: Checkpoint) -> Path:
        """Save a checkpoint to disk. Returns the file path."""
        model_dir = self._base_dir / checkpoint.name
        model_dir.mkdir(parents=True, exist_ok=True)
        path = model_dir / f"{checkpoint.version}.json"
        path.write_text(json.dumps(checkpoint.to_dict(), indent=2))
        return path

    def load(self, name: str, version: str) -> Checkpoint:
        """Load a checkpoint by name and version."""
        path = self._base_dir / name / f"{version}.json"
        if not path.exists():
            msg = f"Checkpoint not found: {name}@{version}"
            raise FileNotFoundError(msg)
        data = json.loads(path.read_text())
        return Checkpoint.from_dict(data)

    def list_models(self) -> list[str]:
        """List all model names with saved checkpoints."""
        if not self._base_dir.exists():
            return []
        return sorted(d.name for d in self._base_dir.iterdir() if d.is_dir())

    def list_versions(self, name: str) -> list[str]:
        """List all versions of a model, newest first."""
        model_dir = self._base_dir / name
        if not model_dir.exists():
            return []
        versions = [p.stem for p in model_dir.glob("*.json")]
        return sorted(versions, reverse=True)

    def latest(self, name: str) -> Checkpoint | None:
        """Load the latest version of a model."""
        versions = self.list_versions(name)
        if not versions:
            return None
        return self.load(name, versions[0])

    def delete(self, name: str, version: str) -> bool:
        """Delete a specific checkpoint. Returns True if it existed."""
        path = self._base_dir / name / f"{version}.json"
        if path.exists():
            path.unlink()
            return True
        return False

    def delete_model(self, name: str) -> int:
        """Delete all checkpoints for a model. Returns count deleted."""
        model_dir = self._base_dir / name
        if not model_dir.exists():
            return 0
        count = 0
        for p in model_dir.glob("*.json"):
            p.unlink()
            count += 1
        if model_dir.exists():
            model_dir.rmdir()
        return count
