(* 创建并启动一个新 domain *)
let () =
  let handler = Domain.spawn (fun () ->
    Printf.printf "Hello from domain %d!\n" (Domain.self () :> int);
    42
  ) in
  let result = Domain.join handler in
  Printf.printf "Result: %d\n" result

(* 创建原子值 *)
let counter = Atomic.make 0

(* 原子操作 *)
let () =
  Atomic.incr counter;                    (* 原子递增 *)
  Atomic.decr counter;                    (* 原子递减 *)
  let old = Atomic.fetch_and_add counter 5 in  (* 原子加法，返回旧值 *)
  Printf.printf "Old: %d, New: %d\n" old (Atomic.get counter);
  Atomic.set counter 100;                 (* 原子设置 *)
  let v = Atomic.get counter in           (* 原子读取 *)
  Printf.printf "Current: %d\n" v

(* CAS（Compare-And-Swap）操作 *)
let rec increment counter =
  let old = Atomic.get counter in
  let new_val = old + 1 in
  if Atomic.compare_and_set counter old new_val then
    new_val  (* 成功 *)
  else
    increment counter  (* 失败，重试 *)

(* 带 CAS 的无锁栈 *)
type 'a stack = 'a list Atomic.t

let create_stack () : 'a stack = Atomic.make []

let push stack item =
  let rec loop () =
    let old = Atomic.get stack in
    if Atomic.compare_and_set stack old (item :: old) then ()
    else loop ()
  in
  loop ()

let pop stack =
  let rec loop () =
    match Atomic.get stack with
    | [] -> None
    | hd :: tl as old ->
      if Atomic.compare_and_set stack old tl then Some hd
      else loop ()
  in
  loop ()

(* 使用 Mutex 保护共享数据 *)
let shared_counter = ref 0
let mutex = Mutex.create ()

let increment () =
  Mutex.lock mutex;
  shared_counter := !shared_counter + 1;
  Mutex.unlock mutex

let get_counter () =
  Mutex.lock mutex;
  let v = !shared_counter in
  Mutex.unlock mutex;
  v

(* 更安全的使用方式：with_lock *)
let safe_increment () =
  Mutex.protect mutex (fun () ->
    shared_counter := !shared_counter + 1
  )

(* 并行递增示例 *)
let () =
  let domains = Array.init 4 (fun _ ->
    Domain.spawn (fun () ->
      for _ = 1 to 10000 do
        safe_increment ()
      done
    )
  ) in
  Array.iter Domain.join domains;
  Printf.printf "Final count: %d\n" (get_counter ())
  (* 输出: Final count: 40000 *)

(* 生产者-消费者队列 *)
type 'a bounded_queue = {
  mutex : Mutex.t;
  not_empty : Condition.t;
  not_full : Condition.t;
  queue : 'a Queue.t;
  capacity : int;
}

let create_queue capacity = {
  mutex = Mutex.create ();
  not_empty = Condition.create ();
  not_full = Condition.create ();
  queue = Queue.create ();
  capacity;
}

let put q item =
  Mutex.protect q.mutex (fun () ->
    while Queue.length q.queue >= q.capacity do
      Condition.wait q.not_full q.mutex
    done;
    Queue.push item q.queue;
    Condition.signal q.not_empty
  )

let take q =
  Mutex.protect q.mutex (fun () ->
    while Queue.is_empty q.queue do
      Condition.wait q.not_empty q.mutex
    done;
    let item = Queue.pop q.queue in
    Condition.signal q.not_full;
    item
  )

(* 使用示例 *)
let () =
  let q = create_queue 10 in
  let producer = Domain.spawn (fun () ->
    for i = 1 to 20 do
      put q i;
      Printf.printf "Produced: %d\n" i
    done
  ) in
  let consumer = Domain.spawn (fun () ->
    for _ = 1 to 20 do
      let v = take q in
      Printf.printf "Consumed: %d\n" v
    done
  ) in
  Domain.join producer;
  Domain.join consumer

