"""Tests for kernel checkpoint — model serialization and versioning."""

from __future__ import annotations

import json
import time

import pytest

from orxaq.kernel.checkpoint import (
    Checkpoint,
    CheckpointStore,
    _content_hash,
    _edge_from_dict,
    _edge_to_dict,
    _entity_from_dict,
    _entity_to_dict,
    create_checkpoint,
    deserialize_dag,
    restore_from_checkpoint,
    serialize_dag,
)
from orxaq.kernel.types import (
    CausalEdge,
    CausalRole,
    Entity,
    SensitivityLevel,
    VariableType,
)
from orxaq.kernel.world_dag import WorldDAG

# ── Fixtures ──────────────────────────────────────────────────


@pytest.fixture()
def sample_dag() -> WorldDAG:
    """Build a small DAG with 3 entities and 2 edges."""
    dag = WorldDAG()
    dag.add_entity(Entity(entity_id="A", name="Income", domain="credit"))
    dag.add_entity(Entity(entity_id="B", name="Score", causal_role=CausalRole.MEDIATOR))
    dag.add_entity(
        Entity(
            entity_id="C",
            name="Default",
            causal_role=CausalRole.OUTCOME,
            sensitivity=SensitivityLevel.CONFIDENTIAL,
            variable_type=VariableType.DYNAMIC,
            units="USD",
            bounds=(0.0, 1.0),
            regulatory_tags=("SR11-7", "CECL"),
            metadata=(("source", "model"), ("version", "1")),
        )
    )
    dag.add_edge(CausalEdge(source="A", target="B", confidence=0.9, mechanism="direct"))
    dag.add_edge(
        CausalEdge(
            source="B",
            target="C",
            confidence=0.85,
            soliton_score=0.7,
            provenance="expert",
            is_immutable=True,
            edge_type="causal",
            temporal_lag=1.5,
        )
    )
    return dag


@pytest.fixture()
def store(tmp_path) -> CheckpointStore:
    return CheckpointStore(tmp_path / "checkpoints")


# ── Entity serialization ─────────────────────────────────────


class TestEntitySerialization:
    def test_roundtrip_simple(self) -> None:
        e = Entity(entity_id="x", name="X")
        d = _entity_to_dict(e)
        restored = _entity_from_dict(d)
        assert restored.entity_id == e.entity_id
        assert restored.name == e.name
        assert restored.causal_role == CausalRole.EXOGENOUS
        assert restored.variable_type == VariableType.STATIC

    def test_roundtrip_full_fields(self) -> None:
        e = Entity(
            entity_id="z",
            name="Z-var",
            description="test desc",
            causal_role=CausalRole.COLLIDER,
            variable_type=VariableType.MACRO,
            sensitivity=SensitivityLevel.RESTRICTED,
            domain="finance",
            units="bps",
            bounds=(-100.0, 100.0),
            regulatory_tags=("SR11-7", "CCAR"),
            metadata=(("k1", "v1"), ("k2", "v2")),
        )
        d = _entity_to_dict(e)
        restored = _entity_from_dict(d)
        assert restored == e

    def test_dict_is_json_serializable(self) -> None:
        e = Entity(entity_id="j", name="JSON", bounds=(None, 5.0))
        d = _entity_to_dict(e)
        json_str = json.dumps(d)
        assert isinstance(json_str, str)


# ── Edge serialization ───────────────────────────────────────


class TestEdgeSerialization:
    def test_roundtrip_simple(self) -> None:
        edge = CausalEdge(source="A", target="B")
        d = _edge_to_dict(edge)
        restored = _edge_from_dict(d)
        assert restored.source == "A"
        assert restored.target == "B"
        assert restored.confidence == 0.5

    def test_roundtrip_full_fields(self) -> None:
        edge = CausalEdge(
            source="X",
            target="Y",
            mechanism="causal",
            confidence=0.95,
            soliton_score=0.8,
            provenance="discovery",
            is_immutable=True,
            edge_type="latent",
            temporal_lag=2.5,
        )
        d = _edge_to_dict(edge)
        restored = _edge_from_dict(d)
        assert restored == edge


# ── DAG serialization ────────────────────────────────────────


