"""Unit tests for the Snakemake runtime backend."""

from __future__ import annotations

import shutil
import subprocess
from pathlib import Path
from unittest.mock import patch

import pytest

from abi.dag import infer_dag
from abi.errors import ABIError
from abi.exporters import SnakemakeExporter
from abi.exporters.snakemake import marker_dir_for
from abi.plugins import get_plugin
from abi.runtimes import SnakemakeRuntime
from abi.runtimes.base import RuntimeOptions
from abi.runtimes.snakemake import (
    _command_rows,
    _cores_for,
    _status_from_marker,
    resolve_snakemake_bin,
)


def _fake_snakemake(tmp_path: Path, *, exit_code: int = 0) -> Path:
    fake_bin = tmp_path / "fake_snakemake"
    fake_bin.write_text(f"#!/bin/sh\nexit {exit_code}\n", encoding="utf-8")
    fake_bin.chmod(0o755)
    return fake_bin


def test_resolve_snakemake_bin_from_option(tmp_path: Path) -> None:
    fake_bin = _fake_snakemake(tmp_path)
    result = resolve_snakemake_bin(fake_bin)
    assert result == fake_bin.resolve()


def test_resolve_snakemake_bin_from_env_var(tmp_path: Path, monkeypatch) -> None:
    fake_bin = _fake_snakemake(tmp_path)
    monkeypatch.setenv("ABI_SNAKEMAKE_BIN", str(fake_bin))
    result = resolve_snakemake_bin(snakemake_bin=None)
    assert result == fake_bin.resolve()


def test_resolve_snakemake_bin_shutil_which_fallback(monkeypatch) -> None:
    monkeypatch.delenv("ABI_SNAKEMAKE_BIN", raising=False)
    monkeypatch.setattr(
        shutil,
        "which",
        lambda name: "/usr/local/bin/snakemake" if name == "snakemake" else None,
    )
    with (
        patch("pathlib.Path.exists", return_value=True),
        patch("os.access", return_value=True),
    ):
        result = resolve_snakemake_bin(snakemake_bin=None)
        assert result.name == "snakemake"


def test_resolve_snakemake_bin_no_binary_found_raises(monkeypatch) -> None:
    monkeypatch.setattr(shutil, "which", lambda name: None)
    monkeypatch.delenv("ABI_SNAKEMAKE_BIN", raising=False)
    with patch("pathlib.Path.exists", return_value=False):
        with pytest.raises(ABIError, match="Snakemake executable was not found"):
            resolve_snakemake_bin(snakemake_bin=None)


def test_status_from_marker_dry_run_returns_fallback(tmp_path: Path) -> None:
    assert _status_from_marker(tmp_path, "RULE_A", fallback="dry_run", dry_run=True) == "dry_run"


def test_status_from_marker_existing_marker_returns_success(tmp_path: Path) -> None:
    (tmp_path / "RULE_A.done").touch()
    assert _status_from_marker(tmp_path, "RULE_A", fallback="failed", dry_run=False) == "success"


def test_status_from_marker_missing_marker_returns_fallback(tmp_path: Path) -> None:
    assert _status_from_marker(tmp_path, "RULE_A", fallback="failed", dry_run=False) == "failed"


def test_cores_for_prefers_cpu_override() -> None:
    options = RuntimeOptions(engine="snakemake", cpu_override=7)
    plan = type("Plan", (), {"threads": 4})()
    assert _cores_for(plan, {"threads": 2}, options) == 7


def test_cores_for_falls_back_to_plan_threads() -> None:
    options = RuntimeOptions(engine="snakemake")
    plan = type("Plan", (), {"threads": 4})()
    assert _cores_for(plan, {}, options) == 4


def test_cores_for_defaults_to_one() -> None:
    options = RuntimeOptions(engine="snakemake")
    plan = type("Plan", (), {"threads": 0})()
    assert _cores_for(plan, {}, options) == 1


def test_snakemake_dry_run_writes_snakefile_and_provenance_without_executing(
    tmp_path, monkeypatch
) -> None:
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)
    runtime = SnakemakeRuntime(plugin, options=RuntimeOptions(engine="snakemake", smoke=True))

    def _forbidden(*args, **kwargs):  # pragma: no cover - must never run
        raise AssertionError("dry_run must not execute snakemake")

    monkeypatch.setattr(subprocess, "run", _forbidden)
    result = runtime.dry_run(plan, config)

    assert result.status == "dry_run"
    assert result.return_code == 0
    snakefile = result.outputs["snakefile"]
    assert snakefile.exists()
    assert snakefile.read_text(encoding="utf-8").startswith("# Generated by ABI SnakemakeExporter.")
    assert (tmp_path / "results" / "provenance" / "commands.tsv").exists()
    assert (tmp_path / "results" / "execution_plan.json").exists()
    # No snakemake logs: nothing was executed.
    assert not (tmp_path / "results" / "snakemake" / "snakemake.stdout.log").exists()


