import json
import stat
from pathlib import Path

import pytest

from scripts import create_root_files

parse_detector_spec = create_root_files.parse_detector_spec
next_shard_index = create_root_files.next_shard_index
plan_jobs = create_root_files.plan_jobs
job_seed = create_root_files.job_seed
run_job = create_root_files.run_job
run_all = create_root_files.run_all
SimJob = create_root_files.SimJob
PlanError = create_root_files.PlanError


def _write_fake_executable(
    path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0
) -> Path:
    """Stand-in for run_pbwo4/run_sampling: writes *output_count* .root files
    into its own cwd (so callers can verify each job gets an isolated workdir
    and that the workdir ends up holding *only* the .root output, matching
    real executables that may also drop other side-effect files). Each .root
    file's content is a JSON record of argv/cwd/timing, so a test can recover
    that info from the moved file after the workdir is gone. Also writes a
    non-.root side file unconditionally, to catch any cleanup code that
    assumes the workdir is empty after the .root is moved out.
    """
    path.write_text(
        f"""#!/usr/bin/env python3
import json, os, sys, time

start = time.time()
time.sleep({sleep})
end = time.time()
payload = json.dumps(
    {{
        "argv": sys.argv[1:],
        "cwd": os.getcwd(),
        "start": start,
        "end": end,
        "seed": os.environ.get("MINICALOSIM_SEED"),
    }}
)
for i in range({output_count}):
    with open(f"out_{{i}}.root", "w") as f:
        f.write(payload)
with open("side_effect.log", "w") as f:
    f.write("not a root file")
sys.exit({exit_code})
"""
    )
    path.chmod(path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
    return path


def test_parse_detector_spec_with_config():
    assert parse_detector_spec("sampling_pb_scint:pb_scint") == (
        "sampling_pb_scint",
        "pb_scint",
    )


def test_parse_detector_spec_without_config():
    assert parse_detector_spec("pbwo4") == ("pbwo4", None)


def test_parse_detector_spec_rejects_empty_parts():
    with pytest.raises(PlanError):
        parse_detector_spec(":pb_scint")
    with pytest.raises(PlanError):
        parse_detector_spec("sampling_pb_scint:")


def test_next_shard_index_missing_dir(tmp_path):
    assert next_shard_index(tmp_path / "nope") == 0


def test_next_shard_index_continues_past_existing(tmp_path):
    d = tmp_path / "pbwo4"
    d.mkdir()
    (d / "shard-000.root").touch()
    (d / "shard-005.root").touch()
    (d / "not-a-shard.root").touch()
    assert next_shard_index(d) == 6


def test_plan_jobs_rejects_missing_gen(tmp_path):
    with pytest.raises(PlanError):
        plan_jobs(
            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1"
        )


def test_plan_jobs_rejects_malformed_gen(tmp_path):
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    with pytest.raises(PlanError):
        plan_jobs(
            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen"
        )


def test_plan_jobs_continues_from_existing_shards(tmp_path):
    gen_dir = tmp_path / "raw" / "steps" / "gen1"
    (gen_dir / "pbwo4").mkdir(parents=True)
    (gen_dir / "pbwo4" / "shard-000.root").touch()
    (gen_dir / "pbwo4" / "shard-001.root").touch()

    jobs = plan_jobs(
        ["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1"
    )

    assert [j.shard_index for j in jobs] == [2, 3, 4]
    assert all(j.detector == "pbwo4" and j.config is None for j in jobs)


def test_plan_jobs_multiple_detectors_each_start_independently(tmp_path):
    gen_dir = tmp_path / "raw" / "steps" / "gen1"
    (gen_dir / "sampling_pb_scint").mkdir(parents=True)
    (gen_dir / "sampling_pb_scint" / "shard-003.root").touch()
    (gen_dir / "sampling_fe_scint").mkdir(parents=True)

    jobs = plan_jobs(
        ["sampling_pb_scint:pb_scint", "sampling_fe_scint:fe_scint"],
        num_files=2,
        dataset_root=tmp_path,
        kind="steps",
        gen="gen1",
    )

    by_detector = {}
    for j in jobs:
        by_detector.setdefault(j.detector, []).append(j.shard_index)
    assert by_detector["sampling_pb_scint"] == [4, 5]
    assert by_detector["sampling_fe_scint"] == [0, 1]


def test_job_seed_deterministic():
    job = SimJob(detector="pbwo4", config=None, shard_index=3)
    assert job_seed("steps", "gen1", job, None) == job_seed("steps", "gen1", job, None)


def test_job_seed_varies_by_shard_index():
    a = SimJob(detector="pbwo4", config=None, shard_index=0)
    b = SimJob(detector="pbwo4", config=None, shard_index=1)
    assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None)


def test_job_seed_varies_by_detector():
    a = SimJob(detector="pbwo4", config=None, shard_index=0)
    b = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
    assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None)


