"""Tests for abreka.run_executor."""

from __future__ import annotations

import io
import json
import subprocess
from pathlib import Path

import pytest
from rich.console import Console

from abreka.models import Config
from abreka.run_executor import LocalRunExecutor, RemoteRunExecutor, build_run_executor


class TestBuildRunExecutor:
    def test_default_mode_is_local(self, abreka_project: Path):
        config = Config()
        executor = build_run_executor(
            abreka_project / "abreka",
            config,
            Console(file=io.StringIO(), force_terminal=False),
        )
        assert isinstance(executor, LocalRunExecutor)

    def test_remote_mode_returns_remote_executor(self, abreka_project: Path):
        config = Config()
        config.run.mode = "remote"
        config.run.remote.host = "gpu-box"
        config.run.remote.project_dir = "/workspace/test-project"

        executor = build_run_executor(
            abreka_project / "abreka",
            config,
            Console(file=io.StringIO(), force_terminal=False),
        )

        assert isinstance(executor, RemoteRunExecutor)
        assert executor.remote.host == "gpu-box"
        assert executor.remote.project_dir == "/workspace/test-project"

    def test_remote_mode_requires_host_and_project_dir(self, abreka_project: Path):
        config = Config()
        config.run.mode = "remote"

        with pytest.raises(ValueError, match="host.*project_dir"):
            build_run_executor(
                abreka_project / "abreka",
                config,
                Console(file=io.StringIO(), force_terminal=False),
            )


class TestLocalRunExecutor:
    def test_execute_streams_to_output_file(self, tmp_path: Path, monkeypatch):
        project_root = tmp_path / "project"
        run_py = project_root / "abreka" / "experiments" / "0001" / "run.py"
        run_py.parent.mkdir(parents=True)
        run_py.write_text("print('hello')\n")
        output_path = run_py.parent / "artifacts" / "run_output.txt"

        class FakeProcess:
            def __init__(self, *args, **kwargs):
                self.returncode = 0
                out = kwargs["stdout"]
                out.write("line 1\n")
                out.write("METRIC:val:accuracy=0.9\n")
                out.write("stderr line\n")
                out.flush()

            def wait(self, timeout=None):
                return 0

            def terminate(self):
                return None

            def kill(self):
                return None

        monkeypatch.setattr("abreka.run_executor.subprocess.Popen", FakeProcess)

        executor = LocalRunExecutor(
            project_root=project_root,
            console=Console(file=io.StringIO(), force_terminal=False),
        )
        result = executor.execute(
            run_py=run_py,
            output_path=output_path,
            timeout_seconds=30,
        )

        assert result.returncode == 0
        assert result.timed_out is False
        assert result.stdout == ""
        output = output_path.read_text()
        assert "line 1" in output
        assert "METRIC:val:accuracy=0.9" in output
        assert "stderr line" in output

    def test_execute_timeout_terminates_and_marks_timed_out(self, tmp_path: Path, monkeypatch):
        project_root = tmp_path / "project"
        run_py = project_root / "abreka" / "experiments" / "0001" / "run.py"
        run_py.parent.mkdir(parents=True)
        run_py.write_text("print('hello')\n")
        output_path = run_py.parent / "artifacts" / "run_output.txt"

        class FakeProcess:
            def __init__(self, *args, **kwargs):
                self.returncode = 143
                self._wait_calls = 0
                kwargs["stdout"].write("running...\n")
                kwargs["stdout"].flush()

            def wait(self, timeout=None):
                self._wait_calls += 1
                if self._wait_calls == 1:
                    raise subprocess.TimeoutExpired(cmd="uv run", timeout=timeout)
                return self.returncode

            def terminate(self):
                return None

            def kill(self):
                return None

        monkeypatch.setattr("abreka.run_executor.subprocess.Popen", FakeProcess)

        executor = LocalRunExecutor(
            project_root=project_root,
            console=Console(file=io.StringIO(), force_terminal=False),
        )
        result = executor.execute(
            run_py=run_py,
            output_path=output_path,
            timeout_seconds=1,
        )

        assert result.timed_out is True
        assert result.returncode == 143
        assert "running..." in output_path.read_text()


