"""Cut a new raw generation or processed schema version for the geant_steps
dataset tree (see scripts/migrate_geant_steps.py for the layout):

    raw/<kind>/<gen>/<detector>/shard-NNN.root
    processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet
    pools/<detector>/<pool>.manifest

`gen` bumps when the underlying ROOT changes (geometry/physics-list/macro).
`schema` bumps when the parquet export (`dwarf convert` or similar) changes,
and is scoped to its gen — a new gen always starts back at schema1.

Creates the new (empty) target directory and appends a dated, reasoned entry
to VERSIONS.md. Defaults to a dry run; pass execute=True to apply.

See `uv run dwarf bump-gen/bump-schema/update-manifest/create-manifest/status
--help` for the CLI.
"""

import datetime as dt
import os
import re
import subprocess
import sys
from pathlib import Path

GEN_RE = re.compile(r"^gen(\d+)$")
SCHEMA_RE = re.compile(r"^schema(\d+)$")

# One color per tree level in `dwarf status` output, so the eye can jump
# straight to e.g. "all the schema rows" or "all the totals".
_LEVEL_COLORS = {
    "kind": "\033[1;36m",  # bold cyan — top-level kind/ header
    "gen": "\033[1;33m",  # bold yellow — genN row + kind total
    "bucket": "\033[34m",  # blue — raw/processed subtotals
    "schema": "\033[32m",  # green — schemaN rows
    "root": "\033[1;35m",  # bold magenta — derived/, pools/, grand total
}
_RESET = "\033[0m"


def _use_color() -> bool:
    return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None


def _colorize(text: str, level: str) -> str:
    if not _use_color():
        return text
    return f"{_LEVEL_COLORS[level]}{text}{_RESET}"


def _find_component_idx(abs_path: Path, pattern: re.Pattern) -> int | None:
    """Return the index of the first path component matching *pattern*, or None."""
    for i, part in enumerate(abs_path.parts):
        if pattern.match(part):
            return i
    return None


def _max_index(parent: Path, pattern: re.Pattern) -> int:
    """Highest N across child dir names matching *pattern* (0 if none/missing)."""
    if not parent.is_dir():
        return 0
    best = 0
    for child in parent.iterdir():
        m = pattern.match(child.name)
        if m and child.is_dir():
            best = max(best, int(m.group(1)))
    return best


def _git_user_name() -> str | None:
    try:
        out = subprocess.run(
            ["git", "config", "user.name"], capture_output=True, text=True, timeout=2
        )
    except OSError:
        return None
    name = out.stdout.strip()
    return name or None


def plan_bump_gen(
    root: Path, kind: str, reason: str, by: str | None, date: str,
    target: str | None = None,
) -> tuple[list[Path], str]:
    """New gen tag is one past the highest seen under raw/ or processed/ for *kind*,
    or *target* if explicitly provided."""
    if target is not None:
        if not GEN_RE.match(target):
            raise SystemExit(f"error: --to must look like 'genN', got {target!r}")
        gen_tag = target
    else:
        next_gen = (
            max(
                _max_index(root / "raw" / kind, GEN_RE),
                _max_index(root / "processed" / kind, GEN_RE),
            )
            + 1
        )
        gen_tag = f"gen{next_gen}"
    new_dirs = [
        root / "raw" / kind / gen_tag,
        root / "processed" / kind / gen_tag / "schema1",
    ]
    by_suffix = f" ({by})" if by else ""
    log_line = f"- `{gen_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
    return new_dirs, log_line


def plan_bump_schema(
    root: Path, kind: str, gen_tag: str, reason: str, by: str | None, date: str,
    target: str | None = None,
) -> tuple[list[Path], str]:
    if not GEN_RE.match(gen_tag):
        raise SystemExit(f"error: --gen must look like 'genN', got {gen_tag!r}")
    raw_gen_dir = root / "raw" / kind / gen_tag
    processed_gen_dir = root / "processed" / kind / gen_tag
    if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir():
        raise SystemExit(
            f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
        )
    if target is not None:
        if not SCHEMA_RE.match(target):
            raise SystemExit(f"error: --to must look like 'schemaN', got {target!r}")
        schema_tag = target
    else:
        next_schema = _max_index(processed_gen_dir, SCHEMA_RE) + 1
        schema_tag = f"schema{next_schema}"
    new_dirs = [processed_gen_dir / schema_tag]
    by_suffix = f" ({by})" if by else ""
    log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
    return new_dirs, log_line


