"""Tests for runtime-cost persistence helpers."""

from __future__ import annotations

from datetime import datetime, timezone

from abreka.models import Experiment
from abreka.persistence import create_experiment_dir, load_experiment, save_experiment
from abreka.runtime_costs import record_stage_runtime


class TestRecordStageRuntime:
    def test_creates_stage_cost_entry(self, abreka_project):
        abreka_dir = abreka_project / "abreka"
        exp = Experiment(id="0001")
        create_experiment_dir(abreka_dir, exp.id)
        save_experiment(exp, abreka_dir)

        record_stage_runtime(
            abreka_dir,
            exp.id,
            stage="propose",
            elapsed_seconds=1.25,
            status="ok",
            started_at=datetime(2026, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
            finished_at=datetime(2026, 1, 1, 0, 0, 1, tzinfo=timezone.utc),
        )

        loaded = load_experiment(abreka_dir, exp.id)
        assert loaded.runtime_costs["propose"]["attempts"] == 1
        assert loaded.runtime_costs["propose"]["seconds_total"] == 1.25
        assert loaded.runtime_costs["propose"]["seconds_last"] == 1.25
        assert loaded.runtime_costs["propose"]["last_status"] == "ok"

    def test_accumulates_attempts_and_total_seconds(self, abreka_project):
        abreka_dir = abreka_project / "abreka"
        exp = Experiment(id="0001")
        create_experiment_dir(abreka_dir, exp.id)
        save_experiment(exp, abreka_dir)

        start = datetime(2026, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
        end = datetime(2026, 1, 1, 0, 0, 1, tzinfo=timezone.utc)

        record_stage_runtime(
            abreka_dir,
            exp.id,
            stage="run",
            elapsed_seconds=2.0,
            status="ok",
            started_at=start,
            finished_at=end,
        )
        record_stage_runtime(
            abreka_dir,
            exp.id,
            stage="run",
            elapsed_seconds=3.5,
            status="timed_out",
            started_at=start,
            finished_at=end,
        )

        loaded = load_experiment(abreka_dir, exp.id)
        assert loaded.runtime_costs["run"]["attempts"] == 2
        assert loaded.runtime_costs["run"]["seconds_total"] == 5.5
        assert loaded.runtime_costs["run"]["seconds_last"] == 3.5
        assert loaded.runtime_costs["run"]["last_status"] == "timed_out"
