open Swim.Types

let gen_node_id : node_id QCheck.Gen.t =
  let open QCheck.Gen in
  let+ id =
    oneof_weighted
      [
        (3, string_size ~gen:printable (int_range 1 64));
        (1, return "");
        (1, string_size ~gen:printable (return 255));
      ]
  in
  node_id_of_string id

let gen_incarnation : incarnation QCheck.Gen.t =
  let open QCheck.Gen in
  let+ i =
    oneof_weighted
      [ (5, int_range 0 1000); (2, int_range 0 max_int); (1, return 0) ]
  in
  incarnation_of_int i

let gen_member_state : member_state QCheck.Gen.t =
  let open QCheck.Gen in
  let alive : member_state = Alive in
  let suspect : member_state = Suspect in
  let dead : member_state = Dead in
  oneof [ return alive; return suspect; return dead ]

let gen_ipv4 : string QCheck.Gen.t =
  let open QCheck.Gen in
  let+ a = int_range 0 255
  and+ b = int_range 0 255
  and+ c = int_range 0 255
  and+ d = int_range 0 255 in
  Printf.sprintf "%d.%d.%d.%d" a b c d

let gen_port : int QCheck.Gen.t =
  let open QCheck.Gen in
  oneof_weighted
    [ (3, int_range 1024 65535); (1, int_range 1 1023); (1, return 7946) ]

