import json
import re
from types import SimpleNamespace

import pytest

from abi.errors import ToolError
from abi.exporters import SnakemakeExporter
from abi.exporters.snakemake import (
    _size_to_mb,
    _walltime_to_minutes,
    marker_dir_for,
)
from abi.plugins import get_plugin
from abi.schemas import PlanStep
from abi.tools import ToolRegistry


def test_snakemake_exporter_renders_rules_from_abi_plan(tmp_path):
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(
        overrides={
            "outdir": str(tmp_path / "results"),
            "log_dir": str(tmp_path / "log"),
        }
    )
    plan = plugin.build_plan(config)

    script = SnakemakeExporter().export(plan, config, plugin.registry())

    assert script.startswith("# Generated by ABI SnakemakeExporter.")
    assert "rule all:" in script
    assert "rule RNA1_QC_FASTP:" in script
    assert "rule RNA1_ALIGNMENT_STAR:" in script
    assert "rule RNA1_EXPRESSION_FEATURECOUNTS:" in script
    assert "fastp -i" in script
    assert "/examples/fixtures/tiny_R1.fastq" in script
    assert "STAR --runThreadN 4" in script
    assert 'featureCounts -T "$1" -p' in script
    assert "retrying featureCounts without -p" in script


def test_snakemake_exporter_rule_all_is_first_rule(tmp_path):
    """Snakemake runs the first rule by default, so ``rule all`` must lead."""
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)

    script = SnakemakeExporter().export(plan, config, plugin.registry())

    assert script.index("rule all:") < script.index("rule RNA1_QC_FASTP:")


def test_snakemake_exporter_wires_dependencies_via_markers(tmp_path):
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)

    script = SnakemakeExporter().export(plan, config, plugin.registry())

    star_rule = script.split("rule RNA1_ALIGNMENT_STAR:", 1)[1]
    assert "input:" in star_rule
    assert ".abi_snakemake/markers/RNA1_QC_FASTP.done" in star_rule
    assert 'touch("' in star_rule
    assert ".abi_snakemake/markers/RNA1_ALIGNMENT_STAR.done" in star_rule


def test_snakemake_exporter_emits_threads_resources_and_conda(tmp_path):
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)

    script = SnakemakeExporter().export(plan, config, plugin.registry())

    assert "threads: 4" in script
    assert "resources:" in script
    assert "mem_mb=" in script
    assert "conda:" in script
    assert "/envs/abi-qc.yml" in script


def test_snakemake_exporter_escapes_literal_braces_in_shell(tmp_path):
    step = PlanStep(
        "brace",
        "Brace",
        "fake_tool",
        "test",
        outputs={"output_dir": str(tmp_path / "out")},
    )
    plan = SimpleNamespace(analysis_type="test", project_name="test", threads=1, steps=[step])
    registry = ToolRegistry([])

    class _BraceSkill:
        def build_command(self, params):
            return ["sh", "-c", "echo ${UNSET_VAR:-default}"]

    registry.has = lambda tool_id: True  # type: ignore[method-assign]
    registry.create = lambda tool_id, mock_tools=True: _BraceSkill()  # type: ignore[method-assign]
    registry.get = lambda tool_id: {}  # type: ignore[method-assign]

    script = SnakemakeExporter().export(plan, {"outdir": str(tmp_path)}, registry)

    assert "${{UNSET_VAR:-default}}" in script


def test_snakemake_exporter_smoke_mode_does_not_call_real_tools(tmp_path):
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)

    script = SnakemakeExporter().export(plan, config, plugin.registry(), smoke=True)

    assert "# Export mode: smoke" in script
    assert "ABI Snakemake smoke step" in script
    assert "fastp -i" not in script
    assert "STAR --runThreadN" not in script
    assert "featureCounts -T" not in script
    assert "conda:" not in script


def test_snakemake_exporter_writes_snakefile(tmp_path):
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)
    output = tmp_path / "Snakefile"

    path = SnakemakeExporter().write(plan, config, plugin.registry(), output)

    assert path == output
    assert output.exists()
    assert output.read_text(encoding="utf-8").startswith("# Generated by ABI SnakemakeExporter.")


def test_snakemake_exporter_overwrites_existing_snakefile_after_successful_export(tmp_path):
    plugin = get_plugin("metatranscriptomics")
    config = plugin.load_config(overrides={"outdir": str(tmp_path / "results")})
    plan = plugin.build_plan(config)
    output = tmp_path / "Snakefile"
    output.write_text("old workflow\n", encoding="utf-8")

    SnakemakeExporter().write(plan, config, plugin.registry(), output)

    assert output.read_text(encoding="utf-8").startswith("# Generated by ABI SnakemakeExporter.")