class TestDAGSerialization:
    def test_roundtrip(self, sample_dag: WorldDAG) -> None:
        entities, edges = serialize_dag(sample_dag)
        restored = deserialize_dag(entities, edges)
        assert len(restored.entities) == 3
        assert len(restored.edges) == 2
        assert restored.get_entity("A") is not None
        assert restored.get_entity("C").sensitivity == SensitivityLevel.CONFIDENTIAL

    def test_empty_dag(self) -> None:
        dag = WorldDAG()
        entities, edges = serialize_dag(dag)
        assert entities == []
        assert edges == []
        restored = deserialize_dag(entities, edges)
        assert len(restored.entities) == 0
        assert len(restored.edges) == 0

    def test_preserves_adjacency(self, sample_dag: WorldDAG) -> None:
        entities, edges = serialize_dag(sample_dag)
        restored = deserialize_dag(entities, edges)
        children_of_a = restored.children("A")
        assert "B" in children_of_a

    def test_preserves_edge_metadata(self, sample_dag: WorldDAG) -> None:
        entities, edges = serialize_dag(sample_dag)
        restored = deserialize_dag(entities, edges)
        bc_edge = None
        for e in restored.edges.values():
            if e.source == "B" and e.target == "C":
                bc_edge = e
                break
        assert bc_edge is not None
        assert bc_edge.is_immutable is True
        assert bc_edge.temporal_lag == 1.5
        assert bc_edge.provenance == "expert"


# ── Checkpoint ───────────────────────────────────────────────


class TestCheckpoint:
    def test_create_checkpoint(self, sample_dag: WorldDAG) -> None:
        cp = create_checkpoint(sample_dag, name="credit_model", version="1.0.0")
        assert cp.name == "credit_model"
        assert cp.version == "1.0.0"
        assert len(cp.entities) == 3
        assert len(cp.edges) == 2
        assert len(cp.content_hash) == 64  # SHA-256 hex

    def test_content_hash_deterministic(self, sample_dag: WorldDAG) -> None:
        cp1 = create_checkpoint(sample_dag, name="m", version="1")
        cp2 = create_checkpoint(sample_dag, name="m", version="1")
        assert cp1.content_hash == cp2.content_hash

    def test_content_hash_changes_on_modification(self, sample_dag: WorldDAG) -> None:
        cp1 = create_checkpoint(sample_dag, name="m", version="1")
        sample_dag.add_entity(Entity(entity_id="D", name="New"))
        cp2 = create_checkpoint(sample_dag, name="m", version="2")
        assert cp1.content_hash != cp2.content_hash

    def test_restore_from_checkpoint(self, sample_dag: WorldDAG) -> None:
        cp = create_checkpoint(sample_dag, name="test", version="1.0.0")
        restored = restore_from_checkpoint(cp)
        assert len(restored.entities) == len(sample_dag.entities)
        assert len(restored.edges) == len(sample_dag.edges)

    def test_to_dict_from_dict_roundtrip(self, sample_dag: WorldDAG) -> None:
        cp = create_checkpoint(
            sample_dag,
            name="model",
            version="2.1.0",
            metadata={"algorithm": "pc", "alpha": 0.05},
        )
        d = cp.to_dict()
        restored = Checkpoint.from_dict(d)
        assert restored.name == cp.name
        assert restored.version == cp.version
        assert restored.content_hash == cp.content_hash
        assert restored.metadata == {"algorithm": "pc", "alpha": 0.05}

    def test_checkpoint_has_timestamp(self, sample_dag: WorldDAG) -> None:
        before = time.time()
        cp = create_checkpoint(sample_dag, name="t", version="1")
        after = time.time()
        assert before <= cp.timestamp <= after

    def test_checkpoint_with_metadata(self, sample_dag: WorldDAG) -> None:
        cp = create_checkpoint(
            sample_dag,
            name="m",
            version="1",
            metadata={"trained_on": "credit_data_2024"},
        )
        assert cp.metadata["trained_on"] == "credit_data_2024"


# ── CheckpointStore ──────────────────────────────────────────