def apply_bump(root: Path, new_dirs: list[Path], log_line: str) -> None:
    for d in new_dirs:
        d.mkdir(parents=True, exist_ok=True)
    versions_path = root / "VERSIONS.md"
    if not versions_path.exists():
        versions_path.write_text("# Dataset versions\n\n")
    with versions_path.open("a") as f:
        f.write(log_line + "\n")


def _du(path: Path) -> int:
    """Total size in bytes of all regular files under *path* (0 if missing)."""
    if not path.is_dir():
        return 0
    total = 0
    for dirpath, _dirnames, filenames in os.walk(path):
        for name in filenames:
            fp = Path(dirpath) / name
            try:
                total += fp.stat().st_size
            except OSError:
                pass
    return total


def _count_files(path: Path) -> int:
    """Total number of regular files under *path*, recursively (0 if missing)."""
    if not path.is_dir():
        return 0
    total = 0
    for _dirpath, _dirnames, filenames in os.walk(path):
        total += len(filenames)
    return total


# Must match giant.data.loader.MANIFEST_SUFFIX.
MANIFEST_SUFFIX = ".manifest"


def _manifest_referenced_files(pools_root: Path) -> set[Path]:
    """Resolved absolute paths of every file listed in any *.manifest under *pools_root*."""
    referenced: set[Path] = set()
    if not pools_root.is_dir():
        return referenced
    for manifest_path in pools_root.rglob(f"*{MANIFEST_SUFFIX}"):
        try:
            referenced.update(_resolve_manifest_files(manifest_path))
        except OSError:
            continue
    return referenced


def _referenced_root_count(raw_gen_dir: Path, processed_gen_dir: Path) -> tuple[int, int]:
    """(total .root files, count with a same-named .parquet under any schema) for one gen."""
    if not raw_gen_dir.is_dir():
        return 0, 0
    parquet_stems: set[tuple[str, str]] = set()
    if processed_gen_dir.is_dir():
        for schema_dir in processed_gen_dir.iterdir():
            if not schema_dir.is_dir():
                continue
            for detector_dir in schema_dir.iterdir():
                if not detector_dir.is_dir():
                    continue
                for f in detector_dir.iterdir():
                    if f.is_file() and f.suffix == ".parquet":
                        parquet_stems.add((detector_dir.name, f.stem))

    total = 0
    referenced = 0
    for detector_dir in raw_gen_dir.iterdir():
        if not detector_dir.is_dir():
            continue
        for f in detector_dir.iterdir():
            if f.is_file() and f.suffix == ".root":
                total += 1
                if (detector_dir.name, f.stem) in parquet_stems:
                    referenced += 1
    return total, referenced


def _referenced_parquet_count(
    schema_dir: Path, manifest_referenced: set[Path]
) -> tuple[int, int]:
    """(total .parquet files, count listed in at least one manifest) for one schema dir."""
    if not schema_dir.is_dir():
        return 0, 0
    total = 0
    referenced = 0
    for detector_dir in schema_dir.iterdir():
        if not detector_dir.is_dir():
            continue
        for f in detector_dir.iterdir():
            if f.is_file() and f.suffix == ".parquet":
                total += 1
                if f.resolve() in manifest_referenced:
                    referenced += 1
    return total, referenced


def _human_size(n: int) -> str:
    size = float(n)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if size < 1024 or unit == "TB":
            return f"{size:.1f} {unit}" if unit != "B" else f"{int(size)} {unit}"
        size /= 1024
    return f"{size:.1f} TB"


_LABEL_WIDTH = 26
_COUNT_WIDTH = 20
_ROW_WIDTH = _LABEL_WIDTH + _COUNT_WIDTH + 10


def _count_str(count: int, referenced: int | None = None) -> str:
    files = "file" if count == 1 else "files"
    if referenced is not None:
        return f"{count} {files} ({referenced} ref)"
    return f"{count} {files}"


