"""Single-process GPU branch runner for OID v2 A/B/C replay diagnostics."""
from __future__ import annotations

import argparse
from contextlib import ExitStack
import json
import os
import sys
from pathlib import Path

import numpy as np
import pycuda.driver as cuda

from LegacyCelerisLab import FlowField
from OID_analysis.configs import CENTER_Y, LEGACY_CFG_DIR, get_scene
from CCD_analysis.utils.cfd_interface import load_legacy_configs
from CCD_analysis.utils.dual_clock import DualClockCollector, field_steps_from_interval
from .artifacts import atomic_save_npz, atomic_write_json, validate_field_artifact
from CCD_analysis.utils.run_contract import global_cfd_lease, preflight_disk
from .contracts import LOGICAL_GPU, PHYSICAL_GPU, SHARED_CFD_LEASE, telemetry_hashes
from .manifest import sha256_file, validate_execution_environment

L0 = 20.0


def build_geometry(scene: str, device: int) -> FlowField:
    cfg = get_scene(scene)
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]), velocity=float(cfg["u0"]))
    ff = FlowField(field_cfg, cuda_cfg, device_id=device)
    if cfg.get("has_disturbance", False):
        ff.add_cylinder((10.0 * L0, CENTER_Y, 0.0), L0)
    for y_offset in (2.0, 0.0, -2.0):
        ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_offset * L0, 0.0), L0 / 4.0)
    ff.add_cylinder((cfg["pinball_front_x"] * L0, CENTER_Y, 0.0), L0 / 2.0)
    ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0)
    ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0)
    return ff


def restore_bundle(ff: FlowField, bundle_path: Path) -> dict[str, np.ndarray]:
    with np.load(bundle_path, allow_pickle=False) as archive:
        bundle = {key: archive[key].copy() for key in archive.files}
    required = {"current", "temp", "action"}
    if not required.issubset(bundle):
        raise ValueError(f"checkpoint bundle missing {sorted(required.difference(bundle))}")
    if bundle["current"].shape != ff.ddf.shape or bundle["temp"].shape != ff.ddf.shape:
        raise ValueError("checkpoint DDF shape differs from FlowField")
    if bundle["action"].shape != ff.action.shape:
        raise ValueError("checkpoint action shape differs from object ordering")
    ff.ddf[:] = bundle["current"]
    cuda.memcpy_htod(ff.ddf_gpu, bundle["current"])
    cuda.memcpy_htod(ff.temp_gpu, bundle["temp"])
    ff.action[:] = bundle["action"]
    cuda.memcpy_htod(ff.action_gpu, bundle["action"])
    return bundle


def scaled_actions(scene: str, normalized_actions: np.ndarray) -> np.ndarray:
    cfg = get_scene(scene)
    actions = np.asarray(normalized_actions, dtype=np.float32)
    output = np.zeros((len(actions), cfg["n_objects_env"]), dtype=np.float32)
    omega = (actions * cfg["action_scale"] + np.asarray(cfg["action_bias"], np.float32)) * cfg["u0"]
    output[:, -3:] = omega
    return output