class TestCheckpointStore:
    def test_save_and_load(self, sample_dag: WorldDAG, store: CheckpointStore) -> None:
        cp = create_checkpoint(sample_dag, name="model_a", version="1.0.0")
        path = store.save(cp)
        assert path.exists()

        loaded = store.load("model_a", "1.0.0")
        assert loaded.name == "model_a"
        assert loaded.content_hash == cp.content_hash

    def test_list_models_empty(self, store: CheckpointStore) -> None:
        assert store.list_models() == []

    def test_list_models(self, sample_dag: WorldDAG, store: CheckpointStore) -> None:
        store.save(create_checkpoint(sample_dag, name="alpha", version="1.0"))
        store.save(create_checkpoint(sample_dag, name="beta", version="1.0"))
        models = store.list_models()
        assert models == ["alpha", "beta"]

    def test_list_versions(self, sample_dag: WorldDAG, store: CheckpointStore) -> None:
        store.save(create_checkpoint(sample_dag, name="m", version="1.0.0"))
        store.save(create_checkpoint(sample_dag, name="m", version="2.0.0"))
        store.save(create_checkpoint(sample_dag, name="m", version="1.1.0"))
        versions = store.list_versions("m")
        assert versions == ["2.0.0", "1.1.0", "1.0.0"]

    def test_latest(self, sample_dag: WorldDAG, store: CheckpointStore) -> None:
        store.save(create_checkpoint(sample_dag, name="m", version="1.0"))
        store.save(create_checkpoint(sample_dag, name="m", version="2.0"))
        latest = store.latest("m")
        assert latest is not None
        assert latest.version == "2.0"

    def test_latest_none_if_missing(self, store: CheckpointStore) -> None:
        assert store.latest("nonexistent") is None

    def test_delete_version(self, sample_dag: WorldDAG, store: CheckpointStore) -> None:
        store.save(create_checkpoint(sample_dag, name="m", version="1.0"))
        assert store.delete("m", "1.0") is True
        assert store.list_versions("m") == []

    def test_delete_nonexistent(self, store: CheckpointStore) -> None:
        assert store.delete("m", "9.9") is False

    def test_delete_model(self, sample_dag: WorldDAG, store: CheckpointStore) -> None:
        store.save(create_checkpoint(sample_dag, name="m", version="1.0"))
        store.save(create_checkpoint(sample_dag, name="m", version="2.0"))
        count = store.delete_model("m")
        assert count == 2
        assert store.list_models() == []

    def test_load_nonexistent_raises(self, store: CheckpointStore) -> None:
        with pytest.raises(FileNotFoundError):
            store.load("missing", "1.0")

    def test_save_creates_directory(self, tmp_path, sample_dag: WorldDAG) -> None:
        store = CheckpointStore(tmp_path / "deep" / "nested" / "checkpoints")
        cp = create_checkpoint(sample_dag, name="m", version="1.0")
        path = store.save(cp)
        assert path.exists()

    def test_roundtrip_preserves_dag_fidelity(
        self, sample_dag: WorldDAG, store: CheckpointStore
    ) -> None:
        cp = create_checkpoint(sample_dag, name="fidelity", version="1.0.0")
        store.save(cp)
        loaded = store.load("fidelity", "1.0.0")
        dag = restore_from_checkpoint(loaded)

        # Verify all entities match.
        for eid, entity in sample_dag.entities.items():
            restored_entity = dag.get_entity(eid)
            assert restored_entity is not None
            assert restored_entity.entity_id == entity.entity_id
            assert restored_entity.name == entity.name
            assert restored_entity.causal_role == entity.causal_role

        # Verify all edges match.
        assert len(dag.edges) == len(sample_dag.edges)


# ── Content hash ─────────────────────────────────────────────


class TestContentHash:
    def test_deterministic(self) -> None:
        entities = [{"entity_id": "A"}, {"entity_id": "B"}]
        edges = [{"source": "A", "target": "B"}]
        h1 = _content_hash(entities, edges)
        h2 = _content_hash(entities, edges)
        assert h1 == h2

    def test_order_independent(self) -> None:
        e1 = [{"entity_id": "B"}, {"entity_id": "A"}]
        e2 = [{"entity_id": "A"}, {"entity_id": "B"}]
        edges = [{"source": "A", "target": "B"}]
        assert _content_hash(e1, edges) == _content_hash(e2, edges)

    def test_different_content_different_hash(self) -> None:
        e1 = [{"entity_id": "A"}]
        e2 = [{"entity_id": "A"}, {"entity_id": "B"}]
        edges: list[dict] = []
        assert _content_hash(e1, edges) != _content_hash(e2, edges)


# ── Targeted coverage (line 243: delete_model for non-existent model) ──


class TestDeleteModelNonExistent:
    """Cover line 243: delete_model returns 0 for non-existent model."""

    def test_delete_nonexistent_returns_zero(self, tmp_path):
        store = CheckpointStore(base_dir=tmp_path)
        count = store.delete_model("nonexistent_model")
        assert count == 0
