from __future__ import annotations

import asyncio
import json
from pathlib import Path

from autocad_bench.orchestration import resume as resume_batch


class _FakeEc2:
    def __init__(self, instances: list[dict]) -> None:
        self.instances = instances

    def describe_instances(self, **_kwargs):
        return {"Reservations": [{"Instances": self.instances}]}


def _interrupted_rollout(batch: Path, task_id: str, instance_id: str) -> None:
    rollout = batch / "rollouts" / task_id
    rollout.mkdir(parents=True)
    (rollout / "run-state.json").write_text(
        json.dumps(
            {
                "task_id": task_id,
                "provider": "openai",
                "model_id": "gpt-test",
            }
        )
    )
    logs = batch / "logs"
    logs.mkdir()
    (logs / f"{task_id}.log").write_text(
        json.dumps(
            {
                "event": "sandbox_instance_launched",
                "instance_id": instance_id,
            }
        )
        + "\n"
    )


def test_discovery_accepts_parent_of_multiple_benchmark_runs(tmp_path: Path) -> None:
    benchmarks = tmp_path / "benchmarks"
    _interrupted_rollout(benchmarks / "run-a", "task-001", "i-a")
    _interrupted_rollout(benchmarks / "run-b", "task-002", "i-b")

    plan = resume_batch._discover_plan([benchmarks])

    assert {entry["rollout_id"] for entry in plan["entries"]} == {
        "run-a/task-001",
        "run-b/task-002",
    }


def test_empty_reconciliation_scan_exits_successfully(tmp_path: Path, capsys) -> None:
    benchmarks = tmp_path / "benchmarks"
    (benchmarks / "finished" / "rollouts").mkdir(parents=True)

    exit_code = resume_batch.main(["--run-root", str(benchmarks)])

    assert exit_code == 0
    assert json.loads(capsys.readouterr().out) == {
        "entries": 0,
        "message": "no interrupted rollout directories were discovered",
        "state": "idle",
    }