def test_snakemake_run_success_writes_logs_and_provenance(tmp_path) -> None:
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)
    fake_bin = _fake_snakemake(tmp_path, exit_code=0)
    runtime = SnakemakeRuntime(
        plugin,
        options=RuntimeOptions(engine="snakemake", smoke=True, snakemake_bin=fake_bin),
    )

    result = runtime.run(plan, config)

    assert result.status == "success"
    assert result.return_code == 0
    assert result.outputs["snakefile"].exists()
    assert result.outputs["snakemake_stdout"].exists()
    assert result.outputs["snakemake_stderr"].exists()
    commands_tsv = (tmp_path / "results" / "provenance" / "commands.tsv").read_text(
        encoding="utf-8"
    )
    assert "success" in commands_tsv


def test_snakemake_run_failure_writes_provenance_before_raising(tmp_path) -> None:
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)
    fake_bin = _fake_snakemake(tmp_path, exit_code=3)
    runtime = SnakemakeRuntime(
        plugin,
        options=RuntimeOptions(engine="snakemake", smoke=True, snakemake_bin=fake_bin),
    )

    with pytest.raises(ABIError, match="Snakemake run failed with exit code 3"):
        runtime.run(plan, config)

    # Provenance was written BEFORE the raise / 失败时也先写溯源再抛错
    commands_tsv = (tmp_path / "results" / "provenance" / "commands.tsv").read_text(
        encoding="utf-8"
    )
    assert "failed" in commands_tsv
    assert "Snakemake rule failed" in commands_tsv
    assert (tmp_path / "results" / "execution_plan.json").exists()
    assert (tmp_path / "results" / "snakemake" / "snakemake.stderr.log").exists()


def test_snakemake_run_marks_steps_with_done_markers_as_success(tmp_path) -> None:
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)
    dag = infer_dag(plan.steps, sequential_fallback=True)
    marker_dir = marker_dir_for(config, plan)
    marker_dir.mkdir(parents=True)
    (marker_dir / f"{dag.bindings[0].process_name}.done").touch()

    rows = _command_rows(
        plan,
        plugin.registry(),
        SnakemakeExporter(),
        dag=dag,
        return_code=1,
        smoke=True,
        marker_dir=marker_dir,
        dry_run=False,
    )

    assert rows[0]["status"] == "success"
    assert rows[0]["return_code"] == 0
    assert all(row["status"] == "failed" for row in rows[1:])


def test_snakemake_runtime_forwards_resource_options_to_exporter(tmp_path) -> None:
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)
    options = RuntimeOptions(
        engine="snakemake",
        smoke=True,
        cpu_override=7,
        memory_override="23GB",
        walltime_override="05:00:00",
    )
    runtime = SnakemakeRuntime(plugin, options=options)

    result = runtime.dry_run(plan, config)
    script = result.outputs["snakefile"].read_text(encoding="utf-8")

    assert "threads: 7" in script
    assert "mem_mb=23552" in script
    assert "runtime=300" in script


def test_snakemake_check_raises_without_binary(tmp_path, monkeypatch) -> None:
    plugin = get_plugin("metatranscriptomics")
    monkeypatch.setattr(shutil, "which", lambda name: None)
    monkeypatch.delenv("ABI_SNAKEMAKE_BIN", raising=False)
    runtime = SnakemakeRuntime(plugin, options=RuntimeOptions(engine="snakemake"))

    with patch("pathlib.Path.exists", return_value=False):
        with pytest.raises(ABIError, match="Snakemake executable was not found"):
            runtime.check()


@pytest.mark.requires_tools
def test_snakemake_runtime_end_to_end_smoke_with_real_snakemake(tmp_path) -> None:
    """End-to-end: real snakemake executes the smoke Snakefile (skipped in CI)."""
    if shutil.which("snakemake") is None:
        pytest.skip("snakemake binary not installed")
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)
    runtime = SnakemakeRuntime(plugin, options=RuntimeOptions(engine="snakemake", smoke=True))

    result = runtime.run(plan, config)

    assert result.status == "success"
    marker_dir = marker_dir_for(config, plan)
    assert any(marker_dir.glob("*.done"))
