type user = [`i | `o | `w]
type chan_a = [`b]
type chan_b = [`k | `o | `v]
type chan_c = [`l]
type chan_d = [`i | `m | `s | `t | `n]
type chan = [chan_a | chan_b | chan_c | chan_d]
type t = [user | chan]

let to_string = function
  | `b -> "b"
  | `i -> "i"
  | `k -> "k"
  | `l -> "l"
  | `m -> "m"
  | `n -> "n"
  | `o -> "o"
  | `s -> "s"
  | `t -> "t"
  | `v -> "v"
  | `w -> "w"

let to_char x = (to_string x).[0]
let pp ppf x = Format.pp_print_string ppf (to_string x)

let of_char_user = function
  | 'i' -> `i
  | 'o' -> `o
  | 'w' -> `w
  | _ -> invalid_arg "Mode.of_char_user"

let of_char_chan = function
  | 'b' -> `b
  | 'i' -> `i
  | 'k' -> `k
  | 'l' -> `l
  | 'm' -> `m
  | 'n' -> `n
  | 'o' -> `o
  | 's' -> `s
  | 't' -> `t
  | 'v' -> `v
  | _ -> invalid_arg "Mode.of_char_chan"

let[@warning "+8"] _check_exhaustive (x : chan) =
  match x with
  | #chan_a -> ()
  | #chan_b -> ()
  | #chan_c -> ()
  | #chan_d -> ()


module Set = struct
  type t = int
  type elt = [user | chan_d]

  let[@inline] singleton = function
    | `i -> 1
    | `m -> 2
    | `n -> 4
    | `o -> 8
    | `s -> 16
    | `t -> 32
    | `w -> 64

  let empty = 0
  let equal = Int.equal
  let union = ( lor )
  let inter = ( land )
  let diff s t = s land Int.lognot t
  let mem elt s = s land (singleton elt) <> 0
  let add elt s = union s (singleton elt)
  let remove elt s = diff s (singleton elt)

  let to_string s =
    let bs = Bytes.create 7 in
    let bit ch elt i =
      if mem elt s then (Bytes.set bs i ch; i + 1)
      else i
    in
    let n =
      0 |> bit 'i' `i |> bit 'm' `m |> bit 'n' `n |> bit 'o' `o |>
      bit 's' `s |> bit 't' `t |> bit 'w' `w
    in
    Bytes.sub_string bs 0 n

  let pp ppf s =
    Format.pp_print_string ppf (to_string s)

  let of_string s =
    let chr = function
      | 'i' -> `i | 'm' -> `m | 'n' -> `n | 'o' -> `o
      | 's' -> `s | 't' -> `t | 'w' -> `w
      | _ -> invalid_arg "Irc.Mode.Set.of_string"
    in
    String.fold_left (fun s c -> add (chr c) s) empty s

  let of_list l =
    List.fold_left (fun s m -> add m s) empty l

  type change = {
    add : t;
    rem : t;
  }

  let pp_change ppf { add; rem } =
    if rem = empty then
      Format.fprintf ppf "+%a" pp add
    else if add = empty then
      Format.fprintf ppf "-%a" pp rem
    else
      Format.fprintf ppf "+%a-%a" pp add pp rem

  let no_change = { add = empty; rem = empty }

  let normalize t { add; rem } =
    let add = diff add t in
    let rem = inter rem t in
    diff (union t add) rem, { add; rem }
end


module Parse = struct
  exception Error

  let fail fmt =
    Format.kasprintf (fun _ -> raise Error) fmt

  let parse_mode_set str ~of_char ~add ~rem ~init =
    let rec loop dir acc i =
      if i >= String.length str then acc
      else match str.[i] with
        | '+' -> loop `add acc (i + 1)
        | '-' -> loop `rem acc (i + 1)
        | _ ->
          let mode = try of_char str.[i]
            with Invalid_argument _ ->
              fail "unrecognized mode char"
          in
          match dir with
          | `add -> loop dir (add acc mode) (i + 1)
          | `rem -> loop dir (rem acc mode) (i + 1)
          | `none -> fail "mode must start with + or -"
    in
    loop `none init 0

  let take = function
    | [] -> fail "expected argument"
    | x :: xs -> x, xs

  let take_int = function
    | [] -> fail "expected argument"
    | x :: xs -> try int_of_string x, xs
      with Invalid_argument _ -> fail "invalid integer"

  type user_modes = Set.change

  let user_modes str =
    parse_mode_set str
      ~of_char:of_char_user
      ~init:Set.{ add = empty; rem = empty }
      ~add:(fun ms m -> Set.{ add = add m ms.add; rem = remove m ms.rem })
      ~rem:(fun ms m -> Set.{ add = remove m ms.add; rem = add m ms.rem })

  type 'a set_or_unset = [
    | `set of 'a
    | `unset
  ]

  type chan_modes = {
    chan_modes : Set.change;
    chan_key : string set_or_unset option;
    chan_limit : int set_or_unset option;
  }

  let chan_modes_add (args, modes) = function
    | `k ->
      (* type B *)
      let key, args = take args in
      let chan_key = Some (`set key) in
      args, { modes with chan_key }
    | `l ->
      (* type C *)
      let limit, args = take_int args in
      let chan_limit = Some (`set limit) in
      args, { modes with chan_limit }
    | `b | `o | `v -> fail "TODO: + ban/op/voice"
    | #chan_d as m ->
      let chan_modes = {
        Set.add = Set.add m modes.chan_modes.add;
        Set.rem = Set.remove m modes.chan_modes.rem;
      } in
      args, { modes with chan_modes }

  let chan_modes_rem (args, modes) = function
    | `k ->
      (* type B *)
      let _key, args = take args in
      let chan_key = Some `unset in
      args, { modes with chan_key }
    | `l ->
      (* type C *)
      let chan_limit = Some `unset in
      args, { modes with chan_limit }
    | `b | `o | `v -> fail "TODO: - ban/op/voice"
    | #chan_d as m ->
      let chan_modes = {
        Set.add = Set.remove m modes.chan_modes.add;
        Set.rem = Set.add m modes.chan_modes.rem;
      } in
      args, { modes with chan_modes }

  let chan_modes str args =
    let modes = {
      chan_modes = Set.no_change;
      chan_key = None;
      chan_limit = None;
    } in
    let _, modes =
      parse_mode_set str
        ~of_char:of_char_chan
        ~init:(args, modes)
        ~add:chan_modes_add
        ~rem:chan_modes_rem
    in
    modes
