"""Tests for abreka.models."""

from datetime import datetime, timezone

from abreka.models import BacklogItem, BacklogStatus, Config, Experiment, ExperimentStatus


class TestExperiment:
    def test_default_status(self):
        exp = Experiment(id="0001")
        assert exp.status == ExperimentStatus.PROPOSED

    def test_round_trip(self):
        exp = Experiment(
            id="0003",
            status=ExperimentStatus.COMPLETED,
            hypothesis="Test hypothesis",
            method="Test method",
            parents=["0001", "0002"],
            tags=["cnn", "baseline"],
            metrics={"val": {"accuracy": 0.95}, "train": {"loss": 0.01}},
            created_at=datetime(2026, 4, 8, 10, 0, 0, tzinfo=timezone.utc),
            completed_at=datetime(2026, 4, 8, 10, 5, 0, tzinfo=timezone.utc),
        )
        d = exp.to_dict()
        restored = Experiment.from_dict(d)
        assert restored.id == exp.id
        assert restored.status == exp.status
        assert restored.hypothesis == exp.hypothesis
        assert restored.method == exp.method
        assert restored.parents == exp.parents
        assert restored.tags == exp.tags
        assert restored.metrics == exp.metrics
        assert restored.created_at == exp.created_at
        assert restored.completed_at == exp.completed_at

    def test_failed_experiment(self):
        exp = Experiment(id="0001", status=ExperimentStatus.FAILED, failure_reason="OOM")
        d = exp.to_dict()
        assert d["failure"]["reason"] == "OOM"
        restored = Experiment.from_dict(d)
        assert restored.failure_reason == "OOM"

    def test_summary(self):
        exp = Experiment(
            id="0001",
            hypothesis="Test",
            method="Do the thing",
            metrics={"val": {"accuracy": 0.9}},
        )
        s = exp.summary()
        assert s["id"] == "0001"
        assert s["hypothesis"] == "Test"
        assert s["method"] == "Do the thing"
        assert "metrics" in s

    def test_summary_includes_method(self):
        exp = Experiment(id="0002", method="Change learning rate to 1e-4")
        s = exp.summary()
        assert s["method"] == "Change learning rate to 1e-4"

    def test_summary_empty_method(self):
        exp = Experiment(id="0003")
        s = exp.summary()
        assert s["method"] == ""

    def test_runtime_costs_round_trip_and_summary_total(self):
        exp = Experiment(
            id="0004",
            runtime_costs={
                "propose": {
                    "attempts": 1,
                    "seconds_total": 2.5,
                    "seconds_last": 2.5,
                    "last_status": "ok",
                },
                "run": {
                    "attempts": 2,
                    "seconds_total": 10.0,
                    "seconds_last": 4.0,
                    "last_status": "timed_out",
                },
            },
        )

        d = exp.to_dict()
        assert "runtime_costs" in d

        restored = Experiment.from_dict(d)
        assert restored.runtime_costs == exp.runtime_costs

        summary = restored.summary()
        assert summary["runtime_seconds_total"] == 12.5

    def test_failures_round_trip(self):
        exp = Experiment(
            id="0005",
            failures=[
                {
                    "stage": "run",
                    "attempt": 1,
                    "repair_attempt": 0,
                    "reason": "run.py exited with code 1",
                    "returncode": 1,
                    "timed_out": False,
                    "output_path": "artifacts/run_failures/run_attempt_1.txt",
                    "occurred_at": "2026-05-02T00:00:00+00:00",
                }
            ],
        )

        d = exp.to_dict()
        assert d["failures"][0]["stage"] == "run"
        restored = Experiment.from_dict(d)
        assert restored.failures == exp.failures