(* 为每个 domain 设置 GC 参数 *)
let setup_gc_for_domain () =
  Gc.set { (Gc.get ()) with
    minor_heap_size = 2 * 1024 * 1024;  (* 2M words *)
    space_overhead = 200;                 (* 允许更多空间 *)
  }

(* 创建带自定义 GC 的 domain *)
let d = Domain.spawn (fun () ->
  setup_gc_for_domain ();
  compute_heavy_task ()
)

let measure f =
  let start = Unix.gettimeofday () in
  let result = f () in
  let elapsed = Unix.gettimeofday () -. start in
  (result, elapsed)

let sum_parallel arr n_domains =
  let len = Array.length arr in
  let chunk_size = len / n_domains in
  let partials = Array.init n_domains (fun i ->
    Domain.spawn (fun () ->
      let start = i * chunk_size in
      let end_ = if i = n_domains - 1 then len else start + chunk_size in
      let sum = ref 0 in
      for j = start to end_ - 1 do
        sum := !sum + arr.(j)
      done; !sum)
  ) in
  Array.fold_left (fun acc d -> acc + Domain.join d) 0 partials

(* 简单 fork-join *)
let par_map f arr =
  let n = Array.length arr in
  let mid = n / 2 in
  let left = Array.sub arr 0 mid in
  let right = Array.sub arr mid (n - mid) in
  let d1 = Domain.spawn (fun () -> Array.map f left) in
  let r2 = Array.map f right in
  Array.append (Domain.join d1) r2

let parallel_for n_domains arr f =
  let len = Array.length arr in
  let chunk = len / n_domains in
  Array.init n_domains (fun i ->
    Domain.spawn (fun () ->
      let lo = i * chunk in
      let hi = if i = n_domains - 1 then len else lo + chunk in
      for j = lo to hi - 1 do arr.(j) <- f arr.(j) done))
  |> Array.iter Domain.join

(* 简单 channel 实现 *)
type 'a channel = { mutex: Mutex.t; cond: Condition.t; queue: 'a Queue.t }

let create_channel () =
  { mutex = Mutex.create (); cond = Condition.create (); queue = Queue.create () }

let send ch msg =
  Mutex.protect ch.mutex (fun () ->
    Queue.push msg ch.queue; Condition.signal ch.cond)

let receive ch =
  Mutex.protect ch.mutex (fun () ->
    while Queue.is_empty ch.queue do
      Condition.wait ch.cond ch.mutex
    done;
    Queue.pop ch.queue)

opam install eio

open Eio.Std

let main env =
  let pool = Eio.Stdenv.domain_mgr env in
  Eio.Domain_manager.run pool (fun () ->
    traceln "Running on domain %d" (Domain.self () :> int))

let () = Eio_main.run main

let process_pipeline data =
  let filtered = par_filter is_valid data in
  let d1 = Domain.spawn (fun () -> process_type_a filtered) in
  let r2 = process_type_b filtered in
  let result_a = Domain.join d1 in
  merge result_a r2

open Lwt.Syntax

(* 将 CPU 密集型任务放到独立 domain *)
let offload_cpu f =
  Lwt_preemptive.detach f ()

(* 或手动使用 Domain *)
let offload_to_domain f =
  let promise, resolver = Lwt.wait () in
  let _ = Domain.spawn (fun () ->
    try
      let result = f () in
      Lwt.wakeup_later resolver result
    with e ->
      Lwt.wakeup_later_exn resolver e
  ) in
  promise

(* 使用 *)
let process_image image_data =
  let* result = offload_to_domain (fun () ->
    heavy_computation image_data
  ) in
  Lwt.return result

let run_server n_domains port =
  let sock = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
  Unix.setsockopt sock Unix.SO_REUSEADDR true;
  Unix.bind sock Unix.(ADDR_INET (inet_addr_any, port));
  Unix.listen sock 128;
  Array.init n_domains (fun _ ->
    Domain.spawn (fun () ->
      while true do
        let client = Unix.accept sock in
        serve_client client
      done)
  ) |> Array.iter Domain.join