end


let%expect_test _ =
  let roundtrip of_char_x m =
    match of_char_x (to_char m) with
    | m' when m = m' -> ()
    | m' -> Printf.printf "fail:%s!=%s\n" (to_string m) (to_string m')
    | exception Invalid_argument _ -> Printf.printf "error:%s\n" (to_string m)
  in
  let roundtrip_user m = roundtrip of_char_user m in
  let roundtrip_chan m = roundtrip of_char_chan m in

  roundtrip_user `i;
  roundtrip_user `w;
  roundtrip_chan `i;
  roundtrip_chan `k;
  roundtrip_chan `l;
  roundtrip_chan `m;
  roundtrip_chan `n;
  roundtrip_chan `o;
  roundtrip_chan `s;
  roundtrip_chan `t;

  let print_set_nl s = Printf.printf "[%s]\n" (Set.to_string s) in
  let print_bool_nl b = print_endline (if b then "true" else "false") in
  let print_change_nl c = Format.kasprintf print_string "%a\n" Set.pp_change c in

  print_set_nl Set.empty; [%expect {| [] |}];
  print_set_nl Set.(of_string "i"); [%expect {| [i] |}];
  print_set_nl Set.(of_string "no"); [%expect {| [no] |}];
  print_set_nl Set.(of_string "sm"); [%expect {| [ms] |}];
  print_set_nl Set.(of_string "wi"); [%expect {| [iw] |}];
  print_bool_nl Set.(mem `i (of_string "ins")); [%expect "true"];
  print_bool_nl Set.(mem `w (of_string "ins")); [%expect "false"];
  print_bool_nl Set.(mem `w (of_string "wwww")); [%expect "true"];
  print_bool_nl Set.(mem `t (of_string "imnosw")); [%expect "false"];

  let expect_parse_error f =
    try f () |> ignore; print_endline "no error"
    with Parse.Error -> ()
       | e -> print_endline (Printexc.to_string e)
  in

  print_change_nl (Parse.user_modes "+iw"); [%expect {| +iw |}];
  print_change_nl (Parse.user_modes "-wo"); [%expect {| -ow |}];
  print_change_nl (Parse.user_modes "+i-w"); [%expect {| +i-w |}];
  print_change_nl (Parse.user_modes "-i+w"); [%expect {| +w-i |}];
  print_change_nl (Parse.user_modes "+i-i"); [%expect {| -i |}];
  print_change_nl (Parse.user_modes "-o+o"); [%expect {| +o |}];
  expect_parse_error (fun () -> Parse.user_modes "+b");
  expect_parse_error (fun () -> Parse.user_modes "w");

  let m, c = Set.normalize (Set.of_string "iw") (Parse.user_modes "-w+io") in
  Format.printf "%a -> [%a]" Set.pp_change c Set.pp m;
  [%expect {| +o-w -> [io] |}];

  let print_chan_modes (m : Parse.chan_modes) =
    Format.printf "[%a" Set.pp_change m.chan_modes;
    begin match m.chan_limit with
      | Some (`set n) -> Format.printf " +l:%d" n
      | Some `unset -> Format.printf " -l"
      | None -> ()
    end;
    begin match m.chan_key with
      | Some (`set k) -> Format.printf " +k:%S" k
      | Some `unset -> Format.printf " -k"
      | None -> ()
    end;
    Format.printf "]\n@."
  in

  print_chan_modes (Parse.chan_modes "+im-nm+s" []);
  [%expect {| [+is-mn] |}];

  print_chan_modes (Parse.chan_modes "+l-ik+lt" ["100"; "*"; "200"]);
  [%expect {| [+t-i +l:200 -k] |}];

  print_chan_modes (Parse.chan_modes "+k-k+k" ["a"; "b"; "c"]);
  [%expect {| [+ +k:"c"] |}];