#!/usr/bin/env python3
"""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
bump_dataset_version.py bump-gen.

Usage:
    create_root_files.py --executable build/run_pbwo4 --detector pbwo4 \\
        --gen gen1 --num-files 4 --events-per-file 10000 --execute

    create_root_files.py --executable build/run_sampling \\
        --detector sampling_pb_scint:pb_scint \\
        --detector sampling_fe_scint:fe_scint \\
        --gen gen1 --num-files 4 --events-per-file 10000 --jobs 8 --execute
"""

import argparse
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 build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--executable", required=True, type=Path, help="Built minicalosim run_* executable"
    )
    parser.add_argument(
        "--detector",
        action="append",
        required=True,
        metavar="NAME[:CONFIG]",
        help="Dataset detector label, optionally with ':CONFIG' to pass as the "
        "executable's config-name argument (e.g. sampling_pb_scint:pb_scint). "
        "Omit ':CONFIG' for executables that take no config selector (e.g. run_pbwo4). "
        "Repeatable.",
    )
    parser.add_argument(
        "--num-files", type=int, required=True, help="New shards to create per detector"
    )
    parser.add_argument(
        "--events-per-file", type=int, required=True, help="nEvents passed to the executable"
    )
    parser.add_argument("--kind", default="steps", help="steps | hits | ... (default: steps)")
    parser.add_argument("--gen", required=True, help="Existing gen tag under raw/<kind>/, e.g. gen1")
    parser.add_argument(
        "--dataset-root",
        default="/ceph/lbogner/geant_steps",
        help="Dataset root (default: /ceph/lbogner/geant_steps)",
    )
    parser.add_argument(
        "-j", "--jobs", type=int, default=4, help="Parallel simulation runs (default: 4)"
    )
    parser.add_argument(
        "--execute", action="store_true", help="Actually run jobs (default: dry run / print plan)"
    )
    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.jobs < 1:
        parser.error("--jobs must be >= 1")
    if args.num_files < 1:
        parser.error("--num-files must be >= 1")
    if args.events_per_file < 1:
        parser.error("--events-per-file must be >= 1")
    if not args.executable.is_file() or not os.access(args.executable, os.X_OK):
        parser.error(f"{args.executable} is not an executable file")

    dataset_root = Path(args.dataset_root)
    try:
        jobs = plan_jobs(args.detector, args.num_files, dataset_root, args.kind, args.gen)
    except PlanError as exc:
        parser.error(str(exc))

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

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

    tmp_root = dataset_root / ".sim-tmp"
    tmp_root.mkdir(parents=True, exist_ok=True)
    results = run_all(
        jobs, args.executable, args.events_per_file, dataset_root, args.kind, args.gen,
        max_workers=args.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)
        sys.exit(1)

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


if __name__ == "__main__":
    main()