"""Online-policy and zero-control paired branches from one complete checkpoint."""
from __future__ import annotations

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

import numpy as np

from CCD_analysis.utils.cfd_interface import build_observation, load_ppo_model
from CCD_analysis.utils.dual_clock import DualClockCollector, field_steps_from_interval
from CCD_analysis.utils.resampling import gen_target_states_at
from OID_analysis.configs import get_scene, model_path_for_scene
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 .collect import field_level_telemetry, telemetry_snapshot
from .contracts import BRANCH_SCHEMA_ID, LOGICAL_GPU, PHYSICAL_GPU, SHARED_CFD_LEASE, telemetry_hashes
from .gpu import build_geometry, restore_bundle
from .manifest import sha256_file, validate_execution_environment


def run_branch(*, scene: str, mode: str, device: int, bundle_path: Path, output_dir: Path,
               count: int, field_interval: int, bootstrap_path: Path | None = None,
               norm_path: Path | None = None, harmonics_path: Path | None = None,
               role: str | None = None, parent_id: str | None = None,
               sample_id_start: int = 0, lattice_dt: float = 1.0) -> 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 mode not in {"online", "zero"}:
        raise ValueError("mode must be online or zero")
    role = role or ("controlled" if mode == "online" else "zero_baseline")
    if role not in {"controlled", "zero_baseline", "baseline_twin"}:
        raise ValueError("role must be a canonical paired branch role")
    if (mode == "online") != (role == "controlled"):
        raise ValueError("controlled role requires online mode; baseline roles require zero mode")
    if not parent_id:
        raise ValueError("formal paired collection requires parent_id")
    if output_dir.exists():
        raise FileExistsError(output_dir)
    partial = output_dir.with_name(f"{output_dir.name}.partial.{os.getpid()}")
    partial.mkdir(parents=True, exist_ok=False)
    ff = model = 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,
        ))
        cfg = get_scene(scene)
        ff = build_geometry(scene, LOGICAL_GPU)
        bundle = restore_bundle(ff, bundle_path)
        norm = json.loads(norm_path.read_text()) if norm_path else None
        harmonics = json.loads(harmonics_path.read_text()) if harmonics_path else None
        bootstrap = None
        if bootstrap_path:
            with np.load(bootstrap_path, allow_pickle=False) as archive:
                bootstrap = archive["actions"][0].copy()
        if mode == "online":
            if norm is None or bootstrap is None:
                raise ValueError("online mode requires norm and bootstrap trajectory")
            model = load_ppo_model(model_path_for_scene(scene), device=f"cuda:{LOGICAL_GPU}",
                                   s_dim=cfg["s_dim"], a_dim=3)
            model.set_random_seed(19 if scene.startswith("illusion") else 0)
        collector = DualClockCollector(ff, control_interval=cfg["sample_interval"], control_count=count,
            u0=cfg["u0"], field_steps=field_steps_from_interval(cfg["sample_interval"] * count, field_interval),
            snapshot_callback=telemetry_snapshot(cfg["obs_slice"]))
        obs = None
        sensors=[]; forces=[]; actions=[]
        for index in range(count):
            if mode == "zero":
                normalized = np.zeros(3, np.float32)
                applied = np.zeros(cfg["n_objects_env"], np.float32)
            else:
                normalized = bootstrap if index == 0 else np.asarray(model.predict(obs, deterministic=True)[0], np.float32).reshape(-1)
                applied = np.zeros(cfg["n_objects_env"], np.float32)
                applied[-3:] = (normalized * cfg["action_scale"] + np.asarray(cfg["action_bias"], np.float32)) * cfg["u0"]
            ff.context.push()
            try:
                collector.run_interval(index, applied)
            finally:
                ff.context.pop()
            raw = ff.obs.copy()[slice(*cfg["obs_slice"])]
            sensors.append(raw[:6]); forces.append(raw[6:12]); actions.append(normalized)
            if mode == "online":
                obs12 = build_observation(raw, norm)
                if cfg["s_dim"] == 14:
                    target = gen_target_states_at(index + 1, harmonics)
                    obs = np.clip(np.hstack((obs12, target[:2] / float(norm["force_norm_fact"]))), -1, 1).astype(np.float32)
                else:
                    obs = obs12
        collector.save(partial / "fields.npz")
        telemetry_arrays = field_level_telemetry(
            collector, sample_id_start=sample_id_start, lattice_dt=lattice_dt)
        telemetry = atomic_save_npz(partial / "telemetry.npz", **telemetry_arrays)
        field_hashes = telemetry_hashes(telemetry_arrays)
        with np.load(partial / "fields.npz", allow_pickle=False) as fields:
            validation=validate_field_artifact(fields, expected_control_interval=cfg["sample_interval"])
        with np.load(partial / "fields.npz", allow_pickle=False) as field_archive:
            field_content = hashlib.sha256()
            for key in ("ux", "uy"):
                value = np.ascontiguousarray(field_archive[key])
                field_content.update(str(value.dtype).encode())
                field_content.update(np.asarray(value.shape, np.int64).tobytes())
                field_content.update(value.tobytes())
        fields_sha256 = field_content.hexdigest()
        telemetry_content = hashlib.sha256("".join(field_hashes[key] for key in sorted(field_hashes)).encode()).hexdigest()
        timeline_sha256 = field_hashes["lattice_steps"]
        report={"schema_id":BRANCH_SCHEMA_ID,"scene":scene,"mode":mode,"role":role,
                "action_kind":"controlled" if role == "controlled" else "zero",
                "parent_id":parent_id,"count":count,"device_physical":device,
                "device_logical":LOGICAL_GPU,
                "control_interval":cfg["sample_interval"],"field_interval":field_interval,
                "checkpoint_sha256":sha256_file(bundle_path),"fields":validation,
                "fields_sha256":fields_sha256,"telemetry_sha256":telemetry_content,
                "timeline_sha256":timeline_sha256,"sample_ids_sha256":field_hashes["sample_ids"],
                "times_sha256":field_hashes["times"],"telemetry_array_hashes":field_hashes,
                "units":{"velocity":"normalized_by_u0","lattice_steps":"LBM steps"},
                "object_order": ["disturbance","sensor_upper","sensor_center","sensor_lower","front","upper","lower"] if cfg.get("has_disturbance") else ["sensor_upper","sensor_center","sensor_lower","front","upper","lower"]}
        atomic_write_json(partial/"branch.json",report); atomic_write_json(partial/"_SUCCESS",{"telemetry_sha256":report["telemetry_sha256"]})
        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 model,ff
        lease_stack.close()