def _row(
    label: str,
    size_bytes: int,
    indent: int = 0,
    level: str | None = None,
    count: int | None = None,
    referenced: int | None = None,
) -> str:
    text = "  " * indent + label
    count_str = _count_str(count, referenced) if count is not None else ""
    size_str = _human_size(size_bytes)
    row = f"{text:<{_LABEL_WIDTH}}{count_str:<{_COUNT_WIDTH}}{size_str:>10}"
    return _colorize(row, level) if level else row


def print_status(root: Path) -> None:
    raw_root = root / "raw"
    if not raw_root.is_dir():
        print(f"no raw/ tree found under {root}")
        return
    manifest_referenced = _manifest_referenced_files(root / "pools")
    grand_total = 0
    grand_files = 0
    for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
        kind = kind_dir.name
        gens = sorted(
            int(m.group(1))
            for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
            if m
        )
        print(_colorize(f"{kind}/", "kind"))
        kind_total = 0
        kind_files = 0
        for gen in gens:
            gen_tag = f"gen{gen}"
            raw_gen_dir = raw_root / kind / gen_tag
            raw_size = _du(raw_gen_dir)
            processed_gen_dir = root / "processed" / kind / gen_tag
            schema_dir = processed_gen_dir
            schemas = sorted(
                int(m.group(1))
                for m in (
                    SCHEMA_RE.match(p.name)
                    for p in (schema_dir.iterdir() if schema_dir.is_dir() else [])
                    if p.is_dir()
                )
                if m
            )
            schema_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas}
            schema_counts = {
                s: _referenced_parquet_count(schema_dir / f"schema{s}", manifest_referenced)
                for s in schemas
            }
            processed_size = sum(schema_sizes.values())
            processed_files = sum(c[0] for c in schema_counts.values())
            processed_referenced = sum(c[1] for c in schema_counts.values())
            raw_files, raw_referenced = _referenced_root_count(raw_gen_dir, processed_gen_dir)
            gen_total = raw_size + processed_size
            gen_files = raw_files + processed_files
            kind_total += gen_total
            kind_files += gen_files

            print(_row(gen_tag, gen_total, indent=1, level="gen", count=gen_files))
            print(
                _row(
                    "raw", raw_size, indent=2, level="bucket",
                    count=raw_files, referenced=raw_referenced,
                )
            )
            print(
                _row(
                    "processed", processed_size, indent=2, level="bucket",
                    count=processed_files, referenced=processed_referenced,
                )
            )
            if schemas:
                for s in schemas:
                    s_total, s_referenced = schema_counts[s]
                    print(
                        _row(
                            f"schema{s}", schema_sizes[s], indent=3, level="schema",
                            count=s_total, referenced=s_referenced,
                        )
                    )
            else:
                print(_colorize("      (none)", "schema"))
        print(_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files))
        print()
        grand_total += kind_total
        grand_files += kind_files

    derived_dir = root / "derived"
    pools_dir = root / "pools"
    derived_size = _du(derived_dir)
    pools_size = _du(pools_dir)
    derived_files = _count_files(derived_dir)
    pools_files = _count_files(pools_dir)
    grand_total += derived_size + pools_size
    grand_files += derived_files + pools_files
    print(_row("derived/", derived_size, level="root", count=derived_files))
    print(_row("pools/", pools_size, level="root", count=pools_files))
    print("-" * _ROW_WIDTH)
    print(_row("grand total", grand_total, level="root", count=grand_files))


# ---------------------------------------------------------------------------
# update-manifest
# ---------------------------------------------------------------------------