def run_recorded_branch(*, scene: str, device: int, bundle_path: Path, trajectory_path: Path,
                        output_dir: Path, field_interval: int | None, count: int | None) -> dict:
    execution = validate_execution_environment(device=device, require_gpu=True)
    if execution["logical_device"] != LOGICAL_GPU:
        raise RuntimeError("isolated physical GPU 2 must be addressed as logical GPU 0")
    preflight_disk(output_dir, 0, reserve_bytes=1024**3)
    if scene not in {"karman_re100", "illusion_1.0L"}:
        raise ValueError("GPU pilot supports only planned scenes")
    if output_dir.exists():
        raise FileExistsError(output_dir)
    partial = output_dir.with_name(f"{output_dir.name}.partial.{os.getpid()}")
    if partial.exists():
        raise FileExistsError(partial)
    partial.mkdir(parents=True)
    ff = None
    lease_stack = ExitStack()
    try:
        lease_stack.enter_context(global_cfd_lease(
            SHARED_CFD_LEASE, run_id=output_dir.name, scene=scene,
            physical_gpu=PHYSICAL_GPU, command=sys.argv,
        ))
        with np.load(trajectory_path, allow_pickle=False) as trajectory:
            normalized = trajectory["actions"].copy()
        if count is not None:
            normalized = normalized[:count]
        actions = scaled_actions(scene, normalized)
        cfg = get_scene(scene)
        ff = build_geometry(scene, LOGICAL_GPU)
        restore_bundle(ff, bundle_path)
        collector = None
        if field_interval is not None:
            collector = DualClockCollector(ff, control_interval=cfg["sample_interval"],
                control_count=len(actions), u0=cfg["u0"],
                field_steps=field_steps_from_interval(cfg["sample_interval"] * len(actions), field_interval))
        sensors, forces = [], []
        obs_start, obs_end = cfg["obs_slice"]
        for index, action in enumerate(actions):
            ff.context.push()
            try:
                if collector is None:
                    ff.run(cfg["sample_interval"], action)
                else:
                    collector.run_interval(index, action)
            finally:
                ff.context.pop()
            raw = ff.obs.copy()[obs_start:obs_end]
            sensors.append(raw[:6]); forces.append(raw[6:12])
        telemetry_arrays = {
            "sensors": np.asarray(sensors, np.float32),
            "forces": np.asarray(forces, np.float32),
            "actions": normalized.astype(np.float32),
            "control_end_lattice_steps": np.arange(1, len(actions) + 1, dtype=np.int64) * cfg["sample_interval"],
            "lattice_steps": np.arange(1, len(actions) + 1, dtype=np.int64) * cfg["sample_interval"],
            "sample_ids": np.arange(len(actions), dtype=np.int64),
            "times": np.arange(1, len(actions) + 1, dtype=np.float64) * cfg["sample_interval"],
        }
        telemetry_path = atomic_save_npz(partial / "telemetry.npz", **telemetry_arrays)
        field_hashes = telemetry_hashes(telemetry_arrays)
        fields_validation = None
        if collector is not None:
            collector.save(partial / "fields.npz")
            with np.load(partial / "fields.npz", allow_pickle=False) as fields:
                fields_validation = validate_field_artifact(fields, expected_control_interval=cfg["sample_interval"])
        report = {"scene": scene, "count": len(actions), "control_interval": cfg["sample_interval"],
                  "field_interval": field_interval, "device_physical": device,
                  "device_logical": LOGICAL_GPU,
                  "checkpoint_bundle_sha256": sha256_file(bundle_path),
                  "trajectory_sha256": sha256_file(trajectory_path), "fields": fields_validation,
                  "telemetry_array_hashes": field_hashes}
        atomic_write_json(partial / "branch.json", report)
        atomic_write_json(partial / "_SUCCESS", {"telemetry_sha256": sha256_file(telemetry_path)})
        os.replace(partial, output_dir)
        return report
    except BaseException:
        if partial.exists():
            os.replace(partial, partial.with_name(partial.name.replace(".partial.", ".quarantine.")))
        raise
    finally:
        if ff is not None:
            try:
                ff.context.detach()
            except Exception:
                pass
            del ff
        lease_stack.close()


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--scene", required=True)
    parser.add_argument("--device", type=int, required=True)
    parser.add_argument("--bundle", type=Path, required=True)
    parser.add_argument("--trajectory", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--field-interval", type=int)
    parser.add_argument("--count", type=int)
    args = parser.parse_args(argv)
    report = run_recorded_branch(scene=args.scene, device=args.device, bundle_path=args.bundle,
        trajectory_path=args.trajectory, output_dir=args.output_dir,
        field_interval=args.field_interval, count=args.count)
    print(json.dumps(report, indent=2))
    return 0


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