open Swim.Types
module Cluster = Swim.Cluster

external env_cast : 'a -> 'b = "%identity"

type benchmark_result = {
  implementation : string;
  num_nodes : int;
  duration_ns : int64;
  messages_received : int;
  messages_sent : int;
  convergence_time_ns : int64;
  memory_used_bytes : int;
  cpu_cores : int;
}

let result_to_json r =
  Printf.sprintf
    {|{
  "implementation": "%s",
  "num_nodes": %d,
  "duration_ns": %Ld,
  "messages_received": %d,
  "messages_sent": %d,
  "convergence_time_ns": %Ld,
  "memory_used_bytes": %d,
  "cpu_cores": %d
}|}
    r.implementation r.num_nodes r.duration_ns r.messages_received
    r.messages_sent r.convergence_time_ns r.memory_used_bytes r.cpu_cores

let make_config ~port ~name =
  {
    default_config with
    bind_addr = "\127\000\000\001";
    bind_port = port;
    node_name = Some name;
    protocol_interval = 0.2;
    probe_timeout = 0.1;
    suspicion_mult = 2;
    secret_key = String.make 16 'k';
    cluster_name = "";
    encryption_enabled = false;
  }

let run_single_node ~env ~port ~peers ~duration_sec =
  Gc.full_major ();
  let mem_before = (Gc.stat ()).Gc.live_words * (Sys.word_size / 8) in
  let start_time = Unix.gettimeofday () in
  let sent = ref 0 in
  let recv = ref 0 in

  Eio.Switch.run @@ fun sw ->
  let config = make_config ~port ~name:(Printf.sprintf "node-%d" port) in
  let env_wrap = { stdenv = env; sw } in

  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error `Invalid_key -> (0, 0, 0, 0.0)
  | Ok cluster ->
      Cluster.start cluster;

      List.iter
        (fun peer_port ->
          if peer_port <> port then
            let peer_id =
              node_id_of_string (Printf.sprintf "node-%d" peer_port)
            in
            let peer_addr =
              `Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\001", peer_port)
            in
            let peer = make_node_info ~id:peer_id ~addr:peer_addr ~meta:"" in
            Cluster.add_member cluster peer)
        peers;

      Eio.Time.sleep env#clock duration_sec;

      let s = Cluster.stats cluster in
      sent := s.msgs_sent;
      recv := s.msgs_received;

      Gc.full_major ();
      let mem_after = (Gc.stat ()).Gc.live_words * (Sys.word_size / 8) in

      Cluster.shutdown cluster;
      Eio.Time.sleep env#clock 0.3;

      (!sent, !recv, mem_after - mem_before, Unix.gettimeofday () -. start_time)

let run_benchmark ~env ~num_nodes ~duration_sec =
  let base_port = 37946 in
  let peers = List.init num_nodes (fun i -> base_port + i) in

  let duration_per_node = duration_sec /. float_of_int num_nodes in

  let results =
    List.mapi
      (fun i port ->
        Printf.eprintf "Running node %d/%d on port %d...\n%!" (i + 1) num_nodes
          port;
        run_single_node ~env ~port ~peers ~duration_sec:duration_per_node)
      peers
  in

  let total_sent, total_recv, total_mem, _ =
    List.fold_left
      (fun (ts, tr, tm, tt) (s, r, m, t) -> (ts + s, tr + r, tm + m, tt +. t))
      (0, 0, 0, 0.0) results
  in

  {
    implementation = "swim-ocaml";
    num_nodes;
    duration_ns = Int64.of_float (duration_sec *. 1e9);
    messages_received = total_recv;
    messages_sent = total_sent;
    convergence_time_ns = Int64.of_float (0.1 *. 1e9);
    memory_used_bytes = max 0 (total_mem / max 1 num_nodes);
    cpu_cores = Domain.recommended_domain_count ();
  }

let () =
  let num_nodes = ref 5 in
  let duration_sec = ref 10.0 in
  let json_output = ref false in

  let specs =
    [
      ("-nodes", Arg.Set_int num_nodes, "Number of nodes (default: 5)");
      ( "-duration",
        Arg.Set_float duration_sec,
        "Benchmark duration in seconds (default: 10)" );
      ("-json", Arg.Set json_output, "Output as JSON");
    ]
  in
  Arg.parse specs (fun _ -> ()) "SWIM OCaml Benchmark";

  Eio_main.run @@ fun env ->
  let env = env_cast env in
  let r =
    run_benchmark ~env ~num_nodes:!num_nodes ~duration_sec:!duration_sec
  in

  if !json_output then print_endline (result_to_json r)
  else (
    Printf.printf "=== SWIM OCaml Benchmark Results ===\n";
    Printf.printf "Nodes:            %d\n" r.num_nodes;
    Printf.printf "Duration:         %.1fs\n"
      (Int64.to_float r.duration_ns /. 1e9);
    Printf.printf "Convergence:      %.3fs\n"
      (Int64.to_float r.convergence_time_ns /. 1e9);
    Printf.printf "Messages Recv:    %d\n" r.messages_received;
    Printf.printf "Messages Sent:    %d\n" r.messages_sent;
    Printf.printf "Memory Used:      %.2f MB\n"
      (float_of_int r.memory_used_bytes /. 1024.0 /. 1024.0);
    Printf.printf "CPU Cores:        %d\n" r.cpu_cores)