class TestConfig:
    def test_round_trip(self):
        config = Config()
        d = config.to_dict()
        restored = Config.from_dict(d)
        assert restored.project.name == config.project.name
        assert restored.agent.provider == config.agent.provider
        assert restored.experiment.primary_metric == config.experiment.primary_metric
        assert restored.run.mode == config.run.mode
        assert restored.run.max_repair_attempts == config.run.max_repair_attempts

    def test_run_repair_attempts_default_and_round_trip(self):
        config = Config()
        assert config.run.max_repair_attempts == 3

        config.run.max_repair_attempts = 2
        d = config.to_dict()
        assert d["run"]["max_repair_attempts"] == 2
        restored = Config.from_dict(d)
        assert restored.run.max_repair_attempts == 2

    def test_run_repair_attempts_missing_in_source_defaults_to_three(self):
        d = Config().to_dict()
        d["run"].pop("max_repair_attempts", None)

        restored = Config.from_dict(d)
        assert restored.run.max_repair_attempts == 3

    def test_remote_run_config_round_trip(self):
        config = Config()
        config.run.mode = "remote"
        config.run.remote.host = "gpu-box"
        config.run.remote.project_dir = "/workspace/test-project"
        config.run.remote.bootstrap_command = "uv sync --frozen"
        config.run.remote.poll_seconds = 7
        config.run.remote.startup_command = "./scripts/start-gpu.sh"
        config.run.remote.teardown_command = "./scripts/stop-gpu.sh"

        d = config.to_dict()
        assert d["run"]["mode"] == "remote"
        assert d["run"]["remote"]["host"] == "gpu-box"

        restored = Config.from_dict(d)
        assert restored.run.mode == "remote"
        assert restored.run.remote.host == "gpu-box"
        assert restored.run.remote.project_dir == "/workspace/test-project"
        assert restored.run.remote.bootstrap_command == "uv sync --frozen"
        assert restored.run.remote.poll_seconds == 7
        assert restored.run.remote.startup_command == "./scripts/start-gpu.sh"
        assert restored.run.remote.teardown_command == "./scripts/stop-gpu.sh"

    def test_optional_fields(self):
        config = Config()
        config.agent.propose_model = "claude-opus-4-20250514"
        d = config.to_dict()
        assert d["agent"]["propose_model"] == "claude-opus-4-20250514"
        restored = Config.from_dict(d)
        assert restored.agent.propose_model == "claude-opus-4-20250514"

    def test_explore_probability_default(self):
        config = Config()
        assert config.experiment.explore_probability == 0.0

    def test_explore_probability_round_trip(self):
        config = Config()
        config.experiment.explore_probability = 0.2
        d = config.to_dict()
        assert d["experiment"]["explore_probability"] == 0.2
        restored = Config.from_dict(d)
        assert restored.experiment.explore_probability == 0.2

    def test_explore_probability_omitted_from_dict_when_zero(self):
        config = Config()
        d = config.to_dict()
        assert "explore_probability" not in d["experiment"]

    def test_explore_probability_missing_in_source(self):
        """Old config.toml files without explore_probability should default to 0.0."""
        d = Config().to_dict()
        # Simulate old config without the field
        d["experiment"].pop("explore_probability", None)
        restored = Config.from_dict(d)
        assert restored.experiment.explore_probability == 0.0

    def test_session_capture_mode_default_and_round_trip(self):
        config = Config()
        assert config.agent.session_capture_mode == "slim"

        config.agent.session_capture_mode = "minimal"
        d = config.to_dict()
        assert d["agent"]["session_capture_mode"] == "minimal"

        restored = Config.from_dict(d)
        assert restored.agent.session_capture_mode == "minimal"

    def test_session_capture_mode_missing_in_source_defaults_to_slim(self):
        d = Config().to_dict()
        d["agent"].pop("session_capture_mode", None)

        restored = Config.from_dict(d)
        assert restored.agent.session_capture_mode == "slim"


class TestBacklogItem:
    def test_default_status(self):
        item = BacklogItem(id="0001", title="Test idea")
        assert item.status == BacklogStatus.OPEN

    def test_round_trip(self):
        item = BacklogItem(
            id="0003",
            title="Try attention mechanism",
            description="Based on findings from exp 0002",
            tags=["attention", "arch"],
            priority=5,
            source_exp="0002",
            resolved_exp="0010",
            status=BacklogStatus.RESOLVED,
            created_at=datetime(2026, 4, 8, 10, 0, 0, tzinfo=timezone.utc),
            updated_at=datetime(2026, 4, 9, 12, 0, 0, tzinfo=timezone.utc),
        )
        d = item.to_dict()
        restored = BacklogItem.from_dict(d)
        assert restored.id == item.id
        assert restored.title == item.title
        assert restored.description == item.description
        assert restored.tags == item.tags
        assert restored.priority == item.priority
        assert restored.source_exp == item.source_exp
        assert restored.resolved_exp == item.resolved_exp
        assert restored.status == item.status
        assert restored.created_at == item.created_at
        assert restored.updated_at == item.updated_at

    def test_dismissed_item(self):
        item = BacklogItem(
            id="0001",
            title="Bad idea",
            status=BacklogStatus.DISMISSED,
            dismiss_reason="Already tried in exp 0005",
        )
        d = item.to_dict()
        assert d["dismiss_reason"] == "Already tried in exp 0005"
        restored = BacklogItem.from_dict(d)
        assert restored.dismiss_reason == "Already tried in exp 0005"

    def test_summary(self):
        item = BacklogItem(
            id="0001",
            title="Try dropout",
            description="Might help with overfitting",
            priority=3,
            source_exp="0002",
            tags=["regularization"],
        )
        s = item.summary()
        assert s["id"] == "0001"
        assert s["title"] == "Try dropout"
        assert s["priority"] == 3
        assert s["source_exp"] == "0002"
        assert s["tags"] == ["regularization"]
        assert s["status"] == "open"

    def test_optional_fields_omitted(self):
        item = BacklogItem(id="0001", title="Simple idea")
        d = item.to_dict()
        assert "source_exp" not in d
        assert "resolved_exp" not in d
        assert "dismiss_reason" not in d
