"""Tests for `abreka exp remote ...` recovery commands."""

from __future__ import annotations

import json
import os
from unittest.mock import patch

from click.testing import CliRunner

from abreka.cli import main
from abreka.models import Experiment, ExperimentStatus
from abreka.persistence import create_experiment_dir, load_experiment, save_experiment, update_index_entry
from abreka.run_executor import RunExecutionResult


def _seed_remote_experiment(abreka_project, *, status: ExperimentStatus = ExperimentStatus.FAILED):
    abreka_dir = abreka_project / "abreka"
    exp = Experiment(id="0001", status=status, hypothesis="Remote baseline")
    exp_dir = create_experiment_dir(abreka_dir, exp.id)
    save_experiment(exp, abreka_dir)
    update_index_entry(abreka_dir, exp)
    (exp_dir / "run.py").write_text("print('hello')\n")

    metadata = {
        "mode": "remote",
        "state": "interrupted",
        "host": "gpu-box",
        "remote_project_dir": "/workspace/test-project",
        "remote_pid": "4242",
        "timeout_seconds": 1800,
        "remote_paths": {
            "artifacts_dir": "/workspace/test-project/abreka/experiments/0001/artifacts",
            "log_path": "/workspace/test-project/abreka/experiments/0001/artifacts/run_output.txt",
            "status_path": "/workspace/test-project/abreka/experiments/0001/artifacts/remote_status.json",
        },
    }
    (exp_dir / "artifacts" / "remote_run.json").write_text(json.dumps(metadata, indent=2) + "\n")
    return abreka_dir, exp_dir


class TestExpRemoteList:
    def test_list_empty(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)

        result = runner.invoke(main, ["exp", "remote", "list"])

        assert result.exit_code == 0, result.output
        assert "No remote runs found" in result.output

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

        result = runner.invoke(main, ["exp", "remote", "list"])

        assert result.exit_code == 0, result.output
        assert "0001" in result.output
        assert "interrupted" in result.output
        assert "gpu-box" in result.output


class TestExpRemoteStatus:
    def test_status_requires_remote_run_metadata(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["exp", "new", "--hypothesis", "Local only"])
        assert result.exit_code == 0, result.output

        result = runner.invoke(main, ["exp", "remote", "status", "1"])
        assert result.exit_code != 0
        assert "remote_run.json" in result.output

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

        with patch(
            "abreka.cli.exp.remote.RemoteRunExecutor.inspect_recovery",
            return_value={
                "remote_reachable": True,
                "pid_alive": True,
                "status": None,
                "status_exists": False,
                "remote_log_size": 1234,
            },
        ):
            result = runner.invoke(main, ["exp", "remote", "status", "1"])

        assert result.exit_code == 0, result.output
        assert "Experiment: 0001" in result.output
        assert "Remote host: gpu-box" in result.output
        assert "Remote process: alive" in result.output
        assert "Remote status file: not present" in result.output


class TestExpRemoteSync:
    def test_sync_pulls_artifacts(self, abreka_project):
        _, exp_dir = _seed_remote_experiment(abreka_project)
        runner = CliRunner()
        os.chdir(abreka_project)

        def fake_sync(self, *, remote_paths, local_artifacts_dir, metadata_path):
            (local_artifacts_dir / "run_output.txt").write_text("epoch 1\n")
            (local_artifacts_dir / "remote_status.json").write_text(
                '{"returncode": 0, "timed_out": false}\n'
            )

        with patch("abreka.cli.exp.remote.RemoteRunExecutor.sync_recovery_artifacts", fake_sync):
            result = runner.invoke(main, ["exp", "remote", "sync", "1"])

        assert result.exit_code == 0, result.output
        assert "Synced remote artifacts for experiment 0001" in result.output
        assert (exp_dir / "artifacts" / "run_output.txt").read_text() == "epoch 1\n"


class TestExpRemoteResume:
    def test_resume_recovers_metrics_and_evaluates_by_default(self, abreka_project):
        abreka_dir, exp_dir = _seed_remote_experiment(abreka_project)
        exp = load_experiment(abreka_dir, "0001")
        exp.failure_reason = "Lost contact with remote run while polling."
        save_experiment(exp, abreka_dir)
        update_index_entry(abreka_dir, exp)

        runner = CliRunner()
        os.chdir(abreka_project)

        def fake_resume(self, *, output_path, metadata_path, remote_paths):
            output_path.write_text("epoch 1\nMETRIC:val:accuracy=0.91\n")
            (output_path.parent / "remote_status.json").write_text(
                '{"returncode": 0, "timed_out": false}\n'
            )
            return RunExecutionResult(
                stdout="epoch 1\nMETRIC:val:accuracy=0.91\n",
                stderr="",
                returncode=0,
                timed_out=False,
            )

        with (
            patch("abreka.cli.exp.remote.RemoteRunExecutor.resume_recovery", fake_resume),
            patch("abreka.cli.exp.remote.AutonomousRunner._evaluate") as mock_evaluate,
        ):
            result = runner.invoke(main, ["exp", "remote", "resume", "1"])

        assert result.exit_code == 0, result.output
        mock_evaluate.assert_called_once_with("0001")

        exp = load_experiment(abreka_dir, "0001")
        assert exp.status == ExperimentStatus.RUNNING
        assert exp.failure_reason == ""
        assert exp.metrics == {"val": {"accuracy": 0.91}}
        assert exp.runtime_costs["run"]["attempts"] == 1
        assert exp.runtime_costs["run"]["last_status"] == "ok"
        assert (exp_dir / "artifacts" / "run_output.txt").is_file()

    def test_resume_can_skip_evaluate(self, abreka_project):
        abreka_dir, _ = _seed_remote_experiment(abreka_project)
        runner = CliRunner()
        os.chdir(abreka_project)

        def fake_resume(self, *, output_path, metadata_path, remote_paths):
            output_path.write_text("METRIC:val:accuracy=0.77\n")
            return RunExecutionResult(
                stdout="METRIC:val:accuracy=0.77\n",
                stderr="",
                returncode=0,
                timed_out=False,
            )

        with (
            patch("abreka.cli.exp.remote.RemoteRunExecutor.resume_recovery", fake_resume),
            patch("abreka.cli.exp.remote.AutonomousRunner._evaluate") as mock_evaluate,
        ):
            result = runner.invoke(main, ["exp", "remote", "resume", "1", "--no-evaluate"])

        assert result.exit_code == 0, result.output
        mock_evaluate.assert_not_called()
        exp = load_experiment(abreka_dir, "0001")
        assert exp.metrics == {"val": {"accuracy": 0.77}}


class TestExpRemoteStop:
    def test_stop_marks_experiment_failed(self, abreka_project):
        abreka_dir, _ = _seed_remote_experiment(abreka_project, status=ExperimentStatus.RUNNING)
        runner = CliRunner()
        os.chdir(abreka_project)

        with patch("abreka.cli.exp.remote.RemoteRunExecutor.stop_recovery"):
            result = runner.invoke(main, ["exp", "remote", "stop", "1"])

        assert result.exit_code == 0, result.output
        exp = load_experiment(abreka_dir, "0001")
        assert exp.status == ExperimentStatus.FAILED
        assert exp.failure_reason == "Remote run stopped by user"