def plan_update_manifest(
    manifest_path: Path,
    target_schema: str | None,
    target_gen: str | None = None,
) -> tuple[list[tuple[str, str | None]], list[Path]]:
    """Parse a manifest and plan gen/schema replacements for each data line.

    Returns:
        lines: list of (original_line, new_relative_path_or_None)
            None means the line is unchanged (comment, blank, or already at target).
        missing: resolved absolute paths that don't exist on disk.
    """
    manifest_path = manifest_path.resolve()
    if not manifest_path.exists():
        raise SystemExit(f"error: manifest not found: {manifest_path}")
    manifest_dir = manifest_path.parent
    raw_lines = manifest_path.read_text().splitlines()

    result: list[tuple[str, str | None]] = []
    missing: list[Path] = []
    schema_cache: dict[Path, str] = {}

    for raw in raw_lines:
        stripped = raw.strip()
        if not stripped or stripped.startswith("#"):
            result.append((raw, None))
            continue

        old_abs = (manifest_dir / stripped).resolve()
        parts = list(old_abs.parts)
        changed = False

        # Apply gen replacement first (shifts subsequent indices).
        if target_gen is not None:
            gen_idx = _find_component_idx(Path(*parts), GEN_RE)
            if gen_idx is not None and parts[gen_idx] != target_gen:
                parts[gen_idx] = target_gen
                changed = True

        schema_idx = _find_component_idx(Path(*parts), SCHEMA_RE)
        if schema_idx is not None:
            old_schema = parts[schema_idx]
            gen_dir = Path(*parts[:schema_idx])

            if target_schema is not None:
                new_schema = target_schema
            else:
                if gen_dir not in schema_cache:
                    n = _max_index(gen_dir, SCHEMA_RE)
                    if n == 0:
                        raise SystemExit(f"error: no schema dirs found under {gen_dir}")
                    schema_cache[gen_dir] = f"schema{n}"
                new_schema = schema_cache[gen_dir]

            if new_schema != old_schema:
                parts[schema_idx] = new_schema
                changed = True

        if not changed:
            result.append((raw, None))
            continue

        new_abs = Path(*parts)
        if not new_abs.exists():
            missing.append(new_abs)

        new_rel = os.path.relpath(new_abs, start=manifest_dir)
        result.append((raw, new_rel))

    return result, missing


def apply_update_manifest(manifest_path: Path, lines: list[tuple[str, str | None]]) -> None:
    out = [replacement if replacement is not None else original for original, replacement in lines]
    manifest_path.write_text("\n".join(out) + "\n")


# ---------------------------------------------------------------------------
# create-manifest
# ---------------------------------------------------------------------------

def _resolve_manifest_files(manifest_path: Path) -> list[Path]:
    """Read a manifest and return its entries as resolved absolute paths."""
    files = []
    for line in manifest_path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        files.append((manifest_path.parent / line).resolve())
    return files


def plan_create_manifest(
    output_path: Path, parquet_files: list[Path]
) -> tuple[list[str], list[Path], list[Path]]:
    """Return (relative_lines, missing_files, resolved_abs_paths)."""
    manifest_dir = output_path.resolve().parent
    lines: list[str] = []
    missing: list[Path] = []
    resolved: list[Path] = []
    for p in parquet_files:
        abs_p = p.resolve()
        resolved.append(abs_p)
        if not abs_p.exists():
            missing.append(abs_p)
        lines.append(os.path.relpath(abs_p, start=manifest_dir))
    return lines, missing, resolved


def check_holdout_overlap(
    output_path: Path, resolved_new_files: list[Path]
) -> list[tuple[str, Path]]:
    """Return (other_manifest_name, file) pairs where new files clash with existing manifests.

    The check is triggered when output_path is (or will be) holdout.manifest, or when a
    holdout.manifest already exists in the same directory — in either case holdout data
    must be strictly isolated from all other pools.
    """
    output_resolved = output_path.resolve()
    manifest_dir = output_resolved.parent
    holdout_path = manifest_dir / "holdout.manifest"

    if output_resolved.name != "holdout.manifest" and not holdout_path.exists():
        return []

    new_set = set(resolved_new_files)
    overlaps: list[tuple[str, Path]] = []
    # When creating holdout, check against all other manifests (dev, full, …).
    # When creating dev/full, only check against holdout — dev vs full overlap is allowed.
    if output_resolved.name == "holdout.manifest":
        candidates = sorted(manifest_dir.glob("*.manifest"))
    else:
        candidates = [holdout_path]
    for existing in candidates:
        if existing.resolve() == output_resolved:
            continue
        try:
            existing_files = set(_resolve_manifest_files(existing))
        except OSError:
            continue
        for f in sorted(new_set & existing_files):
            overlaps.append((existing.name, f))
    return overlaps


def apply_create_manifest(output_path: Path, lines: list[str]) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text("\n".join(lines) + "\n")


# ---------------------------------------------------------------------------
# CLI entry points (called from scripts/dwarf.py)
# ---------------------------------------------------------------------------

def run_status(root: str) -> None:
    root_path = Path(root)
    if not root_path.is_dir():
        raise SystemExit(f"error: {root_path} is not a directory")
    print_status(root_path)