let gen_addr : addr QCheck.Gen.t =
  let open QCheck.Gen in
  let+ ip = gen_ipv4 and+ port = gen_port in
  let ipaddr =
    match Ipaddr.V4.of_string ip with
    | Ok v4 -> v4
    | Error _ -> Ipaddr.V4.localhost
  in
  `Udp (Eio.Net.Ipaddr.of_raw (Ipaddr.V4.to_octets ipaddr), port)

let gen_meta : string QCheck.Gen.t =
  let open QCheck.Gen in
  oneof_weighted
    [
      (3, string_size ~gen:printable (int_range 0 256));
      (1, return "");
      (1, return (String.make 1024 'x'));
    ]

let gen_node_info : node_info QCheck.Gen.t =
  let open QCheck.Gen in
  let+ id = gen_node_id and+ addr = gen_addr and+ meta = gen_meta in
  make_node_info ~id ~addr ~meta

let gen_seq : int QCheck.Gen.t =
  let open QCheck.Gen in
  oneof_weighted
    [ (5, int_range 0 10000); (2, int_range 0 max_int); (1, return 0) ]

let gen_ping : protocol_msg QCheck.Gen.t =
  let open QCheck.Gen in
  let+ seq = gen_seq and+ target = gen_node_id and+ sender = gen_node_info in
  Ping { seq; target; sender }

let gen_ping_req : protocol_msg QCheck.Gen.t =
  let open QCheck.Gen in
  let+ seq = gen_seq and+ target = gen_node_id and+ sender = gen_node_info in
  Ping_req { seq; target; sender }

let gen_payload : string option QCheck.Gen.t =
  let open QCheck.Gen in
  oneof_weighted
    [
      (2, return None);
      (3, map Option.some (string_size ~gen:printable (int_range 0 512)));
    ]

let gen_ack : protocol_msg QCheck.Gen.t =
  let open QCheck.Gen in
  let+ seq = gen_seq
  and+ responder = gen_node_info
  and+ payload = gen_payload in
  Ack { seq; responder; payload }

let gen_alive : protocol_msg QCheck.Gen.t =
  let open QCheck.Gen in
  let+ node = gen_node_info and+ incarnation = gen_incarnation in
  Alive { node; incarnation }

let gen_suspect : protocol_msg QCheck.Gen.t =
  let open QCheck.Gen in
  let+ node = gen_node_id
  and+ incarnation = gen_incarnation
  and+ suspector = gen_node_id in
  Suspect { node; incarnation; suspector }

let gen_dead : protocol_msg QCheck.Gen.t =
  let open QCheck.Gen in
  let+ node = gen_node_id
  and+ incarnation = gen_incarnation
  and+ declarator = gen_node_id in
  Dead { node; incarnation; declarator }

let gen_topic : string QCheck.Gen.t =
  QCheck.Gen.string_size ~gen:QCheck.Gen.printable (QCheck.Gen.int_range 1 64)

let gen_user_payload : string QCheck.Gen.t =
  QCheck.Gen.string_size ~gen:QCheck.Gen.printable (QCheck.Gen.int_range 0 1024)

let gen_user_msg : protocol_msg QCheck.Gen.t =
  let open QCheck.Gen in
  let+ topic = gen_topic
  and+ payload = gen_user_payload
  and+ origin = gen_node_id in
  User_msg { topic; payload; origin }

let gen_protocol_msg : protocol_msg QCheck.Gen.t =
  QCheck.Gen.oneof
    [
      gen_ping;
      gen_ping_req;
      gen_ack;
      gen_alive;
      gen_suspect;
      gen_dead;
      gen_user_msg;
    ]

let gen_cluster_name : string QCheck.Gen.t =
  let open QCheck.Gen in
  oneof_weighted
    [
      (3, string_size ~gen:printable (int_range 1 32));
      (1, return "default");
      (1, return "test-cluster");
    ]

let gen_piggyback : protocol_msg list QCheck.Gen.t =
  let open QCheck.Gen in
  let piggyback_msg =
    oneof [ gen_alive; gen_suspect; gen_dead; gen_user_msg ]
  in
  list_size (int_range 0 8) piggyback_msg

let gen_packet : packet QCheck.Gen.t =
  let open QCheck.Gen in
  let+ cluster = gen_cluster_name
  and+ primary = gen_protocol_msg
  and+ piggyback = gen_piggyback in
  { cluster; primary; piggyback }

let gen_cstruct : Cstruct.t QCheck.Gen.t =
  let open QCheck.Gen in
  let+ len =
    oneof_weighted
      [ (3, int_range 0 1024); (1, return 0); (1, int_range 1024 4096) ]
  and+ fill = char in
  let cs = Cstruct.create len in
  Cstruct.memset cs (Char.code fill);
  cs

let gen_cstruct_sized (size : int) : Cstruct.t QCheck.Gen.t =
  let open QCheck.Gen in
  let+ bytes = string_size ~gen:char (return size) in
  Cstruct.of_string bytes

let gen_config : config QCheck.Gen.t =
  let open QCheck.Gen in
  let+ bind_addr = gen_ipv4
  and+ bind_port = gen_port
  and+ node_name =
    oneof_weighted [ (2, return None); (3, map Option.some gen_topic) ]
  and+ protocol_interval = float_range 0.1 10.0
  and+ probe_timeout = float_range 0.1 5.0
  and+ indirect_checks = int_range 1 10
  and+ suspicion_mult = int_range 1 10
  and+ suspicion_max_timeout = float_range 10.0 120.0
  and+ retransmit_mult = int_range 1 10
  and+ udp_buffer_size =
    oneof [ return 1400; return 1500; return 8192; return 65507 ]
  and+ tcp_timeout = float_range 1.0 30.0
  and+ send_buffer_count = int_range 4 64
  and+ recv_buffer_count = int_range 4 64
  and+ secret_key = gen_cstruct_sized 16
  and+ cluster_name = gen_cluster_name
  and+ label = oneof [ return ""; gen_topic ]
  and+ encryption_enabled = bool
  and+ gossip_verify_incoming = bool
  and+ gossip_verify_outgoing = bool
  and+ max_gossip_queue_depth = int_range 10 10000 in
  {
    bind_addr;
    bind_port;
    node_name;
    protocol_interval;
    probe_timeout;
    indirect_checks;
    suspicion_mult;
    suspicion_max_timeout;
    retransmit_mult;
    udp_buffer_size;
    tcp_timeout;
    send_buffer_count;
    recv_buffer_count;
    secret_key = Cstruct.to_string secret_key;
    cluster_name;
    label;
    encryption_enabled;
    gossip_verify_incoming;
    gossip_verify_outgoing;
    max_gossip_queue_depth;
  }

let gen_decode_error : decode_error QCheck.Gen.t =
  let open QCheck.Gen in
  oneof
    [
      return Invalid_magic;
      map (fun v -> Unsupported_version v) (int_range 0 255);
      return Truncated_message;
      map (fun t -> Invalid_tag t) (int_range 0 255);
      return Decryption_failed;
    ]

let gen_send_error : send_error QCheck.Gen.t =
  let open QCheck.Gen in
  oneof [ return Node_unreachable; return Timeout; return Connection_reset ]

let gen_mtime_span : Mtime.span QCheck.Gen.t =
  let open QCheck.Gen in
  let+ ns = map Int64.of_int (int_range 0 1_000_000_000) in
  Mtime.Span.of_uint64_ns ns

let gen_member_snapshot : member_snapshot QCheck.Gen.t =
  let open QCheck.Gen in
  let+ node = gen_node_info
  and+ state = gen_member_state
  and+ incarnation = gen_incarnation
  and+ state_change = gen_mtime_span in
  { node; state; incarnation; state_change }

let arb_node_id : node_id QCheck.arbitrary =
  QCheck.make ~print:(fun id -> node_id_to_string id) gen_node_id

let arb_incarnation : incarnation QCheck.arbitrary =
  QCheck.make
    ~print:(fun inc -> string_of_int (incarnation_to_int inc))
    ~shrink:(fun inc ->
      let i = incarnation_to_int inc in
      QCheck.Shrink.int i |> QCheck.Iter.map incarnation_of_int)
    gen_incarnation

let arb_member_state : member_state QCheck.arbitrary =
  QCheck.make ~print:member_state_to_string gen_member_state

let format_addr (addr : addr) : string =
  match addr with
  | `Udp (ip, port) -> Fmt.str "%a:%d" Eio.Net.Ipaddr.pp ip port
  | `Unix path -> Printf.sprintf "unix:%s" path

let format_node_info (ni : node_info) : string =
  Printf.sprintf "{ id=%s; addr=%s; meta=%S }" (node_id_to_string ni.id)
    (format_addr ni.addr) ni.meta

let arb_node_info : node_info QCheck.arbitrary =
  QCheck.make ~print:format_node_info gen_node_info

let format_protocol_msg (msg : protocol_msg) : string =
  match msg with
  | Ping { seq; target; sender } ->
      Printf.sprintf "Ping { seq=%d; target=%s; sender=%s }" seq
        (node_id_to_string target) (format_node_info sender)
  | Ping_req { seq; target; sender } ->
      Printf.sprintf "Ping_req { seq=%d; target=%s; sender=%s }" seq
        (node_id_to_string target) (format_node_info sender)
  | Ack { seq; responder; payload } ->
      Printf.sprintf "Ack { seq=%d; responder=%s; payload=%s }" seq
        (format_node_info responder)
        (match payload with
        | None -> "None"
        | Some p -> Printf.sprintf "Some %S" p)
  | Alive { node; incarnation } ->
      Printf.sprintf "Alive { node=%s; incarnation=%d }" (format_node_info node)
        (incarnation_to_int incarnation)
  | Suspect { node; incarnation; suspector } ->
      Printf.sprintf "Suspect { node=%s; incarnation=%d; suspector=%s }"
        (node_id_to_string node)
        (incarnation_to_int incarnation)
        (node_id_to_string suspector)
  | Dead { node; incarnation; declarator } ->
      Printf.sprintf "Dead { node=%s; incarnation=%d; declarator=%s }"
        (node_id_to_string node)
        (incarnation_to_int incarnation)
        (node_id_to_string declarator)
  | User_msg { topic; payload; origin } ->
      Printf.sprintf "User_msg { topic=%S; payload=%S; origin=%s }" topic
        payload (node_id_to_string origin)

let arb_protocol_msg : protocol_msg QCheck.arbitrary =
  QCheck.make ~print:format_protocol_msg gen_protocol_msg

let format_packet (p : packet) : string =
  Printf.sprintf "{ cluster=%S; primary=%s; piggyback=[%d msgs] }" p.cluster
    (format_protocol_msg p.primary)
    (List.length p.piggyback)

let arb_packet : packet QCheck.arbitrary =
  QCheck.make ~print:format_packet gen_packet

let arb_cstruct : Cstruct.t QCheck.arbitrary =
  QCheck.make
    ~print:(fun cs -> Printf.sprintf "<cstruct len=%d>" (Cstruct.length cs))
    gen_cstruct

let arb_decode_error : decode_error QCheck.arbitrary =
  QCheck.make ~print:decode_error_to_string gen_decode_error

let arb_send_error : send_error QCheck.arbitrary =
  QCheck.make ~print:send_error_to_string gen_send_error

let arb_member_snapshot : member_snapshot QCheck.arbitrary =
  QCheck.make
    ~print:(fun ms ->
      Printf.sprintf "{ node=%s; state=%s; incarnation=%d }"
        (format_node_info ms.node)
        (member_state_to_string ms.state)
        (incarnation_to_int ms.incarnation))
    gen_member_snapshot