open Printf
open Yojson.Basic.Util

let ( // ) = Filename.concat
let shell_quote = Filename.quote
let ensure_dir path =
  if path <> "" && not (Sys.file_exists path) then
    let code = Sys.command ("mkdir -p " ^ shell_quote path) in
    if code <> 0 then failwith ("mkdir -p failed for " ^ path)
let ensure_parent_dir path =
  let dir = Filename.dirname path in
  if dir <> "." then ensure_dir dir
let write_lines path lines =
  ensure_parent_dir path;
  let oc = open_out path in
  List.iter (fun line -> output_string oc line; output_char oc '\n') lines;
  close_out oc
let json_to_file path j = ensure_parent_dir path; Yojson.Basic.to_file path j
let f_default obj key default =
  try
    match member key obj with
    | `Float f -> f
    | `Int i -> float_of_int i
    | `String s -> float_of_string s
    | _ -> default
  with _ -> default
let i_default obj key default = try member key obj |> to_int with _ -> default
let s_default obj key default = try member key obj |> to_string with _ -> default
let b_default obj key default = try member key obj |> to_bool with _ -> default
let clamp lo hi x = max lo (min hi x)
let fclamp lo hi x = max lo (min hi x)
let div0 a b = if b <= 0. then 0. else a /. b
let ceil_div a b = if a <= 0 then 0 else (a + b - 1) / b
let csv_escape s =
  if String.exists (function ',' | '"' | '\n' -> true | _ -> false) s then
    "\"" ^ String.concat "\"\"" (String.split_on_char '"' s) ^ "\""
  else s

let poisson_like lambda =
  (* Deterministic stochastic-ish arrival process once Random.init is set.
     We use bounded normal-like jitter rather than exact poisson to keep tests
     stable while still producing congestion bursts. *)
  if lambda <= 0. then 0
  else
    let jitter = (Random.float 2.0 -. 1.0) *. sqrt (max 1. lambda) in
    max 0 (int_of_float (floor (lambda +. jitter +. 0.5)))

let percentile p xs =
  match xs with
  | [] -> 0.
  | _ ->
      let arr = Array.of_list xs in
      Array.sort compare arr;
      let n = Array.length arr in
      let idx = int_of_float (ceil ((float_of_int n *. p) -. 1.0)) in
      arr.(clamp 0 (n - 1) idx)

let queue_to_list q = List.of_seq (Queue.to_seq q)
let push_window q cap v =
  Queue.add v q;
  while Queue.length q > cap do ignore (Queue.take q) done

let mean xs = match xs with [] -> 0. | _ -> List.fold_left ( +. ) 0. xs /. float_of_int (List.length xs)

(* -------------------------------------------------------------------------- *)
(* Scenario and topology                                                       *)
(* -------------------------------------------------------------------------- *)

type system_kind = PCS | USM
type scale_mode = Scale_none | Scale_vertical | Scale_horizontal | Scale_both

type workload = {
  interactive_rps : float;
  batch_rps : float;
  restore_rps : float;
  metadata_ops_rps : float;
  bulk_write_mbps : float;
  kv_hydrate_kb : float;
  kv_write_kb : float;
  object_read_kb : float;
  object_write_kb : float;
  avg_prompt_tokens : int;
  avg_output_tokens : int;
  gpu_flop_t_per_prompt_1k : float;
  gpu_flop_t_per_decode_token : float;
  gpu_hbm_mb_per_req : float;
}

type topology = {
  mutable pods : int;
  mutable racks_per_pod : int;
  mutable gpu_nodes_per_rack : int;
  mutable gpus_per_node : int;
  mutable gpu_tflops_per_gpu : float;
  mutable gpu_hbm_gbps_per_gpu : float;
  mutable gpu_sm_count : int;
  mutable rdma_rails : int;
  mutable roce_gbps_per_rail : float;
  mutable roce_mtu_bytes : int;
  mutable roce_hops : int;
  mutable roce_switch_latency_us : float;
  mutable roce_nic_latency_us : float;
  mutable roce_propagation_us : float;
  mutable roce_ecn_util_threshold : float;
  mutable roce_pfc_util_threshold : float;
  mutable roce_header_bytes : int;
  mutable dpu_paths : int;
  mutable dpu_gbps_per_path : float;
  mutable storage_nodes : int;
  mutable zfs_txg_timeout_us : int;
  mutable zfs_dirty_data_max_mb : float;
  mutable zfs_write_gbps_per_node : float;
  mutable zfs_sync_threads_per_node : int;
  mutable metadata_domains : int;
  mutable metadata_shards_per_domain : int;
  mutable api_nodes : int;
  mutable orchestrators : int;
  mutable interactive_workers : int;
  mutable restore_workers : int;
  mutable bulk_workers : int;
}

type policy = {
  target_roce_p99_us : float;
  max_roce_util : float;
  target_txg_lag_us : float;
  target_gpu_kernel_p99_us : float;
  max_gpu_util : float;
  target_metadata_p99_us : float;
  max_queue_age_us : float;
  scale_cooldown_us : int;
  max_rdma_rails : int;
  max_storage_nodes : int;
  max_gpu_nodes_per_rack : int;
  max_gpus_per_node : int;
  max_metadata_domains : int;
  max_metadata_shards_per_domain : int;
}

type scenario = {
  name : string;
  system : system_kind;
  duration_us : int;
  tick_us : int;
  seed : int;
  scale_mode : scale_mode;
  workload : workload;
  topology : topology;
  policy : policy;
}

let system_of_string = function
  | "pcs" | "PCS" -> PCS
  | "usm" | "USM" -> USM
  | s -> invalid_arg ("unknown system: " ^ s)

let string_of_system = function PCS -> "pcs" | USM -> "usm"

let scale_mode_of_string = function
  | "none" -> Scale_none
  | "vertical" -> Scale_vertical
  | "horizontal" -> Scale_horizontal
  | "both" -> Scale_both
  | s -> invalid_arg ("unknown scale_mode: " ^ s)

let string_of_scale_mode = function
  | Scale_none -> "none"
  | Scale_vertical -> "vertical"
  | Scale_horizontal -> "horizontal"
  | Scale_both -> "both"

let parse_scenario default_system path =
  let j = Yojson.Basic.from_file path in
  let workload_j = member "workload" j in
  let topology_j = member "topology" j in
  let policy_j = member "policy" j in
  {
    name = s_default j "name" "scenario";
    system = system_of_string (s_default j "system" (string_of_system default_system));
    duration_us = i_default j "duration_us" 3_000_000;
    tick_us = i_default j "tick_us" 1_000;
    seed = i_default j "seed" 42;
    scale_mode = scale_mode_of_string (s_default j "scale_mode" "both");
    workload = {
      interactive_rps = f_default workload_j "interactive_rps" 1000.;
      batch_rps = f_default workload_j "batch_rps" 120.;
      restore_rps = f_default workload_j "restore_rps" 20.;
      metadata_ops_rps = f_default workload_j "metadata_ops_rps" 2000.;
      bulk_write_mbps = f_default workload_j "bulk_write_mbps" 9000.;
      kv_hydrate_kb = f_default workload_j "kv_hydrate_kb" 768.;
      kv_write_kb = f_default workload_j "kv_write_kb" 192.;
      object_read_kb = f_default workload_j "object_read_kb" 512.;
      object_write_kb = f_default workload_j "object_write_kb" 1024.;
      avg_prompt_tokens = i_default workload_j "avg_prompt_tokens" 1536;
      avg_output_tokens = i_default workload_j "avg_output_tokens" 128;
      gpu_flop_t_per_prompt_1k = f_default workload_j "gpu_flop_t_per_prompt_1k" 0.9;
      gpu_flop_t_per_decode_token = f_default workload_j "gpu_flop_t_per_decode_token" 0.030;
      gpu_hbm_mb_per_req = f_default workload_j "gpu_hbm_mb_per_req" 128.;
    };
    topology = {
      pods = i_default topology_j "pods" 1;
      racks_per_pod = i_default topology_j "racks_per_pod" 2;
      gpu_nodes_per_rack = i_default topology_j "gpu_nodes_per_rack" 4;
      gpus_per_node = i_default topology_j "gpus_per_node" 8;
      gpu_tflops_per_gpu = f_default topology_j "gpu_tflops_per_gpu" 650.;
      gpu_hbm_gbps_per_gpu = f_default topology_j "gpu_hbm_gbps_per_gpu" 3000.;
      gpu_sm_count = i_default topology_j "gpu_sm_count" 132;
      rdma_rails = i_default topology_j "rdma_rails" 2;
      roce_gbps_per_rail = f_default topology_j "roce_gbps_per_rail" 200.;
      roce_mtu_bytes = i_default topology_j "roce_mtu_bytes" 9000;
      roce_hops = i_default topology_j "roce_hops" 4;
      roce_switch_latency_us = f_default topology_j "roce_switch_latency_us" 0.45;
      roce_nic_latency_us = f_default topology_j "roce_nic_latency_us" 1.20;
      roce_propagation_us = f_default topology_j "roce_propagation_us" 0.20;
      roce_ecn_util_threshold = f_default topology_j "roce_ecn_util_threshold" 0.65;
      roce_pfc_util_threshold = f_default topology_j "roce_pfc_util_threshold" 0.92;
      roce_header_bytes = i_default topology_j "roce_header_bytes" 86;
      dpu_paths = i_default topology_j "dpu_paths" 8;
      dpu_gbps_per_path = f_default topology_j "dpu_gbps_per_path" 100.;
      storage_nodes = i_default topology_j "storage_nodes" 8;
      zfs_txg_timeout_us = i_default topology_j "zfs_txg_timeout_us" 120_000;
      zfs_dirty_data_max_mb = f_default topology_j "zfs_dirty_data_max_mb" 4096.;
      zfs_write_gbps_per_node = f_default topology_j "zfs_write_gbps_per_node" 18.;
      zfs_sync_threads_per_node = i_default topology_j "zfs_sync_threads_per_node" 8;
      metadata_domains = i_default topology_j "metadata_domains" 4;
      metadata_shards_per_domain = i_default topology_j "metadata_shards_per_domain" 8;
      api_nodes = i_default topology_j "api_nodes" 6;
      orchestrators = i_default topology_j "orchestrators" 3;
      interactive_workers = i_default topology_j "interactive_workers" 16;
      restore_workers = i_default topology_j "restore_workers" 8;
      bulk_workers = i_default topology_j "bulk_workers" 8;
    };
    policy = {
      target_roce_p99_us = f_default policy_j "target_roce_p99_us" 18.;
      max_roce_util = f_default policy_j "max_roce_util" 0.72;
      target_txg_lag_us = f_default policy_j "target_txg_lag_us" 250_000.;
      target_gpu_kernel_p99_us = f_default policy_j "target_gpu_kernel_p99_us" 1400.;
      max_gpu_util = f_default policy_j "max_gpu_util" 0.78;
      target_metadata_p99_us = f_default policy_j "target_metadata_p99_us" 1500.;
      max_queue_age_us = f_default policy_j "max_queue_age_us" 80_000.;
      scale_cooldown_us = i_default policy_j "scale_cooldown_us" 80_000;
      max_rdma_rails = i_default policy_j "max_rdma_rails" 16;
      max_storage_nodes = i_default policy_j "max_storage_nodes" 64;
      max_gpu_nodes_per_rack = i_default policy_j "max_gpu_nodes_per_rack" 32;
      max_gpus_per_node = i_default policy_j "max_gpus_per_node" 16;
      max_metadata_domains = i_default policy_j "max_metadata_domains" 32;
      max_metadata_shards_per_domain = i_default policy_j "max_metadata_shards_per_domain" 64;
    };
  }

(* -------------------------------------------------------------------------- *)
(* SST-style event/state model                                                 *)
(* -------------------------------------------------------------------------- *)

module Sst = struct
  type component_kind =
    | WorkloadGen | Scheduler | MetadataRouter | GpuExecutor | RdmaNic
    | RoceSwitch | ZfsPool | BackpressureController | ScalingController
    | StorageAdapter

  let string_of_component_kind = function
    | WorkloadGen -> "WorkloadGen"
    | Scheduler -> "Scheduler"
    | MetadataRouter -> "MetadataRouter"
    | GpuExecutor -> "GpuExecutor"
    | RdmaNic -> "RdmaNic"
    | RoceSwitch -> "RoceSwitch"
    | ZfsPool -> "ZfsPool"
    | BackpressureController -> "BackpressureController"
    | ScalingController -> "ScalingController"
    | StorageAdapter -> "StorageAdapter"

  type event_kind =
    | WorkloadArrival | AdmissionDecision | MetadataLookup | GpuKernelSubmit
    | RdmaPacketBatch | ZfsTxgWrite | ZfsTxgStateTransition | ScaleEvaluate
    | ScaleApply | BackpressureSignal

  let string_of_event_kind = function
    | WorkloadArrival -> "WorkloadArrival"
    | AdmissionDecision -> "AdmissionDecision"
    | MetadataLookup -> "MetadataLookup"
    | GpuKernelSubmit -> "GpuKernelSubmit"
    | RdmaPacketBatch -> "RdmaPacketBatch"
    | ZfsTxgWrite -> "ZfsTxgWrite"
    | ZfsTxgStateTransition -> "ZfsTxgStateTransition"
    | ScaleEvaluate -> "ScaleEvaluate"
    | ScaleApply -> "ScaleApply"
    | BackpressureSignal -> "BackpressureSignal"

  type component = { id : string; kind : component_kind; ports : string list }

  type event = {
    time_us : int;
    src : string;
    dst : string;
    kind : event_kind;
    bytes : int;
    detail : string;
  }

  let make_components system =
    let common = [
      { id = "workload"; kind = WorkloadGen; ports = ["out"] };
      { id = "scheduler"; kind = Scheduler; ports = ["in"; "gpu"; "rdma"; "scale"] };
      { id = "rdma-nic"; kind = RdmaNic; ports = ["host"; "fabric"] };
      { id = "roce-switch"; kind = RoceSwitch; ports = ["fabric"] };
      { id = "zfs-pool"; kind = ZfsPool; ports = ["nvmeof"; "txg"] };
      { id = "scaler"; kind = ScalingController; ports = ["metrics"; "actions"] };
      { id = "backpressure"; kind = BackpressureController; ports = ["metrics"; "admission"] };
    ] in
    match system with
    | PCS -> { id = "gpu-executor"; kind = GpuExecutor; ports = ["kernels"; "rdma"] } :: common
    | USM ->
        { id = "metadata-router"; kind = MetadataRouter; ports = ["lookup"; "journal"] } ::
        { id = "storage-adapter"; kind = StorageAdapter; ports = ["io"; "rdma"] } :: common
end

(* -------------------------------------------------------------------------- *)
(* Simulation state                                                            *)
(* -------------------------------------------------------------------------- *)

type txg_phase = Open | Quiescing | Syncing | Synced
let string_of_txg_phase = function Open -> "OPEN" | Quiescing -> "QUIESCING" | Syncing -> "SYNCING" | Synced -> "SYNCED"

type txg = {
  id : int;
  open_us : int;
  mutable phase : txg_phase;
  mutable dirty_mb : float;
  mutable blocks : int;
  mutable quiesce_us : int option;
  mutable sync_start_us : int option;
  mutable sync_done_us : int option;
}

type sim_state = {
  sc : scenario;
  mutable now_us : int;
  mutable admitted : int;
  mutable rejected : int;
  mutable total_requests : int;
  mutable total_packets : int;
  mutable total_cnp : int;
  mutable total_ecn_marks : int;
  mutable total_pfc_pause_us : float;
  mutable scaling_actions : int;
  mutable last_scale_us : int;
  mutable gpu_queue_us : float;
  mutable metadata_queue_age_us : float;
  mutable rdma_backlog_bytes : float;
  mutable current_txg : txg;
  mutable syncing_txgs : txg list;
  mutable next_txg_id : int;
  mutable packet_batch_id : int;
  mutable kernel_id : int;
  mutable sst_event_id : int;
  roce_latency_samples : float Queue.t;
  roce_util_samples : float Queue.t;
  gpu_kernel_samples : float Queue.t;
  gpu_util_samples : float Queue.t;
  txg_lag_samples : float Queue.t;
  metadata_latency_samples : float Queue.t;
  metadata_queue_samples : float Queue.t;
  timeline_lines : string list ref;
  rdma_lines : string list ref;
  txg_lines : string list ref;
  gpu_lines : string list ref;
  sst_lines : string list ref;
  scale_lines : string list ref;
  component_lines : string list ref;
}

let total_gpu_nodes t = t.pods * t.racks_per_pod * t.gpu_nodes_per_rack
let total_gpus t = total_gpu_nodes t * t.gpus_per_node
let total_metadata_shards t = t.metadata_domains * t.metadata_shards_per_domain

let new_txg id now = { id; open_us = now; phase = Open; dirty_mb = 0.; blocks = 0; quiesce_us = None; sync_start_us = None; sync_done_us = None }

let new_state sc =
  Random.init sc.seed;
  let component_lines =
    Sst.make_components sc.system
    |> List.map (fun c -> sprintf "%s,%s,%s" c.Sst.id (Sst.string_of_component_kind c.Sst.kind) (String.concat ":" c.Sst.ports))
  in
  {
    sc;
    now_us = 0;
    admitted = 0; rejected = 0; total_requests = 0;
    total_packets = 0; total_cnp = 0; total_ecn_marks = 0; total_pfc_pause_us = 0.;
    scaling_actions = 0; last_scale_us = -1_000_000_000;
    gpu_queue_us = 0.; metadata_queue_age_us = 0.; rdma_backlog_bytes = 0.;
    current_txg = new_txg 1 0; syncing_txgs = []; next_txg_id = 2;
    packet_batch_id = 0; kernel_id = 0; sst_event_id = 0;
    roce_latency_samples = Queue.create (); roce_util_samples = Queue.create ();
    gpu_kernel_samples = Queue.create (); gpu_util_samples = Queue.create ();
    txg_lag_samples = Queue.create (); metadata_latency_samples = Queue.create (); metadata_queue_samples = Queue.create ();
    timeline_lines = ref ["time_us,system,requests,admitted,rejected,rdma_packets,roce_util,roce_p99_us,gpu_util,gpu_p99_us,txg_dirty_mb,txg_lag_us,metadata_p99_us,metadata_queue_age_us,pods,gpu_nodes,gpus,rdma_rails,roce_gbps_per_rail,storage_nodes,metadata_domains,metadata_shards"];
    rdma_lines = ref ["time_us,batch_id,traffic_class,opcode,qp,psn_start,payload_bytes,packet_count,wire_bytes,mtu,rails,gbps_per_rail,utilization,serialization_us,queue_delay_us,base_latency_us,p99_latency_us,ecn_marks,cnp_packets,pfc_pause_us,retransmits"];
    txg_lines = ref ["time_us,txg_id,event,phase,dirty_mb,blocks,open_us,quiesce_us,sync_start_us,sync_done_us,storage_nodes,write_gbps_per_node,sync_threads_per_node,estimated_lag_us"];
    gpu_lines = ref ["time_us,kernel_id,kernel_type,batch_size,gpus,node_count,flops_t,hbm_bytes,rdma_bytes,occupancy,compute_us,mem_us,rdma_wait_us,queue_us,duration_us"];
    sst_lines = ref ["event_id,time_us,src,dst,event_kind,bytes,detail"];
    scale_lines = ref ["time_us,system,scale_mode,action,dimension,reason,before,after,roce_p99_us,roce_util,gpu_p99_us,gpu_util,txg_lag_us,metadata_p99_us,metadata_queue_age_us"];
    component_lines = ref (List.rev ("component_id,component_kind,ports" :: component_lines));
  }

let add_line r line = r := line :: !(r)

let emit_sst st src dst kind bytes detail =
  st.sst_event_id <- st.sst_event_id + 1;
  add_line st.sst_lines
    (sprintf "%d,%d,%s,%s,%s,%d,%s" st.sst_event_id st.now_us src dst (Sst.string_of_event_kind kind) bytes (csv_escape detail))

(* -------------------------------------------------------------------------- *)
(* RoCE-v2 / RDMA packet model                                                 *)
(* -------------------------------------------------------------------------- *)

type rdma_flow = { traffic_class : string; opcode : string; qp : int; payload_bytes : int }

type rdma_tick_metrics = { roce_util : float; roce_p99_us : float; packets : int; ecn : int; cnp : int; pfc_us : float; wire_bytes : int }

let model_rdma st flows =
  let t = st.sc.topology in
  let tick_us = st.sc.tick_us in
  let payload_per_packet = max 256 (t.roce_mtu_bytes - t.roce_header_bytes) in
  let flow_infos =
    List.mapi (fun idx f ->
      let packets = ceil_div f.payload_bytes payload_per_packet in
      let wire = f.payload_bytes + packets * (t.roce_header_bytes + 20) in
      (idx, f, packets, wire)) flows
  in
  let total_wire = List.fold_left (fun acc (_, _, _, wire) -> acc + wire) 0 flow_infos in
  let capacity_bytes_per_us = float_of_int t.rdma_rails *. t.roce_gbps_per_rail *. 125. in
  let tick_capacity = capacity_bytes_per_us *. float_of_int tick_us in
  let util = div0 (float_of_int total_wire +. st.rdma_backlog_bytes) tick_capacity in
  let overload = max 0. (util -. 1.0) in
  let queue_delay_us = overload *. float_of_int tick_us *. 0.35 in
  let base_latency_us = t.roce_nic_latency_us *. 2.0 +. float_of_int t.roce_hops *. t.roce_switch_latency_us +. t.roce_propagation_us in
  let ecn_ratio = if util > t.roce_ecn_util_threshold then min 0.25 ((util -. t.roce_ecn_util_threshold) *. 0.15) else 0. in
  let pfc_pause_us = if util > t.roce_pfc_util_threshold then min (0.75 *. float_of_int tick_us) ((util -. t.roce_pfc_util_threshold) *. float_of_int tick_us *. 0.80) else 0. in
  let retransmit_ratio = if util > 1.20 then min 0.05 ((util -. 1.20) *. 0.02) else 0. in
  let latencies = ref [] in
  let total_packets = ref 0 in
  let total_ecn = ref 0 in
  let total_cnp = ref 0 in
  List.iter (fun (idx, f, packets, wire) ->
    if packets > 0 then begin
      st.packet_batch_id <- st.packet_batch_id + 1;
      let serialization_us = div0 (float_of_int wire) capacity_bytes_per_us in
      let qp_hash_jitter = float_of_int ((f.qp * 13 + idx * 7) mod 17) *. 0.03 in
      let p99_latency = base_latency_us +. serialization_us +. queue_delay_us +. pfc_pause_us +. qp_hash_jitter in
      let ecn_marks = int_of_float (float_of_int packets *. ecn_ratio) in
      let cnp_packets = if ecn_marks > 0 then max 1 (ecn_marks / 64) else 0 in
      let retransmits = int_of_float (float_of_int packets *. retransmit_ratio) in
      let psn_start = st.packet_batch_id * 100000 in
      total_packets := !total_packets + packets;
      total_ecn := !total_ecn + ecn_marks;
      total_cnp := !total_cnp + cnp_packets;
      latencies := p99_latency :: !latencies;
      st.total_packets <- st.total_packets + packets;
      st.total_ecn_marks <- st.total_ecn_marks + ecn_marks;
      st.total_cnp <- st.total_cnp + cnp_packets;
      st.total_pfc_pause_us <- st.total_pfc_pause_us +. pfc_pause_us;
      add_line st.rdma_lines
        (sprintf "%d,%d,%s,%s,%d,%d,%d,%d,%d,%d,%d,%.3f,%.6f,%.6f,%.6f,%.6f,%.6f,%d,%d,%.6f,%d"
           st.now_us st.packet_batch_id f.traffic_class f.opcode f.qp psn_start f.payload_bytes packets wire
           t.roce_mtu_bytes t.rdma_rails t.roce_gbps_per_rail util serialization_us queue_delay_us base_latency_us
           p99_latency ecn_marks cnp_packets pfc_pause_us retransmits);
      emit_sst st "rdma-nic" "roce-switch" Sst.RdmaPacketBatch wire
        (sprintf "%s opcode=%s packets=%d util=%.3f" f.traffic_class f.opcode packets util)
    end) flow_infos;
  st.rdma_backlog_bytes <- max 0. ((float_of_int total_wire +. st.rdma_backlog_bytes) -. tick_capacity);
  let p99 = percentile 0.99 !latencies in
  push_window st.roce_latency_samples 512 p99;
  push_window st.roce_util_samples 512 util;
  { roce_util = util; roce_p99_us = p99; packets = !total_packets; ecn = !total_ecn; cnp = !total_cnp; pfc_us = pfc_pause_us; wire_bytes = total_wire }

(* -------------------------------------------------------------------------- *)
(* GPU kernel model                                                            *)
(* -------------------------------------------------------------------------- *)

type kernel_spec = { kernel_type : string; batch_size : int; flops_t : float; hbm_bytes : int; rdma_bytes : int; occupancy : float }

type gpu_tick_metrics = { gpu_util : float; gpu_p99_us : float; kernel_count : int }

let model_gpu_kernels st kernels roce_wait_us =
  let t = st.sc.topology in
  let gpus = max 1 (total_gpus t) in
  let node_count = max 1 (total_gpu_nodes t) in
  let tick_us = float_of_int st.sc.tick_us in
  let total_work_us = ref 0. in
  let durations = ref [] in
  List.iter (fun k ->
    if k.batch_size > 0 then begin
      st.kernel_id <- st.kernel_id + 1;
      let effective_gpus = max 1. (float_of_int gpus *. fclamp 0.05 1.0 k.occupancy) in
      let compute_us = div0 k.flops_t (effective_gpus *. t.gpu_tflops_per_gpu) *. 1_000_000. in
      let hbm_bytes_per_us = float_of_int gpus *. t.gpu_hbm_gbps_per_gpu *. 125. in
      let mem_us = div0 (float_of_int k.hbm_bytes) hbm_bytes_per_us in
      let queue_us = min 20_000. st.gpu_queue_us in
      let launch_us = 7.5 in
      let duration_us = launch_us +. max compute_us mem_us +. roce_wait_us +. queue_us in
      total_work_us := !total_work_us +. duration_us *. fclamp 0.01 1.0 k.occupancy;
      durations := duration_us :: !durations;
      push_window st.gpu_kernel_samples 512 duration_us;
      add_line st.gpu_lines
        (sprintf "%d,%d,%s,%d,%d,%d,%.6f,%d,%d,%.4f,%.6f,%.6f,%.6f,%.6f,%.6f"
           st.now_us st.kernel_id k.kernel_type k.batch_size gpus node_count k.flops_t k.hbm_bytes k.rdma_bytes
           k.occupancy compute_us mem_us roce_wait_us queue_us duration_us);
      emit_sst st "scheduler" "gpu-executor" Sst.GpuKernelSubmit k.rdma_bytes
        (sprintf "%s batch=%d duration_us=%.3f" k.kernel_type k.batch_size duration_us)
    end) kernels;
  let gpu_util = div0 !total_work_us (float_of_int gpus *. tick_us) in
  st.gpu_queue_us <- min 20_000. (max 0. (st.gpu_queue_us +. (!total_work_us /. float_of_int gpus) -. tick_us));
  push_window st.gpu_util_samples 512 gpu_util;
  { gpu_util; gpu_p99_us = percentile 0.99 !durations; kernel_count = List.length kernels }

(* -------------------------------------------------------------------------- *)
(* OpenZFS transaction group model                                             *)
(* -------------------------------------------------------------------------- *)

let txg_lag_us st =
  let open_lag = float_of_int (st.now_us - st.current_txg.open_us) in
  let syncing_lag =
    st.syncing_txgs
    |> List.map (fun txg -> float_of_int (st.now_us - txg.open_us))
    |> percentile 0.99
  in
  max open_lag syncing_lag

let emit_txg st txg event estimated_lag_us =
  let opt_int = function None -> "" | Some x -> string_of_int x in
  add_line st.txg_lines
    (sprintf "%d,%d,%s,%s,%.6f,%d,%d,%s,%s,%s,%d,%.6f,%d,%.6f"
       st.now_us txg.id event (string_of_txg_phase txg.phase) txg.dirty_mb txg.blocks txg.open_us
       (opt_int txg.quiesce_us) (opt_int txg.sync_start_us) (opt_int txg.sync_done_us)
       st.sc.topology.storage_nodes st.sc.topology.zfs_write_gbps_per_node st.sc.topology.zfs_sync_threads_per_node estimated_lag_us);
  emit_sst st "zfs-pool" "zfs-pool" Sst.ZfsTxgStateTransition 0
    (sprintf "txg=%d event=%s phase=%s dirty_mb=%.3f" txg.id event (string_of_txg_phase txg.phase) txg.dirty_mb)

let start_txg_sync st reason =
  let t = st.sc.topology in
  let txg = st.current_txg in
  if txg.dirty_mb > 0. then begin
    txg.phase <- Quiescing;
    txg.quiesce_us <- Some st.now_us;
    emit_txg st txg ("QUIESCE_" ^ reason) (txg_lag_us st);
    txg.phase <- Syncing;
    txg.sync_start_us <- Some (st.now_us + 250);
    let write_gbps_total = float_of_int t.storage_nodes *. t.zfs_write_gbps_per_node in
    let sync_parallelism = sqrt (float_of_int (max 1 (t.zfs_sync_threads_per_node * t.storage_nodes))) in
    let media_us = div0 (txg.dirty_mb *. 8. *. 1000.) write_gbps_total in
    let checksum_us = txg.dirty_mb *. 0.75 /. sync_parallelism in
    let flush_us = 850. +. float_of_int (txg.blocks mod 37) *. 6.0 in
    let dpu_limit_gbps = float_of_int t.dpu_paths *. t.dpu_gbps_per_path in
    let dpu_penalty = if write_gbps_total > dpu_limit_gbps then (write_gbps_total /. dpu_limit_gbps -. 1.) *. 1000. else 0. in
    let completion = st.now_us + 250 + int_of_float (ceil (media_us +. checksum_us +. flush_us +. dpu_penalty)) in
    txg.sync_done_us <- Some completion;
    emit_txg st txg "SYNC_START" (txg_lag_us st);
    st.syncing_txgs <- st.syncing_txgs @ [txg];
    st.current_txg <- new_txg st.next_txg_id st.now_us;
    st.next_txg_id <- st.next_txg_id + 1;
    emit_txg st st.current_txg "OPEN" 0.
  end

let update_zfs_txg st dirty_mb blocks =
  let txg = st.current_txg in
  if dirty_mb > 0. then begin
    txg.dirty_mb <- txg.dirty_mb +. dirty_mb;
    txg.blocks <- txg.blocks + blocks;
    emit_sst st "scheduler" "zfs-pool" Sst.ZfsTxgWrite (int_of_float (dirty_mb *. 1024. *. 1024.))
      (sprintf "txg=%d dirty_mb=%.3f blocks=%d" txg.id dirty_mb blocks)
  end;
  let age = st.now_us - txg.open_us in
  if txg.dirty_mb >= st.sc.topology.zfs_dirty_data_max_mb then start_txg_sync st "DIRTY_LIMIT"
  else if age >= st.sc.topology.zfs_txg_timeout_us && txg.dirty_mb > 0. then start_txg_sync st "TIMEOUT";
  let remaining, completed =
    List.partition (fun txg -> match txg.sync_done_us with Some done_us -> done_us > st.now_us | None -> true) st.syncing_txgs
  in
  List.iter (fun txg -> txg.phase <- Synced; emit_txg st txg "SYNC_DONE_COMMIT" 0.) completed;
  st.syncing_txgs <- remaining;
  let lag = txg_lag_us st in
  push_window st.txg_lag_samples 512 lag;
  lag

(* -------------------------------------------------------------------------- *)
(* Metadata and backpressure model                                             *)
(* -------------------------------------------------------------------------- *)

let model_metadata st ops =
  let t = st.sc.topology in
  let shards = max 1 (total_metadata_shards t) in
  let workers = max 1 (t.interactive_workers + t.restore_workers + t.bulk_workers) in
  let service_ops_per_tick = float_of_int shards *. 42.0 *. (float_of_int st.sc.tick_us /. 1000.) in
  let worker_ops_per_tick = float_of_int workers *. 16.0 *. (float_of_int st.sc.tick_us /. 1000.) in
  let cap = min service_ops_per_tick worker_ops_per_tick in
  let util = div0 (float_of_int ops) cap in
  let base_us = 130. +. 220. /. sqrt (float_of_int shards) in
  let queue_us = max 0. (util -. 0.75) *. float_of_int st.sc.tick_us *. 0.30 +. st.metadata_queue_age_us *. 0.05 in
  let p99_us = base_us +. queue_us +. float_of_int (ops mod 13) *. 3.0 in
  if util > 1.0 then st.metadata_queue_age_us <- st.metadata_queue_age_us +. (util -. 1.0) *. float_of_int st.sc.tick_us
  else st.metadata_queue_age_us <- max 0. (st.metadata_queue_age_us -. float_of_int st.sc.tick_us *. 0.5);
  push_window st.metadata_latency_samples 512 p99_us;
  push_window st.metadata_queue_samples 512 st.metadata_queue_age_us;
  emit_sst st "scheduler" "metadata-router" Sst.MetadataLookup ops
    (sprintf "ops=%d shards=%d p99_us=%.3f queue_age_us=%.3f" ops shards p99_us st.metadata_queue_age_us);
  (p99_us, st.metadata_queue_age_us)

(* -------------------------------------------------------------------------- *)
(* Scaling controller                                                          *)
(* -------------------------------------------------------------------------- *)

let before_after before after = sprintf "%s,%s" (csv_escape before) (csv_escape after)

let apply_scale st action dimension reason metrics before after =
  let (roce_p99, roce_util, gpu_p99, gpu_util, txg_lag, meta_p99, meta_queue) = metrics in
  st.scaling_actions <- st.scaling_actions + 1;
  st.last_scale_us <- st.now_us;
  add_line st.scale_lines
    (sprintf "%d,%s,%s,%s,%s,%s,%s,%s,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f"
       st.now_us (string_of_system st.sc.system) (string_of_scale_mode st.sc.scale_mode)
       action dimension (csv_escape reason) (csv_escape before) (csv_escape after)
       roce_p99 roce_util gpu_p99 gpu_util txg_lag meta_p99 meta_queue);
  emit_sst st "scaler" "scheduler" Sst.ScaleApply 0 (sprintf "%s %s: %s" dimension action reason)

let can_scale st = st.sc.scale_mode <> Scale_none && st.now_us - st.last_scale_us >= st.sc.policy.scale_cooldown_us

let scale_roce st metrics reason =
  let t = st.sc.topology in
  let before = sprintf "rails=%d gbps_per_rail=%.1f mtu=%d" t.rdma_rails t.roce_gbps_per_rail t.roce_mtu_bytes in
  begin match st.sc.scale_mode with
  | Scale_horizontal -> if t.rdma_rails < st.sc.policy.max_rdma_rails then t.rdma_rails <- t.rdma_rails + 1
  | Scale_vertical -> t.roce_gbps_per_rail <- min 800. (t.roce_gbps_per_rail *. 1.25); if t.roce_mtu_bytes < 9000 then t.roce_mtu_bytes <- 9000
  | Scale_both ->
      if t.rdma_rails < st.sc.policy.max_rdma_rails then t.rdma_rails <- t.rdma_rails + 1;
      t.roce_gbps_per_rail <- min 800. (t.roce_gbps_per_rail *. 1.15)
  | Scale_none -> () end;
  let after = sprintf "rails=%d gbps_per_rail=%.1f mtu=%d" t.rdma_rails t.roce_gbps_per_rail t.roce_mtu_bytes in
  if before <> after then apply_scale st "scale_roce_v2" "network" reason metrics before after

let scale_gpu st metrics reason =
  let t = st.sc.topology in
  let before = sprintf "gpu_nodes_per_rack=%d gpus_per_node=%d tflops_per_gpu=%.1f" t.gpu_nodes_per_rack t.gpus_per_node t.gpu_tflops_per_gpu in
  begin match st.sc.scale_mode with
  | Scale_horizontal -> if t.gpu_nodes_per_rack < st.sc.policy.max_gpu_nodes_per_rack then t.gpu_nodes_per_rack <- t.gpu_nodes_per_rack + 1
  | Scale_vertical ->
      if t.gpus_per_node < st.sc.policy.max_gpus_per_node then t.gpus_per_node <- t.gpus_per_node + 1;
      t.gpu_tflops_per_gpu <- t.gpu_tflops_per_gpu *. 1.08
  | Scale_both ->
      if t.gpu_nodes_per_rack < st.sc.policy.max_gpu_nodes_per_rack then t.gpu_nodes_per_rack <- t.gpu_nodes_per_rack + 1;
      if t.gpus_per_node < st.sc.policy.max_gpus_per_node then t.gpus_per_node <- t.gpus_per_node + 1;
      t.gpu_tflops_per_gpu <- t.gpu_tflops_per_gpu *. 1.05
  | Scale_none -> () end;
  let after = sprintf "gpu_nodes_per_rack=%d gpus_per_node=%d tflops_per_gpu=%.1f" t.gpu_nodes_per_rack t.gpus_per_node t.gpu_tflops_per_gpu in
  if before <> after then apply_scale st "scale_gpu_capacity" "compute" reason metrics before after

let scale_zfs st metrics reason =
  let t = st.sc.topology in
  let before = sprintf "storage_nodes=%d write_gbps_per_node=%.1f sync_threads_per_node=%d" t.storage_nodes t.zfs_write_gbps_per_node t.zfs_sync_threads_per_node in
  begin match st.sc.scale_mode with
  | Scale_horizontal -> if t.storage_nodes < st.sc.policy.max_storage_nodes then t.storage_nodes <- t.storage_nodes + 1
  | Scale_vertical -> t.zfs_write_gbps_per_node <- t.zfs_write_gbps_per_node *. 1.15; t.zfs_sync_threads_per_node <- t.zfs_sync_threads_per_node + 1
  | Scale_both ->
      if t.storage_nodes < st.sc.policy.max_storage_nodes then t.storage_nodes <- t.storage_nodes + 1;
      t.zfs_write_gbps_per_node <- t.zfs_write_gbps_per_node *. 1.10;
      t.zfs_sync_threads_per_node <- t.zfs_sync_threads_per_node + 1
  | Scale_none -> () end;
  let after = sprintf "storage_nodes=%d write_gbps_per_node=%.1f sync_threads_per_node=%d" t.storage_nodes t.zfs_write_gbps_per_node t.zfs_sync_threads_per_node in
  if before <> after then apply_scale st "scale_zfs_txg_writeback" "storage" reason metrics before after

let scale_metadata st metrics reason =
  let t = st.sc.topology in
  let before = sprintf "metadata_domains=%d shards_per_domain=%d api_nodes=%d orchestrators=%d workers=%d/%d/%d"
      t.metadata_domains t.metadata_shards_per_domain t.api_nodes t.orchestrators t.interactive_workers t.restore_workers t.bulk_workers in
  begin match st.sc.scale_mode with
  | Scale_horizontal ->
      if t.metadata_domains < st.sc.policy.max_metadata_domains then t.metadata_domains <- t.metadata_domains + 1;
      t.api_nodes <- t.api_nodes + 1; t.orchestrators <- t.orchestrators + 1
  | Scale_vertical ->
      if t.metadata_shards_per_domain < st.sc.policy.max_metadata_shards_per_domain then t.metadata_shards_per_domain <- t.metadata_shards_per_domain + 2;
      t.interactive_workers <- t.interactive_workers + 4
  | Scale_both ->
      if t.metadata_domains < st.sc.policy.max_metadata_domains then t.metadata_domains <- t.metadata_domains + 1;
      if t.metadata_shards_per_domain < st.sc.policy.max_metadata_shards_per_domain then t.metadata_shards_per_domain <- t.metadata_shards_per_domain + 2;
      t.api_nodes <- t.api_nodes + 1; t.orchestrators <- t.orchestrators + 1;
      t.interactive_workers <- t.interactive_workers + 4; t.restore_workers <- t.restore_workers + 2
  | Scale_none -> () end;
  let after = sprintf "metadata_domains=%d shards_per_domain=%d api_nodes=%d orchestrators=%d workers=%d/%d/%d"
      t.metadata_domains t.metadata_shards_per_domain t.api_nodes t.orchestrators t.interactive_workers t.restore_workers t.bulk_workers in
  if before <> after then apply_scale st "scale_metadata_mesh" "metadata" reason metrics before after

let evaluate_scaling st rdma gpu txg_lag metadata_p99 metadata_queue =
  if can_scale st then begin
    let roce_p99 = percentile 0.99 (queue_to_list st.roce_latency_samples) in
    let roce_util = percentile 0.99 (queue_to_list st.roce_util_samples) in
    let gpu_p99 = percentile 0.99 (queue_to_list st.gpu_kernel_samples) in
    let gpu_util = percentile 0.99 (queue_to_list st.gpu_util_samples) in
    let meta_p99 = percentile 0.99 (queue_to_list st.metadata_latency_samples) in
    let meta_queue = percentile 0.99 (queue_to_list st.metadata_queue_samples) in
    let txg_p99 = percentile 0.99 (queue_to_list st.txg_lag_samples) in
    let metrics = (roce_p99, roce_util, gpu_p99, gpu_util, txg_p99, meta_p99, meta_queue) in
    emit_sst st "scheduler" "scaler" Sst.ScaleEvaluate 0
      (sprintf "roce_p99=%.3f roce_util=%.3f gpu_p99=%.3f txg_lag=%.3f meta_p99=%.3f" roce_p99 roce_util gpu_p99 txg_p99 meta_p99);
    if roce_p99 > st.sc.policy.target_roce_p99_us || roce_util > st.sc.policy.max_roce_util then
      scale_roce st metrics (sprintf "RoCEv2 latency/util breach p99=%.3f util=%.3f target=%.3f/%.3f" roce_p99 roce_util st.sc.policy.target_roce_p99_us st.sc.policy.max_roce_util);
    if gpu_p99 > st.sc.policy.target_gpu_kernel_p99_us || gpu_util > st.sc.policy.max_gpu_util then
      scale_gpu st metrics (sprintf "GPU kernel/util breach p99=%.3f util=%.3f target=%.3f/%.3f" gpu_p99 gpu_util st.sc.policy.target_gpu_kernel_p99_us st.sc.policy.max_gpu_util);
    if txg_p99 > st.sc.policy.target_txg_lag_us then
      scale_zfs st metrics (sprintf "OpenZFS TXG lag breach p99=%.3f target=%.3f" txg_p99 st.sc.policy.target_txg_lag_us);
    if st.sc.system = USM && (meta_p99 > st.sc.policy.target_metadata_p99_us || meta_queue > st.sc.policy.max_queue_age_us) then
      scale_metadata st metrics (sprintf "USM metadata/backpressure breach p99=%.3f queue=%.3f" meta_p99 meta_queue);
    ignore (rdma, gpu, txg_lag, metadata_p99, metadata_queue)
  end

(* -------------------------------------------------------------------------- *)
(* System tick functions                                                       *)
(* -------------------------------------------------------------------------- *)

let make_kernel_specs st interactive batch restore =
  let w = st.sc.workload in
  let prompt_scale = float_of_int w.avg_prompt_tokens /. 1000. in
  let prefill_flops = float_of_int interactive *. prompt_scale *. w.gpu_flop_t_per_prompt_1k in
  let decode_flops = float_of_int interactive *. float_of_int w.avg_output_tokens *. w.gpu_flop_t_per_decode_token in
  let batch_flops = float_of_int batch *. prompt_scale *. w.gpu_flop_t_per_prompt_1k *. 1.65 in
  let restore_flops = float_of_int restore *. 0.18 in
  let mk typ batch_size flops_t hbm_mb rdma_kb occ =
    { kernel_type = typ; batch_size; flops_t; hbm_bytes = int_of_float (hbm_mb *. 1024. *. 1024.); rdma_bytes = int_of_float (rdma_kb *. 1024.); occupancy = occ }
  in
  [
    mk "prefill_attention_gemm" interactive prefill_flops (float_of_int interactive *. w.gpu_hbm_mb_per_req) (float_of_int interactive *. w.kv_hydrate_kb) 0.72;
    mk "decode_attention_kv" interactive decode_flops (float_of_int interactive *. w.gpu_hbm_mb_per_req *. 0.25) (float_of_int interactive *. w.kv_write_kb) 0.58;
    mk "batch_prefill_gemm" batch batch_flops (float_of_int batch *. w.gpu_hbm_mb_per_req *. 1.1) (float_of_int batch *. w.object_read_kb) 0.80;
    mk "checksum_compress_or_vector_index" restore restore_flops (float_of_int restore *. 32.) (float_of_int restore *. w.object_write_kb) 0.38;
  ]

let tick_pcs st =
  let w = st.sc.workload in
  let tick_s = float_of_int st.sc.tick_us /. 1_000_000. in
  let interactive = poisson_like (w.interactive_rps *. tick_s) in
  let batch = poisson_like (w.batch_rps *. tick_s) in
  let restore = poisson_like (w.restore_rps *. tick_s) in
  let requests = interactive + batch + restore in
  st.total_requests <- st.total_requests + requests;
  emit_sst st "workload" "scheduler" Sst.WorkloadArrival requests (sprintf "pcs interactive=%d batch=%d restore=%d" interactive batch restore);
  let admissions_capacity = max 1 (total_gpus st.sc.topology / 2 + st.sc.topology.interactive_workers) in
  let admitted = min requests admissions_capacity in
  let rejected = max 0 (requests - admitted) in
  st.admitted <- st.admitted + admitted;
  st.rejected <- st.rejected + rejected;
  emit_sst st "scheduler" "backpressure" Sst.AdmissionDecision admitted (sprintf "admitted=%d rejected=%d" admitted rejected);
  let kv_bytes = int_of_float (float_of_int admitted *. w.kv_hydrate_kb *. 1024.) in
  let kv_write_bytes = int_of_float (float_of_int admitted *. w.kv_write_kb *. 1024.) in
  let model_load_bytes = int_of_float (float_of_int batch *. w.object_read_kb *. 1024.) in
  let bulk_bytes = int_of_float (w.bulk_write_mbps *. 1_000_000. /. 8. *. tick_s) in
  let collective_bytes = if total_gpu_nodes st.sc.topology > 1 then int_of_float (float_of_int admitted *. 64. *. 1024.) else 0 in
  let flows = [
    { traffic_class = "kv_hydration"; opcode = "RDMA_READ"; qp = 1001; payload_bytes = kv_bytes };
    { traffic_class = "kv_persist"; opcode = "RDMA_WRITE"; qp = 1002; payload_bytes = kv_write_bytes };
    { traffic_class = "model_object_load"; opcode = "RDMA_READ"; qp = 1003; payload_bytes = model_load_bytes };
    { traffic_class = "zfs_nvmeof_write"; opcode = "RDMA_WRITE"; qp = 1004; payload_bytes = bulk_bytes };
    { traffic_class = "gpu_collective"; opcode = "SEND"; qp = 1005; payload_bytes = collective_bytes };
  ] in
  let rdma = model_rdma st flows in
  let kernels = make_kernel_specs st interactive batch restore in
  let gpu = model_gpu_kernels st kernels rdma.roce_p99_us in
  let dirty_mb = float_of_int (kv_write_bytes + bulk_bytes) /. (1024. *. 1024.) in
  let txg_lag = update_zfs_txg st dirty_mb (ceil_div (kv_write_bytes + bulk_bytes) 131072) in
  let metadata_p99, metadata_queue = model_metadata st (poisson_like (w.metadata_ops_rps *. tick_s) + admitted / 2) in
  evaluate_scaling st rdma gpu txg_lag metadata_p99 metadata_queue;
  (requests, admitted, rejected, rdma, gpu, txg_lag, metadata_p99, metadata_queue)

let tick_usm st =
  let w = st.sc.workload in
  let tick_s = float_of_int st.sc.tick_us /. 1_000_000. in
  let interactive = poisson_like (w.interactive_rps *. tick_s) in
  let batch = poisson_like (w.batch_rps *. tick_s) in
  let restore = poisson_like (w.restore_rps *. tick_s) in
  let metadata_ops = poisson_like (w.metadata_ops_rps *. tick_s) + interactive * 2 + batch in
  let requests = interactive + batch + restore in
  st.total_requests <- st.total_requests + requests;
  emit_sst st "workload" "scheduler" Sst.WorkloadArrival requests (sprintf "usm interactive=%d batch=%d restore=%d metadata=%d" interactive batch restore metadata_ops);
  let metadata_p99, metadata_queue = model_metadata st metadata_ops in
  let receiver_pressure = metadata_queue > st.sc.policy.max_queue_age_us in
  let admission_capacity = st.sc.topology.api_nodes * 8 + st.sc.topology.interactive_workers in
  let admitted = if receiver_pressure then min requests (admission_capacity / 2) else min requests admission_capacity in
  let rejected = max 0 (requests - admitted) in
  st.admitted <- st.admitted + admitted;
  st.rejected <- st.rejected + rejected;
  if receiver_pressure then emit_sst st "backpressure" "scheduler" Sst.BackpressureSignal rejected "receiver pressure: demote bulk/restore";
  let object_read_bytes = int_of_float (float_of_int interactive *. w.object_read_kb *. 1024.) in
  let restore_bytes = int_of_float (float_of_int restore *. w.object_read_kb *. 4096.) in
  let object_write_bytes = int_of_float (float_of_int (interactive + batch) *. w.object_write_kb *. 1024.) in
  let bulk_bytes = int_of_float (w.bulk_write_mbps *. 1_000_000. /. 8. *. tick_s) in
  let bulk_bytes = if receiver_pressure then bulk_bytes / 3 else bulk_bytes in
  let flows = [
    { traffic_class = "interactive_object_read"; opcode = "RDMA_READ"; qp = 2001; payload_bytes = object_read_bytes };
    { traffic_class = "restore_stream"; opcode = "RDMA_READ"; qp = 2002; payload_bytes = restore_bytes };
    { traffic_class = "metadata_journal"; opcode = "RDMA_WRITE"; qp = 2003; payload_bytes = metadata_ops * 512 };
    { traffic_class = "bulk_transfer"; opcode = "RDMA_WRITE"; qp = 2004; payload_bytes = bulk_bytes };
    { traffic_class = "zfs_nvmeof_write"; opcode = "RDMA_WRITE"; qp = 2005; payload_bytes = object_write_bytes };
  ] in
  let rdma = model_rdma st flows in
  let kernels = make_kernel_specs st (interactive / 2) (batch / 2) restore in
  let gpu = model_gpu_kernels st kernels rdma.roce_p99_us in
  let dirty_mb = float_of_int (object_write_bytes + bulk_bytes / 2) /. (1024. *. 1024.) in
  let txg_lag = update_zfs_txg st dirty_mb (ceil_div (object_write_bytes + bulk_bytes / 2) 131072) in
  evaluate_scaling st rdma gpu txg_lag metadata_p99 metadata_queue;
  (requests, admitted, rejected, rdma, gpu, txg_lag, metadata_p99, metadata_queue)

let append_timeline st requests admitted rejected rdma gpu txg_lag metadata_p99 metadata_queue =
  let t = st.sc.topology in
  add_line st.timeline_lines
    (sprintf "%d,%s,%d,%d,%d,%d,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%d,%d,%d,%d,%.6f,%d,%d,%d"
       st.now_us (string_of_system st.sc.system) requests admitted rejected rdma.packets rdma.roce_util rdma.roce_p99_us
       gpu.gpu_util gpu.gpu_p99_us st.current_txg.dirty_mb txg_lag metadata_p99 metadata_queue
       t.pods (total_gpu_nodes t) (total_gpus t) t.rdma_rails t.roce_gbps_per_rail t.storage_nodes t.metadata_domains (total_metadata_shards t))

let run_sim sc out_dir =
  ensure_dir out_dir;
  let st = new_state sc in
  emit_txg st st.current_txg "OPEN" 0.;
  while st.now_us <= sc.duration_us do
    let requests, admitted, rejected, rdma, gpu, txg_lag, meta_p99, meta_queue =
      match sc.system with PCS -> tick_pcs st | USM -> tick_usm st
    in
    append_timeline st requests admitted rejected rdma gpu txg_lag meta_p99 meta_queue;
    st.now_us <- st.now_us + sc.tick_us
  done;
  (* Force close any final open TXG so output always contains a complete transition. *)
  if st.current_txg.dirty_mb > 0. then start_txg_sync st "FINAL_FLUSH";
  st.now_us <- st.now_us + 10_000_000;
  let remaining = st.syncing_txgs in
  st.syncing_txgs <- [];
  List.iter (fun txg -> txg.phase <- Synced; emit_txg st txg "FINAL_SYNC_DONE_COMMIT" 0.) remaining;
  let rev = List.rev in
  write_lines (out_dir // "timeline.csv") (rev !(st.timeline_lines));
  write_lines (out_dir // "rdma_packets.csv") (rev !(st.rdma_lines));
  write_lines (out_dir // "openzfs_txg.csv") (rev !(st.txg_lines));
  write_lines (out_dir // "gpu_kernels.csv") (rev !(st.gpu_lines));
  write_lines (out_dir // "sst_events.csv") (rev !(st.sst_lines));
  write_lines (out_dir // "sst_components.csv") (rev !(st.component_lines));
  write_lines (out_dir // "scaling_decisions.csv") (rev !(st.scale_lines));
  let t = sc.topology in
  let summary = `Assoc [
    "scenario", `String sc.name;
    "system", `String (string_of_system sc.system);
    "duration_us", `Int sc.duration_us;
    "tick_us", `Int sc.tick_us;
    "seed", `Int sc.seed;
    "scale_mode", `String (string_of_scale_mode sc.scale_mode);
    "requests_total", `Int st.total_requests;
    "admitted_total", `Int st.admitted;
    "rejected_total", `Int st.rejected;
    "rdma_packets_total", `Int st.total_packets;
    "roce_p99_us", `Float (percentile 0.99 (queue_to_list st.roce_latency_samples));
    "roce_util_p99", `Float (percentile 0.99 (queue_to_list st.roce_util_samples));
    "ecn_marks_total", `Int st.total_ecn_marks;
    "cnp_packets_total", `Int st.total_cnp;
    "pfc_pause_us_total", `Float st.total_pfc_pause_us;
    "gpu_kernel_p99_us", `Float (percentile 0.99 (queue_to_list st.gpu_kernel_samples));
    "gpu_util_p99", `Float (percentile 0.99 (queue_to_list st.gpu_util_samples));
    "txg_lag_p99_us", `Float (percentile 0.99 (queue_to_list st.txg_lag_samples));
    "metadata_p99_us", `Float (percentile 0.99 (queue_to_list st.metadata_latency_samples));
    "metadata_queue_age_p99_us", `Float (percentile 0.99 (queue_to_list st.metadata_queue_samples));
    "scaling_actions", `Int st.scaling_actions;
    "final_topology", `Assoc [
      "pods", `Int t.pods;
      "gpu_nodes", `Int (total_gpu_nodes t);
      "gpus", `Int (total_gpus t);
      "rdma_rails", `Int t.rdma_rails;
      "roce_gbps_per_rail", `Float t.roce_gbps_per_rail;
      "storage_nodes", `Int t.storage_nodes;
      "zfs_write_gbps_per_node", `Float t.zfs_write_gbps_per_node;
      "zfs_sync_threads_per_node", `Int t.zfs_sync_threads_per_node;
      "metadata_domains", `Int t.metadata_domains;
      "metadata_shards", `Int (total_metadata_shards t);
      "api_nodes", `Int t.api_nodes;
      "orchestrators", `Int t.orchestrators
    ]
  ] in
  json_to_file (out_dir // "summary.json") summary;
  printf "%s simulation complete: %s\n%!" (String.uppercase_ascii (string_of_system sc.system)) (out_dir // "summary.json")

(* -------------------------------------------------------------------------- *)
(* CLI                                                                         *)
(* -------------------------------------------------------------------------- *)

let () =
  let system = ref None in
  let scenario_path = ref "" in
  let out_dir = ref "out/run" in
  let anon s = match !system with None -> system := Some (system_of_string s) | Some _ -> raise (Arg.Bad ("unexpected argument: " ^ s)) in
  let speclist = [
    "--scenario", Arg.Set_string scenario_path, "Path to scenario JSON";
    "--out", Arg.Set_string out_dir, "Output directory";
  ] in
  Arg.parse speclist anon "storage-sim <pcs|usm> --scenario scenario.json --out out_dir";
  let sys = match !system with Some s -> s | None -> prerr_endline "missing system: pcs|usm"; exit 2 in
  if !scenario_path = "" then (prerr_endline "missing --scenario"; exit 2);
  let sc = parse_scenario sys !scenario_path in
  if sc.system <> sys then eprintf "warning: CLI system differs from scenario system; using scenario=%s\n%!" (string_of_system sc.system);
  run_sim sc !out_dir