def _run_bump(
    kind: str,
    reason: str,
    by: str | None,
    date: str | None,
    execute: bool,
    root: str,
    gen: str | None,
    to: str | None,
) -> None:
    root_path = Path(root)
    if not root_path.is_dir():
        raise SystemExit(f"error: {root_path} is not a directory")

    date = date or dt.date.today().isoformat()
    by = by if by is not None else _git_user_name()
    if gen is None:
        new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to)
    else:
        new_dirs, log_line = plan_bump_schema(root_path, kind, gen, reason, by, date, to)

    print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
    print("new directories:")
    for d in new_dirs:
        print(f"  {d}")
    print("VERSIONS.md entry:")
    print(f"  {log_line}")

    if not execute:
        print("\nDry run only — pass --execute to apply.")
        return
    apply_bump(root_path, new_dirs, log_line)
    print("\nDone.")


def run_bump_gen(
    kind: str,
    reason: str,
    by: str | None,
    date: str | None,
    execute: bool,
    root: str,
    to: str | None = None,
) -> None:
    _run_bump(kind, reason, by, date, execute, root, gen=None, to=to)


def run_bump_schema(
    kind: str,
    gen: str,
    reason: str,
    by: str | None,
    date: str | None,
    execute: bool,
    root: str,
    to: str | None = None,
) -> None:
    _run_bump(kind, reason, by, date, execute, root, gen=gen, to=to)


def run_update_manifest(
    manifests: list[str],
    schema: str | None,
    execute: bool,
    gen: str | None = None,
) -> None:
    if schema and not SCHEMA_RE.match(schema):
        raise SystemExit(f"error: --schema must look like 'schemaN', got {schema!r}")
    if gen and not GEN_RE.match(gen):
        raise SystemExit(f"error: --gen must look like 'genN', got {gen!r}")

    all_plans: list[tuple[Path, list[tuple[str, str | None]]]] = []
    all_missing: list[Path] = []

    for raw in manifests:
        mp = Path(raw)
        plan, missing = plan_update_manifest(mp, schema, gen)
        all_plans.append((mp.resolve(), plan))
        all_missing.extend(missing)

    print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
    for mp, plan in all_plans:
        changes = [(old, new) for old, new in plan if new is not None]
        print(f"\n{mp}  ({len(changes)} path(s) to update)")
        for old, new in changes:
            print(f"  - {old.strip()}")
            print(f"  + {new}")

    if all_missing:
        print(f"\nMISSING ({len(all_missing)} file(s) — target paths do not exist):")
        for p in all_missing:
            print(f"  {p}")
        if execute:
            raise SystemExit("error: refusing to write manifests with missing targets")

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

    for mp, plan in all_plans:
        apply_update_manifest(mp, plan)
    print("\nDone.")


def run_create_manifest(
    files: list[str],
    execute: bool,
    output: str | None = None,
    pool: str | None = None,
    type_: str | None = None,
    root: str = "/ceph/lbogner/geant_steps",
) -> None:
    if (output is None) == (pool is None):
        raise SystemExit("error: exactly one of --output or --pool is required")
    if pool is not None and type_ is None:
        raise SystemExit("error: --type is required when --pool is given")

    if pool is not None:
        output_path = Path(root) / "pools" / pool / f"{type_}.manifest"
    else:
        assert output is not None  # guaranteed by the exclusivity check above
        output_path = Path(output)

    parquet_files = [Path(f) for f in files]
    lines, missing, resolved = plan_create_manifest(output_path, parquet_files)
    overlaps = check_holdout_overlap(output_path, resolved)

    print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
    print(f"manifest: {output_path.resolve()}")
    for line in lines:
        print(f"  {line}")

    if missing:
        print(f"\nMISSING ({len(missing)} file(s) do not exist):")
        for p in missing:
            print(f"  {p}")

    if overlaps:
        print(f"\nHOLDOUT OVERLAP ({len(overlaps)} file(s) appear in other manifests):")
        for name, f in overlaps:
            print(f"  {f}  (also in {name})")

    if (missing or overlaps) and execute:
        raise SystemExit("error: refusing to write manifest (see above)")

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

    apply_create_manifest(output_path, lines)
    print("\nDone.")
