"""One bounded GPU case, with explicit opt-in experimental explosion."""

import argparse
import hashlib
import json
import os
import shutil
from pathlib import Path
import subprocess
import sys
import time

import numpy as np
from oracle import U, TIMES, LIMITS, initial, metrics, rectangles

NUMERICAL_REVISION = "a135252ae78b47b57d7e77d19001700b3bedac8b"
SOURCE_SHA256 = "a5dbfd804befce868067d79a23e5fd47dce7f505b3f9b1bb2145fb57cb5f2178"
BUILD_CAPTURE_SHA256 = "3dbaa6547cb0a4cb1541054f3b64ad8b6e18cc70527ac6166e56b944e9bd464a"
REFERENCE_HARNESS = "11901a7f77557a901e017a171eeae1758d584fa5"
INITIALIZATION = "physical FP64 equilibrium G at actual cell centers; rho=1+3p"


def sample_steps(n, early_pressure=False):
    """Root synchronization steps; early ladder matches the 707afe6 diagnostic."""
    main = [round(tstar * n / U) for tstar in TIMES]
    early = list(range(0, 4 * n + 1, n // 16)) if early_pressure else []
    return main, early


def runtime_source(repo):
    """Fail before CUDA if Python or the compiler snapshot comes from another tree."""
    import CelerisLab
    from CelerisLab.cuda import compiler_v2

    paths = dict(module=Path(CelerisLab.__file__).resolve(),
                 compiler=Path(compiler_v2.__file__).resolve(),
                 kernel_root=Path(compiler_v2.kernel_path("blocks.cu")).resolve().parent)
    package = repo / "src/CelerisLab"
    expected = dict(module=package / "__init__.py", compiler=package / "cuda/compiler_v2.py",
                    kernel_root=package / "lbm/kernels")
    if paths != expected:
        raise RuntimeError("Wrong imported source or CUDA source/include root")
    return dict(checkout=str(repo), **{key: str(path) for key, path in paths.items()})


def provenance(repo):
    git = lambda *a: subprocess.check_output(
        ["git", "-C", str(repo), *a], text=True
    ).strip()
    source_hash = hashlib.sha256()
    for path in sorted((repo / "src").rglob("*")):
        if path.is_file() and "__pycache__" not in path.parts and path.suffix != ".pyc":
            source_hash.update(
                str(path.relative_to(repo)).encode() + b"\0" + path.read_bytes() + b"\0"
            )
    return dict(
        revision=git("rev-parse", "HEAD"),
        dirty=bool(git("status", "--porcelain")),
        source_tree_sha256=source_hash.hexdigest(),
        harness_sha256={
            p.name: hashlib.sha256(p.read_bytes()).hexdigest()
            for p in Path(__file__).parent.glob("*.py")
        },
    )


def verify_source(repo, revision):
    """Exact clean harness commit, fixed production tree, and reviewed capture helper."""
    import oracle
    from tests.validation.multigrid import check_projected_pairs as gate

    harness = repo / "tests/validation/taylor_green"
    if (Path(__file__).resolve() != harness / "run_case.py"
            or Path(oracle.__file__).resolve() != harness / "oracle.py"):
        raise RuntimeError("Wrong runner or oracle checkout")
    helper = repo / "tests/validation/multigrid/check_projected_pairs.py"
    if (Path(gate.__file__).resolve() != helper
            or hashlib.sha256(helper.read_bytes()).hexdigest() != BUILD_CAPTURE_SHA256):
        raise RuntimeError("Wrong build-capture helper")
    if gate.preflight(repo, revision, "candidate") != NUMERICAL_REVISION:
        raise RuntimeError("Wrong production revision")
    result = provenance(repo)
    if result["source_tree_sha256"] != SOURCE_SHA256:
        raise RuntimeError("Wrong numerical source fingerprint")
    result["build_capture_sha256"] = BUILD_CAPTURE_SHA256
    return result


def build_capture(report, save):
    from tests.validation.multigrid.check_projected_pairs import capture_builds
    return capture_builds(report, save)


def runtime_identity(ctx):
    import pycuda
    import pycuda.driver as cuda

    nvcc = shutil.which("nvcc") or "/usr/local/cuda/bin/nvcc"
    return dict(device=dict(name=ctx.device.name(), pci_bus_id=ctx.device.pci_bus_id(),
                            capability=list(ctx.device.compute_capability()),
                            driver=cuda.get_driver_version(), runtime=list(cuda.get_version())),
                nvcc_path=str(Path(nvcc).resolve()),
                nvcc_version=subprocess.check_output([nvcc, "--version"], text=True).strip(),
                python=sys.version, numpy=np.__version__, pycuda=pycuda.VERSION_TEXT)


def config(n, re):
    from CelerisLab.config import LBMConfig

    cfg = LBMConfig(
        nx=n,
        ny=n,
        viscosity=U * n / re,
        velocity=0.0,
        rho=1.0,
        lattice_model="D2Q9",
        collision="SRT",
        streaming="double_buffer",
        store_precision="FP32",
        ddf_shifting=True,
        les_enabled=False,
    )
    cfg.validate()
    return cfg


def workload(masks, base_steps):
    cells = [int(np.count_nonzero(m)) for m in masks]
    hops = [base_steps * 2**level for level in range(len(masks))]
    updates = [count * hop for count, hop in zip(cells, hops)]
    return dict(
        active_cells_by_level=cells,
        level_hops=hops,
        active_collision_updates_by_level=updates,
        active_collision_updates=sum(updates),
        root_intervals=base_steps,
        excludes="ghost streaming, halo copies, transfers, initialization and diagnostics",
    )


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--repo", type=Path, required=True)
    p.add_argument("--expected-revision", required=True)
    p.add_argument("--run-gpu", action="store_true")
    p.add_argument("--dry-run", action="store_true")
    p.add_argument("--re", type=int, choices=[100, 1000], required=True)
    p.add_argument("--n", type=int, choices=[32, 64, 128, 256], required=True)
    p.add_argument("--depth", type=int, choices=[1, 2, 3], required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--explosion", choices=["constant", "projected_linear"], default="constant")
    p.add_argument("--early-pressure", action="store_true",
                   help="also save 65 scalar pressure samples at t*=k*.001875, k=0..64")
    args = p.parse_args()
    if args.output.exists():
        p.error("Output already exists; choose a new path to retain prior evidence")
    if args.explosion != "constant" and args.depth != 2:
        p.error("projected_linear requires depth 2")
    selected = os.environ.get("CUDA_VISIBLE_DEVICES", "")
    if not args.dry_run and (not args.run_gpu or not selected or "," in selected or selected == "-1"):
        p.error("--run-gpu and exactly one explicit CUDA_VISIBLE_DEVICES are required")
    repo = args.repo.resolve()
    sys.path[:0] = [str(repo / "src"), str(repo)]
    source = verify_source(repo, args.expected_revision)
    runtime = runtime_source(repo)
    if args.dry_run:
        print(json.dumps(dict(source=source, runtime_source=runtime, dry_run=True)))
        return 0
    from CelerisLab.cuda.context import CudaContext
    from CelerisLab.lbm.blocks import BlockLayout2D
    from CelerisLab.lbm.block_stepper import PeriodicBlockStepper
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
    from CelerisLab.lbm.refinement import ACTIVE
    from CelerisLab.lbm.hierarchy import build_hierarchy, leaf_partition

    cfg = config(args.n, args.re)
    refinements = [] if args.depth == 1 else rectangles(args.n, args.depth)
    h = build_hierarchy(cfg, refinements)
    np.testing.assert_array_equal(leaf_partition(h), np.ones((args.n, args.n)))
    masks = [k == ACTIVE for k in h.kinds]
    main_steps, early_steps = sample_steps(args.n, args.early_pressure)
    report = dict(
        schema=1,
        source=source,
        numerical_revision=NUMERICAL_REVISION,
        reference_harness=REFERENCE_HARNESS,
        actual_build_sidecars=[],
        runtime_identity=None,
        runtime_source=runtime,
        historical_baseline_revision="9805383f9c50d867e8555471baf5dbda8333a652",
        explosion=args.explosion,
        n=args.n,
        re_L=args.re,
        depth=args.depth,
        refinements=refinements,
        U=U,
        mach=U * np.sqrt(3),
        viscosity=cfg.viscosity,
        tau_by_level=[0.5 + 3 * c.viscosity for c in h.configs],
        initialization=INITIALIZATION,
        numerics=dict(lattice_model=cfg.lattice_model, collision=cfg.collision,
                      store_precision=cfg.store_precision, ddf_shifting=cfg.ddf_shifting,
                      streaming=cfg.streaming, les_enabled=cfg.les_enabled,
                      boundary="periodic", block_shape=[8, 8],
                      sample_phase="postcollision G at root synchronization"),
        cuda_visible_devices=selected,
        device_id=0,
        limits=LIMITS,
        workload=workload(masks, round(TIMES[-1] * args.n / U)),
        rows=[],
        early_pressure=dict(enabled=args.early_pressure, rows=[],
                            sampled_max_pressure_l2=None),
        completed=False,
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    def save():
        temporary = args.output.with_suffix(args.output.suffix + ".tmp")
        text = json.dumps(report, indent=2, allow_nan=False) + "\n"
        if len(text.encode()) > 1024 * 1024:
            raise RuntimeError("Scalar receipt exceeds 1 MiB cap")
        temporary.write_text(text)
        temporary.replace(args.output)

    save()
    started = time.monotonic()
    with CudaContext(0) as ctx, build_capture(report, save):
        report["runtime_identity"] = runtime_identity(ctx)
        save()
        if args.depth == 1:
            layout = BlockLayout2D.build(args.n, args.n, 8, 8, periodic=(True, True))
            run = PeriodicBlockStepper(cfg, layout, context=ctx)
        else:
            run = FixedPatchStepper(
                cfg,
                context=ctx,
                coarse_boundary="periodic",
                refinements=refinements,
                level_blocks=[(8, 8)] * args.depth,
                explosion=args.explosion,
            )
        with run:
            run.initialize(
                *[
                    initial(m.shape, level, args.n, cfg.viscosity)
                    for level, m in enumerate(masks)
                ]
            )
            initial_mass = None
            for step in sorted(set(main_steps + early_steps)):
                run.step(step - run.step_count)
                arrays = (run.download(),) if args.depth == 1 else run.download()
                row = metrics(arrays, masks, step, args.n, cfg.viscosity, initial_mass)
                if not all(np.isfinite(value) for value in row.values()):
                    raise FloatingPointError("nonfinite sample diagnostics")
                row["step"] = step
                if initial_mass is None:
                    initial_mass = row["mass"]
                if step in main_steps:
                    report["rows"].append(row)
                if step in early_steps:
                    trace = report["early_pressure"]
                    trace["rows"].append({key: row[key] for key in ("step", "tstar", "pressure_l2")})
                    trace["sampled_max_pressure_l2"] = max(r["pressure_l2"] for r in trace["rows"])
                save()
                print(
                    json.dumps(
                        dict(
                            event="PROGRESS" if step in main_steps else "EARLY_PRESSURE",
                            n=args.n,
                            re_L=args.re,
                            depth=args.depth,
                            **(row if step in main_steps else report["early_pressure"]["rows"][-1]),
                        )
                    ),
                    flush=True,
                )
                if step == 0:
                    # Exercise real kernels/halos/transfers on every level before READY.
                    # This is the first measured trajectory step, not discarded warmup.
                    run.step(1)
                    probe = (run.download(),) if args.depth == 1 else run.download()
                    probe_metrics = metrics(probe, masks, 1, args.n, cfg.viscosity, initial_mass)
                    if not all(np.isfinite(value) for value in probe_metrics.values()):
                        raise FloatingPointError("nonfinite READY diagnostics")
                    report["ready_after_step"] = 1
                    save()
                    print(
                        json.dumps(
                            dict(
                                event="READY",
                                step=1,
                                n=args.n,
                                re_L=args.re,
                                depth=args.depth,
                                workload=report["workload"],
                            )
                        ),
                        flush=True,
                    )
    report.update(
        completed=True,
        seconds=time.monotonic() - started,
        screening_pass=all(
            row[key] <= limit for row in report["rows"] for key, limit in LIMITS.items()
        ),
    )
    save()
    print(
        json.dumps(
            {k: report[k] for k in ("n", "re_L", "depth", "screening_pass", "seconds")}
        ),
        flush=True,
    )
    return 0 if report["screening_pass"] else 1


if __name__ == "__main__":
    raise SystemExit(main())