def test_job_seed_varies_by_gen():
    job = SimJob(detector="pbwo4", config=None, shard_index=0)
    assert job_seed("steps", "gen1", job, None) != job_seed("steps", "gen2", job, None)


def test_job_seed_varies_by_energy():
    job = SimJob(detector="pbwo4", config=None, shard_index=0)
    assert job_seed("steps", "gen1", job, 1.0) != job_seed("steps", "gen1", job, 10.0)


def test_run_job_passes_deterministic_seed_env_var(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py")
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    job = SimJob(detector="pbwo4", config=None, shard_index=5)
    result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)

    assert result.dest is not None
    payload = json.loads(result.dest.read_text())
    assert payload["seed"] == str(job_seed("steps", "gen1", job, None))


def test_run_job_moves_output_to_correct_shard_path(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py")
    gen_dir = tmp_path / "raw" / "steps" / "gen1"
    gen_dir.mkdir(parents=True)
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    job = SimJob(detector="pbwo4", config=None, shard_index=7)
    result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)

    assert result.ok
    assert result.dest == gen_dir / "pbwo4" / "shard-007.root"
    assert result.dest is not None
    assert result.dest.is_file()
    assert not any(tmp_root.iterdir())  # workdir cleaned up


def test_run_job_passes_config_arg_and_isolates_cwd(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py")
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    job = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
    result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)

    assert result.ok
    assert result.dest is not None
    payload = json.loads(result.dest.read_text())
    assert payload["argv"] == ["pb_scint", "10000"]
    # ran in its own scratch workdir under .sim-tmp, not directly in dataset_root
    assert payload["cwd"] != str(tmp_path)
    assert str(tmp_root) in payload["cwd"]


def test_run_job_omits_config_arg_when_none(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py")
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    job = SimJob(detector="pbwo4", config=None, shard_index=0)
    result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)

    assert result.dest is not None
    payload = json.loads(result.dest.read_text())
    assert payload["argv"] == ["10000"]


def test_run_job_appends_energy_arg_when_given(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py")
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    job = SimJob(detector="pbwo4_10gev", config=None, shard_index=0)
    result = run_job(job, fake, 10000, 10.0, tmp_path, "steps", "gen1", tmp_root)

    assert result.dest is not None
    payload = json.loads(result.dest.read_text())
    assert payload["argv"] == ["10000", "10.0"]


def test_run_job_fails_when_executable_errors(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py", exit_code=1)
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    job = SimJob(detector="pbwo4", config=None, shard_index=0)
    result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)

    assert not result.ok
    assert "exited 1" in result.message


def test_run_job_fails_when_no_root_file_produced(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py", output_count=0)
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    job = SimJob(detector="pbwo4", config=None, shard_index=0)
    result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)

    assert not result.ok
    assert "found 0" in result.message


def test_run_job_fails_when_multiple_root_files_produced(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py", output_count=2)
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    job = SimJob(detector="pbwo4", config=None, shard_index=0)
    result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)

    assert not result.ok
    assert "found 2" in result.message


def test_run_job_refuses_to_overwrite_existing_shard(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py")
    detector_dir = tmp_path / "raw" / "steps" / "gen1" / "pbwo4"
    detector_dir.mkdir(parents=True)
    (detector_dir / "shard-000.root").write_text("already here")
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    job = SimJob(detector="pbwo4", config=None, shard_index=0)
    result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)

    assert not result.ok
    assert "overwrite" in result.message
    assert (detector_dir / "shard-000.root").read_text() == "already here"


def test_run_all_caps_concurrency(tmp_path):
    fake = _write_fake_executable(tmp_path / "fake_exe.py", sleep=0.2)
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    tmp_root = tmp_path / ".sim-tmp"
    tmp_root.mkdir()

    jobs = [SimJob(detector="pbwo4", config=None, shard_index=i) for i in range(6)]
    results = run_all(
        jobs,
        fake,
        10000,
        None,
        tmp_path,
        "steps",
        "gen1",
        max_workers=2,
        tmp_root=tmp_root,
    )

    assert all(r.ok and r.dest is not None for r in results)
    dests = [r.dest for r in results if r.dest is not None]
    assert {d.name for d in dests} == {f"shard-{i:03d}.root" for i in range(6)}

    intervals = [json.loads(d.read_text()) for d in dests]
    events = sorted(
        [(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals]
    )
    concurrent = 0
    peak = 0
    for _, delta in events:
        concurrent += delta
        peak = max(peak, concurrent)
    assert peak <= 2