"""Generate new ROOT shards by running a minicalosim executable (e.g.
run_pbwo4, run_sampling) and filing the output into the dataset's raw/ tree:

    raw/<kind>/<gen>/<detector>/shard-NNN.root

These executables take `[configName] nEvents` and always write a fixed-name
*.root file into the current directory — so running several in parallel
needs separate working directories, and the output filename has to be
discovered rather than assumed (it differs per executable: run_pbwo4 writes
pbwo4_<n>events_hits.root, run_sampling writes sampling_<config>_<n>events_hits.root,
others may differ again). This script gives each run its own scratch
directory under <dataset-root>/.sim-tmp/, requires exactly one *.root to
appear there, and moves it to the next free shard index for that detector
(existing shards are never overwritten).

--gen must already exist under raw/<kind>/ — create one first with
`dwarf bump-gen`.

See `uv run dwarf make-root --help` for the CLI.
"""

import os
import re
import shutil
import subprocess
import sys
import uuid
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

# Must match scripts/bump_dataset_version.py's GEN_RE.
GEN_RE = re.compile(r"^gen\d+$")
SHARD_RE = re.compile(r"^shard-(\d+)\.root$")


class PlanError(ValueError):
    pass


@dataclass(frozen=True)
class SimJob:
    detector: str
    config: str | None
    shard_index: int


@dataclass(frozen=True)
class JobResult:
    job: SimJob
    ok: bool
    dest: Path | None
    message: str
    stdout: str
    stderr: str


def parse_detector_spec(spec: str) -> tuple[str, str | None]:
    """'sampling_pb_scint:pb_scint' -> ('sampling_pb_scint', 'pb_scint');
    'pbwo4' -> ('pbwo4', None)."""
    if ":" in spec:
        label, config = spec.split(":", 1)
        if not label or not config:
            raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG")
        return label, config
    return spec, None


def next_shard_index(detector_dir: Path) -> int:
    """One past the highest existing shard-NNN.root in *detector_dir* (0 if none/missing)."""
    if not detector_dir.is_dir():
        return 0
    best = -1
    for child in detector_dir.iterdir():
        m = SHARD_RE.match(child.name)
        if m and child.is_file():
            best = max(best, int(m.group(1)))
    return best + 1


def plan_jobs(
    detector_specs: list[str],
    num_files: int,
    dataset_root: Path,
    kind: str,
    gen: str,
) -> list[SimJob]:
    if not GEN_RE.match(gen):
        raise PlanError(f"--gen must look like 'genN', got {gen!r}")
    gen_dir = dataset_root / "raw" / kind / gen
    if not gen_dir.is_dir():
        raise PlanError(
            f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first"
        )

    jobs = []
    for spec in detector_specs:
        label, config = parse_detector_spec(spec)
        start = next_shard_index(gen_dir / label)
        for i in range(num_files):
            jobs.append(SimJob(detector=label, config=config, shard_index=start + i))
    return jobs


def run_job(
    job: SimJob,
    executable: Path,
    events_per_file: int,
    dataset_root: Path,
    kind: str,
    gen: str,
    tmp_root: Path,
) -> JobResult:
    workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
    workdir.mkdir(parents=True)

    cmd = [str(executable)]
    if job.config:
        cmd.append(job.config)
    cmd.append(str(events_per_file))

    result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True)

    if result.returncode != 0:
        return JobResult(
            job, False, None,
            f"executable exited {result.returncode}",
            result.stdout, result.stderr,
        )

    produced = sorted(workdir.glob("*.root"))
    if len(produced) != 1:
        return JobResult(
            job, False, None,
            f"expected exactly one .root output in {workdir}, found {len(produced)}: "
            f"{[p.name for p in produced]}",
            result.stdout, result.stderr,
        )

    dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
    if dest.exists():
        return JobResult(
            job, False, None, f"refusing to overwrite existing {dest}",
            result.stdout, result.stderr,
        )

    dest.parent.mkdir(parents=True, exist_ok=True)
    shutil.move(str(produced[0]), str(dest))
    # rmtree, not rmdir: the executable may leave other side-effect files
    # (logs, seed state, ...) behind in its scratch workdir besides the .root.
    shutil.rmtree(workdir)
    return JobResult(job, True, dest, "ok", result.stdout, result.stderr)


def run_all(
    jobs: list[SimJob],
    executable: Path,
    events_per_file: int,
    dataset_root: Path,
    kind: str,
    gen: str,
    max_workers: int,
    tmp_root: Path,
) -> list[JobResult]:
    results = []
    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(
                run_job, job, executable, events_per_file, dataset_root, kind, gen, tmp_root
            ): job
            for job in jobs
        }
        for future in as_completed(futures):
            result = future.result()
            label = f"{result.job.detector} shard-{result.job.shard_index:03d}"
            status = "ok" if result.ok else f"FAILED: {result.message}"
            print(f"\n=== {label}: {status} ===")
            if result.stdout:
                print(result.stdout, end="")
            if result.stderr:
                print(result.stderr, end="", file=sys.stderr)
            results.append(result)
    return results


def run_make_root(
    executable: Path,
    detector: list[str],
    num_files: int,
    events_per_file: int,
    kind: str,
    gen: str,
    dataset_root: str,
    jobs: int,
    execute: bool,
) -> None:
    if jobs < 1:
        raise SystemExit("error: --jobs must be >= 1")
    if num_files < 1:
        raise SystemExit("error: --num-files must be >= 1")
    if events_per_file < 1:
        raise SystemExit("error: --events-per-file must be >= 1")
    if not executable.is_file() or not os.access(executable, os.X_OK):
        raise SystemExit(f"error: {executable} is not an executable file")

    dataset_root_path = Path(dataset_root)
    try:
        planned_jobs = plan_jobs(detector, num_files, dataset_root_path, kind, gen)
    except PlanError as exc:
        raise SystemExit(f"error: {exc}")

    print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
    print(f"executable: {executable}")
    for job in planned_jobs:
        cmd = [str(executable)] + ([job.config] if job.config else []) + [str(events_per_file)]
        dest = dataset_root_path / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
        print(f"  {' '.join(cmd)}  ->  {dest}")

    if not execute:
        print("\nDry run only — pass --execute to apply.")
        return

    tmp_root = dataset_root_path / ".sim-tmp"
    tmp_root.mkdir(parents=True, exist_ok=True)
    results = run_all(
        planned_jobs, executable, events_per_file, dataset_root_path, kind, gen,
        max_workers=jobs, tmp_root=tmp_root,
    )
    if tmp_root.is_dir() and not any(tmp_root.iterdir()):
        tmp_root.rmdir()

    failures = [r for r in results if not r.ok]
    if failures:
        print(f"\n{len(failures)} of {len(results)} job(s) failed:", file=sys.stderr)
        for r in failures:
            print(f"  {r.job.detector} shard-{r.job.shard_index:03d}: {r.message}", file=sys.stderr)
        raise SystemExit(1)

    print(f"\nAll {len(results)} job(s) completed.")