def test_completed_execution_with_incomplete_evaluation_is_cleanup_only(
    tmp_path: Path,
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-a")
    rollout = batch / "rollouts" / "task-001"
    (rollout / "result.json").write_text(
        json.dumps(
            {
                "task_id": "task-001",
                "completed": True,
                "stop_reason": "done",
                "artifact_bytes": 100,
            }
        )
    )

    plan = resume_batch._discover_plan([tmp_path / "benchmarks"])

    assert [entry["rollout_id"] for entry in plan["entries"]] == [
        "run-a/task-001"
    ]
    assert plan["entries"][0]["cleanup_only"] is True


def test_deferred_vision_only_retry_does_not_hold_a_windows_session(
    tmp_path: Path,
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-a")
    rollout = batch / "rollouts" / "task-001"
    (rollout / "result.json").write_text(
        json.dumps(
            {
                "task_id": "task-001",
                "completed": True,
                "stop_reason": "done",
                "artifact_bytes": 100,
                "evaluation": {"status": "retry_pending"},
            }
        )
    )

    try:
        resume_batch._discover_plan([tmp_path / "benchmarks"])
    except resume_batch.NoInterruptedRollouts:
        pass
    else:
        raise AssertionError("vision-only retry unexpectedly retained the worker")


def test_budget_exhaustion_does_not_restart_the_model(tmp_path: Path) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-a")
    rollout = batch / "rollouts" / "task-001"
    (batch / "batch-state.json").write_text(
        json.dumps({"batch_id": "batch-a"})
    )
    (rollout / "result.json").write_text(
        json.dumps(
            {
                "task_id": "task-001",
                "execution_status": "budget_exhausted",
                "stop_reason": "max_actions",
                "artifact_bytes": 100,
                "evaluation": {"status": "not_started"},
            }
        )
    )
    ec2 = _FakeEc2(
        [
            {
                "InstanceId": "i-a",
                "Tags": [
                    {"Key": "BatchId", "Value": "batch-a"},
                    {"Key": "RolloutId", "Value": "task-001"},
                ],
            }
        ]
    )

    plan = resume_batch._discover_plan([batch.parent], ec2=ec2)

    assert plan["entries"][0]["cleanup_only"] is True


def test_malformed_model_failure_mislabeled_as_infrastructure_is_cleanup_only(
    tmp_path: Path,
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-a")
    rollout = batch / "rollouts" / "task-001"
    (rollout / "result.json").write_text(
        json.dumps(
            {
                "execution_status": "infrastructure_failed",
                "stop_reason": "invalid_action_limit",
                "failure": "Model returned invalid computer-action arguments",
                "artifact_bytes": 0,
            }
        )
    )

    plan = resume_batch._discover_plan([batch.parent])

    assert plan["entries"][0]["cleanup_only"] is True
    assert plan["entries"][0]["recovery_exhausted"] is False


def test_recovery_attempt_limit_reaps_worker_without_another_model_turn(
    tmp_path: Path,
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-a")
    rollout = batch / "rollouts" / "task-001"
    result = {
        "stop_reason": "infrastructure_error",
        "failure": "BrokerError: transport unavailable",
        "artifact_bytes": 0,
    }
    (rollout / "result.json").write_text(json.dumps(result))
    (batch / "cloud-resume-state.json").write_text(
        json.dumps(
            {
                "entries": {
                    "run-a/task-001": {
                        "state": "infrastructure_failed",
                        "recovery_attempts": 3,
                        "last_result_fingerprint": (
                            resume_batch._result_recovery_fingerprint(result)
                        ),
                    }
                }
            }
        )
    )

    plan = resume_batch._discover_plan([batch.parent])

    assert plan["entries"][0]["cleanup_only"] is True
    assert plan["entries"][0]["recovery_exhausted"] is True


def test_discovery_caps_subprocess_recovery_concurrency(tmp_path: Path) -> None:
    benchmarks = tmp_path / "benchmarks"
    for index in range(6):
        _interrupted_rollout(
            benchmarks / f"run-{index}",
            f"basic-00{index + 1}",
            f"i-{index}",
        )

    plan = resume_batch._discover_plan([benchmarks])

    assert len(plan["entries"]) == 6
    assert plan["max_concurrency"] == 1


def test_discovery_uses_ec2_tags_when_launch_log_was_not_persisted(
    tmp_path: Path,
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-from-log")
    (batch / "logs" / "task-001.log").write_text("")
    (batch / "batch-state.json").write_text(
        json.dumps({"batch_id": "batch-a"})
    )
    ec2 = _FakeEc2(
        [
            {
                "InstanceId": "i-from-tags",
                "Tags": [
                    {"Key": "BatchId", "Value": "batch-a"},
                    {"Key": "RolloutId", "Value": "task-001"},
                ],
            }
        ]
    )

    plan = resume_batch._discover_plan([batch.parent], ec2=ec2)

    assert plan["entries"][0]["instance_id"] == "i-from-tags"


def test_discovery_skips_rollout_owned_by_live_primary_process(
    tmp_path: Path, monkeypatch
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-a")
    rollout = batch / "rollouts" / "task-001"
    (batch / "batch-state.json").write_text(
        json.dumps(
            {
                "batch_id": "batch-a",
                "rollouts": {
                    "task-001": {
                        "state": "running",
                        "pid": 1234,
                    }
                },
            }
        )
    )
    monkeypatch.setattr(resume_batch, "_pid_owns_rollout", lambda *_args, **_kwargs: True)

    try:
        resume_batch._discover_plan([batch.parent])
    except resume_batch.NoInterruptedRollouts:
        pass
    else:
        raise AssertionError("live primary rollout was scheduled for duplicate resume")


def test_shared_primary_pid_is_bound_to_batch_root(
    tmp_path: Path, monkeypatch
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    rollout = batch / "rollouts" / "task-001"
    rollout.mkdir(parents=True)
    checked_paths: list[Path] = []

    def owns(_pid: object, path: Path) -> bool:
        checked_paths.append(path)
        return True

    monkeypatch.setattr(resume_batch, "_pid_owns_rollout", owns)

    assert resume_batch._primary_rollout_is_active(
        {
            "rollouts": {
                "task-001": {
                    "state": "running",
                    "pid": 1234,
                    "controller_execution_mode": "shared",
                }
            }
        },
        rollout_id="task-001",
        output_dir=rollout,
    )
    assert checked_paths == [batch]


def test_discovery_recovers_active_row_after_primary_process_dies(
    tmp_path: Path, monkeypatch
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-a")
    (batch / "batch-state.json").write_text(
        json.dumps(
            {
                "batch_id": "batch-a",
                "rollouts": {
                    "task-001": {
                        "state": "running",
                        "pid": 1234,
                    }
                },
            }
        )
    )
    monkeypatch.setattr(
        resume_batch,
        "_pid_owns_rollout",
        lambda *_args, **_kwargs: False,
    )
    ec2 = _FakeEc2(
        [
            {
                "InstanceId": "i-a",
                "Tags": [
                    {"Key": "BatchId", "Value": "batch-a"},
                    {"Key": "RolloutId", "Value": "task-001"},
                ],
            }
        ]
    )

    plan = resume_batch._discover_plan([batch.parent], ec2=ec2)

    assert [entry["rollout_id"] for entry in plan["entries"]] == [
        "run-a/task-001"
    ]


def test_discovery_skips_rollout_owned_by_live_resume_process(
    tmp_path: Path, monkeypatch
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-a")
    (batch / "cloud-resume-state.json").write_text(
        json.dumps(
            {
                "entries": {
                    "run-a/task-001": {
                        "state": "running",
                        "pid": 5678,
                    }
                }
            }
        )
    )
    monkeypatch.setattr(resume_batch, "_pid_owns_rollout", lambda *_args, **_kwargs: True)

    try:
        resume_batch._discover_plan([batch.parent])
    except resume_batch.NoInterruptedRollouts:
        pass
    else:
        raise AssertionError("live recovery rollout was scheduled for duplicate resume")


def test_discovery_does_not_reopen_terminated_tagged_worker(tmp_path: Path) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-terminated")
    (batch / "batch-state.json").write_text(
        json.dumps({"batch_id": "batch-a"})
    )

    try:
        resume_batch._discover_plan([batch.parent], ec2=_FakeEc2([]))
    except resume_batch.NoInterruptedRollouts:
        pass
    else:
        raise AssertionError("terminated tagged worker was scheduled for resume")


def test_discovery_reaps_a_live_worker_only_after_evaluation_is_durable(
    tmp_path: Path,
) -> None:
    batch = tmp_path / "benchmarks" / "run-a"
    _interrupted_rollout(batch, "task-001", "i-a")
    (batch / "batch-state.json").write_text(
        json.dumps({"batch_id": "batch-a"})
    )
    rollout = batch / "rollouts" / "task-001"
    (rollout / "result.json").write_text(
        json.dumps(
            {
                "task_id": "task-001",
                "completed": True,
                "artifact_bytes": 100,
                "evaluation": {"status": "completed"},
            }
        )
    )
    ec2 = _FakeEc2(
        [
            {
                "InstanceId": "i-a",
                "Tags": [{"Key": "RolloutId", "Value": "task-001"}],
            }
        ]
    )

    plan = resume_batch._discover_plan([batch.parent], ec2=ec2)

    assert plan["entries"][0]["cleanup_only"] is True


def test_released_without_checkpoint_can_restart_fresh_on_same_worker(
    tmp_path: Path,
) -> None:
    log = tmp_path / "resume.log"
    log.write_text("BrokerError: checkpoint is required\n")
    plan = {"aws_profile": "test-profile", "aws_region": "us-east-1"}
    entry = {
        "task_id": "task-001",
        "instance_id": "i-a",
        "output_dir": str(tmp_path / "replacement"),
        "provider": "openai",
        "model_id": "gpt-test",
        "resume": False,
    }

    command = resume_batch._command(plan, entry)

    assert "--resume" not in command
    assert resume_batch._needs_fresh_same_worker_recovery(
        exit_code=1,
        result=None,
        log_path=log,
    )


def test_ready_session_before_first_action_restarts_fresh_on_same_worker(
    tmp_path: Path,
) -> None:
    log = tmp_path / "resume.log"
    log.write_text("resume session is not running: ready\n")

    assert resume_batch._needs_fresh_same_worker_recovery(
        exit_code=1,
        result=None,
        log_path=log,
    )


def test_openrouter_recovery_preserves_required_tool_choice(tmp_path: Path) -> None:
    plan = {"aws_profile": "test-profile", "aws_region": "us-east-1"}
    entry = {
        "task_id": "task-001",
        "instance_id": "i-a",
        "output_dir": str(tmp_path / "replacement"),
        "provider": "openrouter",
        "model_id": "moonshotai/kimi-k3",
        "resume": False,
    }

    command = resume_batch._command(plan, entry)

    index = command.index("--tool-choice-mode")
    assert command[index + 1] == "required"


def test_legacy_aws_plan_checks_and_cleans_worker_through_plugin(
    tmp_path: Path,
) -> None:
    class Ec2:
        def __init__(self) -> None:
            self.terminated: list[str] = []

        def describe_instances(self, **_kwargs):
            return {
                "Reservations": [
                    {
                        "Instances": [
                            {
                                "InstanceId": "i-a",
                                "State": {"Name": "running"},
                            }
                        ]
                    }
                ]
            }

        def terminate_instances(self, *, InstanceIds):
            self.terminated.extend(InstanceIds)
            return {}

    output_dir = tmp_path / "rollout"
    output_dir.mkdir()
    (output_dir / "result.json").write_text(
        json.dumps(
            {
                "completed": True,
                "artifact_bytes": 1,
                "evaluation": {"status": "completed"},
            }
        ),
        encoding="utf-8",
    )
    state_path = tmp_path / "resume-state.json"
    plan_path = tmp_path / "resume-plan.json"
    plan = {
        "schema_version": 1,
        "aws_profile": "test-profile",
        "aws_region": "us-east-1",
        "state_path": str(state_path),
        "max_concurrency": 1,
        "entries": [
            {
                "rollout_id": "run-a/task-001",
                "task_id": "task-001",
                "instance_id": "i-a",
                "output_dir": str(output_dir),
                "provider": "openai",
                "model_id": "gpt-test",
                "cleanup_only": True,
            }
        ],
    }
    ec2 = Ec2()

    exit_code = asyncio.run(
        resume_batch.execute_plan(
            plan,
            plan_path=plan_path,
            sessions={"aws": ec2},
        )
    )
    state = json.loads(state_path.read_text(encoding="utf-8"))

    assert exit_code == 0
    assert ec2.terminated == ["i-a"]
    assert state["entries"]["run-a/task-001"]["worker_id"] == "i-a"
    assert state["entries"]["run-a/task-001"]["instance_id"] == "i-a"