(** Test solid pill lifecycle with all three Nock implementations using Domainslib **)

open Nock_lib

module T = Domainslib.Task

let time_it name f =
  Printf.printf "  [%s] Starting on domain %d...\n%!" name (Domain.self () :> int);
  let start = Unix.gettimeofday () in
  try
    let result = f () in
    let elapsed = Unix.gettimeofday () -. start in
    Printf.printf "  [%s] ✓ Complete in %.4fs (domain %d)\n%!" name elapsed (Domain.self () :> int);
    Ok (name, elapsed, result)
  with
  | Stack_overflow ->
      let elapsed = Unix.gettimeofday () -. start in
      Printf.printf "  [%s] ✗ Stack overflow after %.4fs (domain %d)\n%!" name elapsed (Domain.self () :> int);
      Error (name, "Stack overflow")
  | e ->
      let elapsed = Unix.gettimeofday () -. start in
      Printf.printf "  [%s] ✗ Error: %s (after %.4fs, domain %d)\n%!" name (Printexc.to_string e) elapsed (Domain.self () :> int);
      Error (name, Printexc.to_string e)

let main () =
  Printf.printf "\n╔═══════════════════════════════════════════════════════╗\n";
  Printf.printf "║  Parallel Solid Pill Test (Domainslib)               ║\n";
  Printf.printf "╚═══════════════════════════════════════════════════════╝\n\n";

  Printf.printf "[1] Loading solid.pill...\n";
  let ic = open_in_bin "solid.pill" in
  let len = in_channel_length ic in
  let bytes = Bytes.create len in
  really_input ic bytes 0 len;
  close_in ic;
  Printf.printf "    Size: %d bytes (%.1f MB)\n" len (float_of_int len /. 1024.0 /. 1024.0);

  Printf.printf "[2] Cuing solid pill...\n";
  let start_time = Unix.gettimeofday () in
  let pill = Serial.cue bytes in
  let elapsed = Unix.gettimeofday () -. start_time in
  Printf.printf "    ✓ Cued in %.2fs\n" elapsed;

  Printf.printf "[3] Parsing pill structure...\n";

  (* Parse: [tag type [bot mod use]] *)
  match pill with
  | Noun.Cell { h = _tag; t = rest; _ } ->
      (match rest with
       | Noun.Cell { h = _ptype; t = events_triple; _ } ->
           (match events_triple with
            | Noun.Cell { h = bot; t = rest2; _ } ->
                (match rest2 with
                 | Noun.Cell { h = _mod; t = rest3; _ } ->
                     (match rest3 with
                      | Noun.Cell { h = use; t = _; _ } ->
                          (* Concatenate bot and use *)
                          let rec to_list acc n =
                            match n with
                            | Noun.Atom _ -> List.rev acc
                            | Noun.Cell { h; t; _ } -> to_list (h :: acc) t
                          in
                          let bot_list = to_list [] bot in
                          let use_list = to_list [] use in
                          let all_events = bot_list @ use_list in

                          Printf.printf "    Bot events: %d\n" (List.length bot_list);
                          Printf.printf "    Use events: %d\n" (List.length use_list);
                          Printf.printf "    Total: %d events\n" (List.length all_events);

                          (* Convert to proper Nock list format *)
                          let rec from_list = function
                            | [] -> Noun.atom 0
                            | h :: t -> Noun.cell h (from_list t)
                          in
                          let event_list = from_list all_events in

                          Printf.printf "\n[4] Building lifecycle formula...\n";
                          let formula = Noun.cell
                            (Noun.atom 2)
                            (Noun.cell
                              (Noun.cell (Noun.atom 0) (Noun.atom 3))
                              (Noun.cell (Noun.atom 0) (Noun.atom 2))) in
                          Printf.printf "    Formula: [2 [0 3] [0 2]]\n";

                          Printf.printf "\n[5] Setting up Domainslib pool with 3 domains...\n";
                          let num_domains = 18 in
                          let pool = T.setup_pool ~num_domains:(num_domains - 1) () in

                          Printf.printf "\n[6] Running lifecycle with all three implementations:\n\n";

                          (* Run in T.run context for effect handlers *)
                          let res1, res2, res3 =
                            T.run pool (fun () ->
                              (* Launch all three as async tasks *)
                              let task1 = T.async pool (fun () ->
                                time_it "nock.ml" (fun () -> Nock.nock_on event_list formula)) in

                              let task2 = T.async pool (fun () ->
                                time_it "nock_iter.ml" (fun () -> Nock_iter.nock_on event_list formula)) in

                              let task3 = T.async pool (fun () ->
                                time_it "nock_tail.ml" (fun () -> Nock_tail.nock_on event_list formula)) in

                              (* Await all results *)
                              Printf.printf "\n[7] Waiting for results...\n%!";
                              let r1 = T.await pool task1 in
                              let r2 = T.await pool task2 in
                              let r3 = T.await pool task3 in
                              (r1, r2, r3)
                            )
                          in

                          T.teardown_pool pool;

                          Printf.printf "\n[8] Results:\n";

                          let successes = ref [] in
                          let failures = ref [] in

                          List.iter (fun res ->
                            match res with
                            | Ok (name, time, result) ->
                                let mug = Noun.mug result in
                                Printf.printf "  %s: 0x%08lx (%.4fs) ✓\n" name mug time;
                                successes := (name, mug, time) :: !successes
                            | Error (name, err) ->
                                Printf.printf "  %s: %s ✗\n" name err;
                                failures := (name, err) :: !failures
                          ) [res1; res2; res3];

                          Printf.printf "\n[9] Verification:\n";
                          Printf.printf "  Successes: %d\n" (List.length !successes);
                          Printf.printf "  Failures: %d\n" (List.length !failures);

                          if List.length !successes >= 2 then begin
                            let mugs = List.map (fun (_, mug, _) -> mug) !successes in
                            let all_same = List.for_all (fun m -> m = List.hd mugs) mugs in
                            if all_same then
                              Printf.printf "  ✓ All successful implementations produce identical kernels!\n"
                            else
                              Printf.printf "  ✗ Mug mismatch between successful implementations!\n"
                          end;

                          Printf.printf "\n╔═══════════════════════════════════════════════════════╗\n";
                          Printf.printf "║  Summary:                                             ║\n";
                          if List.length !successes > 0 then begin
                            let fastest = List.fold_left
                              (fun (acc_name, acc_time) (name, _, time) ->
                                if time < acc_time then (name, time) else (acc_name, acc_time))
                              ("", max_float)
                              !successes in
                            Printf.printf "║  Fastest (of successful): %-28s║\n" (fst fastest);
                            Printf.printf "║  Time: %.2fs                                          ║\n" (snd fastest);
                          end;
                          Printf.printf "╚═══════════════════════════════════════════════════════╝\n\n"

                      | _ -> Printf.printf "  ✗ Unexpected structure at use level\n")
                 | _ -> Printf.printf "  ✗ Unexpected structure at mod level\n")
            | _ -> Printf.printf "  ✗ Unexpected structure at bot level\n")
       | _ -> Printf.printf "  ✗ Unexpected structure at events level\n")
  | _ -> Printf.printf "  ✗ Unexpected pill structure\n"

let _ = main ()