class TestRemoteRunExecutor:
    def test_sync_project_to_remote_includes_shared_artifacts_only(self, tmp_path: Path, monkeypatch):
        project_root = tmp_path / "project"
        project_root.mkdir()

        config = Config()
        config.run.mode = "remote"
        config.run.remote.host = "gpu-box"
        config.run.remote.project_dir = "/workspace/test-project"

        executor = RemoteRunExecutor(
            project_root=project_root,
            remote=config.run.remote,
            console=Console(file=io.StringIO(), force_terminal=False),
        )

        events: list[object] = []

        monkeypatch.setattr(
            executor,
            "_ssh",
            lambda command, **kwargs: events.append(("ssh", command, kwargs)),
        )

        def fake_run(command, check=True, **kwargs):
            events.append(("run", command, check, kwargs))
            return None

        monkeypatch.setattr("abreka.run_executor.subprocess.run", fake_run)

        executor._sync_project_to_remote()

        assert events[0] == ("ssh", "mkdir -p /workspace/test-project", {})
        rsync_command = events[1][1]
        assert rsync_command[:3] == ["rsync", "-az", "--delete"]
        assert rsync_command[-2:] == [
            f"{project_root}/",
            "gpu-box:/workspace/test-project/",
        ]

        rules = list(zip(rsync_command[3:-2:2], rsync_command[4:-2:2]))
        assert rules == [
            ("--exclude", ".git/"),
            ("--exclude", ".venv/"),
            ("--exclude", ".pytest_cache/"),
            ("--exclude", ".mypy_cache/"),
            ("--exclude", ".ruff_cache/"),
            ("--exclude", "__pycache__/"),
            ("--include", "abreka/artifacts/"),
            ("--include", "abreka/artifacts/shared/"),
            ("--include", "abreka/artifacts/shared/**"),
            ("--exclude", "abreka/artifacts/**"),
            ("--exclude", "abreka/experiments/*/artifacts/"),
        ]

    def test_execute_runs_remote_lifecycle_and_writes_metadata(self, tmp_path: Path, monkeypatch):
        project_root = tmp_path / "project"
        run_py = project_root / "abreka" / "experiments" / "0001" / "run.py"
        artifacts_dir = run_py.parent / "artifacts"
        artifacts_dir.mkdir(parents=True)
        run_py.write_text("print('hello')\n")
        output_path = artifacts_dir / "run_output.txt"

        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 = 1
        config.run.remote.startup_command = "./scripts/start-gpu.sh"
        config.run.remote.teardown_command = "./scripts/stop-gpu.sh"

        executor = RemoteRunExecutor(
            project_root=project_root,
            remote=config.run.remote,
            console=Console(file=io.StringIO(), force_terminal=False),
        )

        events: list[object] = []

        monkeypatch.setattr(executor, "_check_local_dependencies", lambda: events.append("check_deps"))
        monkeypatch.setattr(
            executor,
            "_run_local_hook",
            lambda label, command: events.append((label, command)),
        )
        monkeypatch.setattr(executor, "_sync_project_to_remote", lambda: events.append("sync_to_remote"))
        monkeypatch.setattr(executor, "_run_remote_bootstrap", lambda: events.append("bootstrap"))
        monkeypatch.setattr(
            executor,
            "_launch_remote_process",
            lambda run_relative_path, timeout_seconds, remote_paths: (
                events.append(("launch", str(run_relative_path), timeout_seconds, remote_paths["log_path"])),
                "4242",
            )[1],
        )

        def fake_poll(remote_paths, output_path):
            events.append(("poll", remote_paths["status_path"]))
            output_path.write_text("epoch 1\nMETRIC:val:accuracy=0.9\n")
            return {"returncode": 0, "timed_out": False}

        monkeypatch.setattr(executor, "_poll_remote_run", fake_poll)

        def fake_sync_back(remote_paths, local_artifacts_dir):
            events.append(("sync_back", remote_paths["artifacts_dir"], local_artifacts_dir))
            (local_artifacts_dir / "remote_status.json").write_text(
                '{"returncode": 0, "timed_out": false}\n'
            )

        monkeypatch.setattr(executor, "_sync_artifacts_from_remote", fake_sync_back)

        result = executor.execute(
            run_py=run_py,
            output_path=output_path,
            timeout_seconds=60,
        )

        assert result.stdout == "epoch 1\nMETRIC:val:accuracy=0.9\n"
        assert result.returncode == 0
        assert result.timed_out is False
        assert events == [
            "check_deps",
            ("startup", "./scripts/start-gpu.sh"),
            "sync_to_remote",
            "bootstrap",
            (
                "launch",
                "abreka/experiments/0001/run.py",
                60,
                "/workspace/test-project/abreka/experiments/0001/artifacts/run_output.txt",
            ),
            ("poll", "/workspace/test-project/abreka/experiments/0001/artifacts/remote_status.json"),
            (
                "sync_back",
                "/workspace/test-project/abreka/experiments/0001/artifacts",
                artifacts_dir,
            ),
            ("teardown", "./scripts/stop-gpu.sh"),
        ]

        metadata = json.loads((artifacts_dir / "remote_run.json").read_text())
        assert metadata["host"] == "gpu-box"
        assert metadata["remote_project_dir"] == "/workspace/test-project"
        assert metadata["remote_pid"] == "4242"
        assert metadata["state"] == "completed"
        assert metadata["status"]["returncode"] == 0