def test_snakemake_exporter_preserves_existing_snakefile_when_export_fails(tmp_path):
    output = tmp_path / "Snakefile"
    output.write_text("old workflow\n", encoding="utf-8")
    bad_step = PlanStep("bad", "Bad", "missing", "test")
    plan = SimpleNamespace(analysis_type="test", project_name="test", threads=1, steps=[bad_step])

    with pytest.raises(ToolError, match="unknown tool"):
        SnakemakeExporter().write(plan, {}, ToolRegistry([]), output)

    assert output.read_text(encoding="utf-8") == "old workflow\n"


def test_snakemake_worker_internal_rule_payload_matches_step_runner_contract(tmp_path, monkeypatch):
    step = PlanStep(
        "normalize",
        "Normalize",
        "internal",
        "normalization",
        inputs={"source": "input.tsv"},
        outputs={"output_dir": str(tmp_path / "normalized")},
        params={
            "threshold": 2,
            "_internal_handler": {
                "handler_id": "normalize",
                "execution_scope": "worker",
            },
        },
    )
    plan = SimpleNamespace(
        analysis_type="example_plugin", project_name="test", threads=1, steps=[step]
    )
    config = {"outdir": str(tmp_path / "results"), "feature": {"enabled": True}}

    script = SnakemakeExporter().export(plan, config, ToolRegistry([]), project_root=tmp_path)

    assert "rule NORMALIZE:" in script
    assert '"plugin_id": "example_plugin"' in script
    assert '"step_name": "Normalize"' in script
    assert "chmod 600 .abi_payload.json" in script
    assert "abi run-step --payload-file .abi_payload.json" in script

    payload_match = re.search(
        r"<<'ABI_PAYLOAD'\n(?P<payload>.*?)\nABI_PAYLOAD", script, flags=re.DOTALL
    )
    assert payload_match is not None
    # Braces are doubled for Snakemake shell substitution; undo before parsing.
    payload = json.loads(payload_match.group("payload").replace("{{", "{").replace("}}", "}"))
    assert set(payload) == {
        "plugin_id",
        "step",
        "config",
        "provenance_dir",
        "result_path",
    }

    from abi import step_runner

    monkeypatch.chdir(tmp_path)
    monkeypatch.setattr(step_runner, "get_plugin", lambda plugin_id: object())
    monkeypatch.setattr(
        step_runner,
        "execute_step",
        lambda plugin, actual_step, actual_config, provenance_dir: step_runner.StepExecutionResult(
            step_id=actual_step.step_id,
            tool_id=actual_step.tool_id,
            status="success",
        ),
    )
    payload_path = tmp_path / ".abi_payload.json"
    payload_path.write_text(json.dumps(payload), encoding="utf-8")

    result = step_runner.execute_step_payload(payload_path)

    assert result.status == "success"
    assert (tmp_path / ".abi_result.json").exists()


def test_snakemake_driver_internal_step_fails_fast(tmp_path):
    step = PlanStep(
        "prepare",
        "Prepare",
        "internal",
        "preflight",
        params={
            "_internal_handler": {
                "handler_id": "prepare_inputs",
                "execution_scope": "driver",
            }
        },
    )
    plan = SimpleNamespace(
        analysis_type="example_plugin", project_name="test", threads=1, steps=[step]
    )

    with pytest.raises(ToolError, match=r"prepare.*prepare_inputs"):
        SnakemakeExporter().export(plan, {"outdir": str(tmp_path)}, ToolRegistry([]))


def test_snakemake_exporter_serializes_metagenomic_plasmid_steps(tmp_path):
    """Worker-scoped internal handlers generate Snakemake rules via abi run-step."""
    plugin = get_plugin("metagenomic_plasmid")
    config = plugin.load_config(
        "examples/config_minimal.yaml",
        overrides={"outdir": str(tmp_path / "plasmid"), "log_dir": str(tmp_path / "log")},
    )
    plan = plugin.build_plan(config)

    script = SnakemakeExporter().export(plan, config, plugin.registry())

    assert "rule " in script
    assert "abi run-step" in script or "abi-snakemake-smoke" in script


def test_marker_dir_for_resolves_relative_outdir_against_project_root(tmp_path):
    plan = SimpleNamespace(outdir="")
    marker_dir = marker_dir_for({"outdir": "results"}, plan, project_root=tmp_path)
    assert marker_dir == tmp_path / "results" / ".abi_snakemake" / "markers"


def test_size_to_mb_conversions():
    assert _size_to_mb("16GB") == 16384
    assert _size_to_mb("500MB") == 500
    assert _size_to_mb("1TB") == 1048576
    assert _size_to_mb("bogus") is None


def test_walltime_to_minutes_conversions():
    assert _walltime_to_minutes("04:00:00") == 240
    assert _walltime_to_minutes("01:30:00") == 90
    assert _walltime_to_minutes("45") == 45
    assert _walltime_to_minutes("bogus") is None
