"""Tests for abreka.cli.run_cmd — explore flag resolution and recovery."""

import os
from unittest.mock import patch

import pytest
from click.testing import CliRunner

from abreka.cli import main
from abreka.models import Experiment, ExperimentStatus
from abreka.persistence import create_experiment_dir, save_experiment, update_index_entry
from abreka.runner import parse_duration
from abreka.stage_sessions import stage_session_path


class TestParseDuration:
    def test_hours(self):
        assert parse_duration("2h") == 7200

    def test_minutes(self):
        assert parse_duration("30m") == 1800

    def test_seconds(self):
        assert parse_duration("90s") == 90

    def test_invalid(self):
        with pytest.raises(ValueError):
            parse_duration("abc")


class TestResolveExploreProbability:
    """Test the _resolve_explore_probability helper."""

    def test_explore_flag_sets_1(self):
        from abreka.cli.run_cmd import _resolve_explore_probability
        assert _resolve_explore_probability(explore=True, explore_prob=None, config_prob=0.0) == 1.0

    def test_explore_prob_option(self):
        from abreka.cli.run_cmd import _resolve_explore_probability
        assert _resolve_explore_probability(explore=False, explore_prob=0.1, config_prob=0.0) == 0.1

    def test_config_fallback(self):
        from abreka.cli.run_cmd import _resolve_explore_probability
        assert _resolve_explore_probability(explore=False, explore_prob=None, config_prob=0.2) == 0.2

    def test_no_explore(self):
        from abreka.cli.run_cmd import _resolve_explore_probability
        assert _resolve_explore_probability(explore=False, explore_prob=None, config_prob=0.0) is None

    def test_explore_and_prob_mutually_exclusive(self):
        from abreka.cli.run_cmd import _resolve_explore_probability
        with pytest.raises(ValueError, match="mutually exclusive"):
            _resolve_explore_probability(explore=True, explore_prob=0.5, config_prob=0.0)

    def test_explore_prob_must_be_positive(self):
        from abreka.cli.run_cmd import _resolve_explore_probability
        with pytest.raises(ValueError, match="must be in"):
            _resolve_explore_probability(explore=False, explore_prob=0.0, config_prob=0.0)

    def test_explore_prob_must_be_at_most_1(self):
        from abreka.cli.run_cmd import _resolve_explore_probability
        with pytest.raises(ValueError, match="must be in"):
            _resolve_explore_probability(explore=False, explore_prob=1.5, config_prob=0.0)

    def test_explore_prob_exactly_1(self):
        from abreka.cli.run_cmd import _resolve_explore_probability
        assert _resolve_explore_probability(explore=False, explore_prob=1.0, config_prob=0.0) == 1.0


def _seed_stage(abreka_project, exp_id: str, status: ExperimentStatus, stage: str | None = None):
    abreka_dir = abreka_project / "abreka"
    exp = Experiment(id=exp_id, status=status, hypothesis=f"exp {exp_id}")
    create_experiment_dir(abreka_dir, exp_id)
    save_experiment(exp, abreka_dir)
    update_index_entry(abreka_dir, exp)
    if stage is not None:
        path = stage_session_path(abreka_dir, exp_id, stage)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text('{"type":"session","version":3}\n')


class TestRecoverDiscovery:
    def test_finds_recoverable_local_stage_candidates(self, abreka_project):
        from abreka.cli.run_cmd import _find_recoverable_local_stage_candidates

        _seed_stage(abreka_project, "0001", ExperimentStatus.PROPOSING, "propose")
        _seed_stage(abreka_project, "0002", ExperimentStatus.IMPLEMENTING, "implement")
        _seed_stage(abreka_project, "0003", ExperimentStatus.TESTING, "test")
        _seed_stage(abreka_project, "0004", ExperimentStatus.RUNNING, "evaluate")
        _seed_stage(abreka_project, "0005", ExperimentStatus.RUNNING, None)

        found = _find_recoverable_local_stage_candidates(abreka_project / "abreka")

        assert [(c.kind, c.exp_id, c.stage) for c in found] == [
            ("local_stage", "0001", "propose"),
            ("local_stage", "0002", "implement"),
            ("local_stage", "0003", "test"),
            ("local_stage", "0004", "evaluate"),
        ]


class TestRecoverMode:
    def test_recover_can_resume_local_stage_without_duration(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        _seed_stage(abreka_project, "0001", ExperimentStatus.IMPLEMENTING, "implement")

        with (
            patch("abreka.cli.run_cmd._recover_candidate") as mock_recover,
            patch("abreka.cli.run_cmd.AutonomousRunner.run") as mock_run,
        ):
            result = runner.invoke(main, ["run", "--recover"])

        assert result.exit_code == 0, result.output
        mock_recover.assert_called_once()
        mock_run.assert_not_called()

    def test_recover_errors_when_none_found(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)

        result = runner.invoke(main, ["run", "--recover"])

        assert result.exit_code != 0
        assert "No recoverable runs found" in result.output

    def test_recover_errors_when_multiple_found(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        _seed_stage(abreka_project, "0001", ExperimentStatus.IMPLEMENTING, "implement")
        _seed_stage(abreka_project, "0002", ExperimentStatus.TESTING, "test")

        result = runner.invoke(main, ["run", "--recover"])

        assert result.exit_code != 0
        assert "Multiple recoverable runs found" in result.output
        assert "0001 implement" in result.output
        assert "0002 test" in result.output

    def test_recover_errors_when_local_and_remote_candidates_overlap(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        _seed_stage(abreka_project, "0001", ExperimentStatus.IMPLEMENTING, "implement")

        with patch("abreka.cli.run_cmd._find_recoverable_remote_run_ids", return_value=["0009"]):
            result = runner.invoke(main, ["run", "--recover"])

        assert result.exit_code != 0
        assert "Multiple recoverable runs found" in result.output
        assert "0001 implement" in result.output
        assert "0009 remote-run" in result.output

    def test_recover_then_continue_running_when_duration_or_step_given(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        _seed_stage(abreka_project, "0001", ExperimentStatus.TESTING, "test")

        with (
            patch("abreka.cli.run_cmd._recover_candidate") as mock_recover,
            patch("abreka.cli.run_cmd.AutonomousRunner.run") as mock_run,
        ):
            result = runner.invoke(main, ["run", "--step", "--recover"])

        assert result.exit_code == 0, result.output
        mock_recover.assert_called_once()
        mock_run.assert_called_once()