def main(argv=None):
    p=argparse.ArgumentParser(); p.add_argument("--scene",required=True); p.add_argument("--mode",choices=("online","zero"),required=True)
    p.add_argument("--device",type=int,required=True); p.add_argument("--bundle",type=Path,required=True); p.add_argument("--output-dir",type=Path,required=True)
    p.add_argument("--count",type=int,required=True); p.add_argument("--field-interval",type=int,required=True); p.add_argument("--bootstrap",type=Path)
    p.add_argument("--norm",type=Path); p.add_argument("--harmonics",type=Path)
    p.add_argument("--role",choices=("controlled","zero_baseline","baseline_twin"),required=True)
    p.add_argument("--parent-id",required=True); p.add_argument("--sample-id-start",type=int,default=0)
    p.add_argument("--lattice-dt",type=float,default=1.0); a=p.parse_args(argv)
    print(json.dumps(run_branch(scene=a.scene,mode=a.mode,device=a.device,bundle_path=a.bundle,output_dir=a.output_dir,
        count=a.count,field_interval=a.field_interval,bootstrap_path=a.bootstrap,norm_path=a.norm,harmonics_path=a.harmonics,
        role=a.role,parent_id=a.parent_id,sample_id_start=a.sample_id_start,lattice_dt=a.lattice_dt),indent=2)); return 0

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