open Types

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 of_list l =
    List.fold_left (fun s m -> add m s) empty l

  let to_string s =
    Seq.filter_map
      (fun m -> if mem m s then Some (to_char m) else None)
      (List.to_seq [`i; `m; `n; `o; `s; `t; `w]) |>
    String.of_seq

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

  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 change dir m chg =
    match dir with
    | `add -> { add = add m chg.add; rem = remove m chg.rem }
    | `rem -> { add = remove m chg.add; rem = add m chg.rem }

  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 Unknown_mode of char
  exception Missing_args

  let parse_mode_flags f str =
    let rec loop dir acc i =
      if i >= String.length str then List.rev acc
      else match str.[i] with
        | '+' -> loop `add acc (i + 1)
        | '-' -> loop `rem acc (i + 1)
        | ch ->
          let mode = try f ch
            with Invalid_argument _ ->
              raise (Unknown_mode ch)
          in
          loop dir ((dir, mode) :: acc) (i + 1)
    in
    loop `add [] 0

  type user_modes = Set.change

  let no_user_mdoes = Set.no_change

  let user_modes str =
    let parse ms (dir, m) = Set.change dir m ms in
    List.fold_left
      parse
      no_user_mdoes
      (parse_mode_flags of_char_user str)

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

  type add_or_rem = [
    | `add
    | `rem
  ]

  type priv = [`o | `v]

  type chan_modes = {
    chan_modes : Set.change;
    chan_key : string set_or_unset option;
    chan_limit : int set_or_unset option;
    chan_privs : (add_or_rem * priv * name) list;
  }

  let no_chan_modes = {
    chan_modes = Set.no_change;
    chan_key = None;
    chan_limit = None;
    chan_privs = []
  }

  let chan_modes str args =
    let parse (acc, args) (dir, m) = match m, dir, args with
      (* Type A: Modes that add or remove an address to or from a list. These modes MUST
         always have a parameter when sent from the server to a client. A client MAY issue
         this type of mode without an argument to obtain the current contents of the
         list. *)
      | #chan_a as m, _dir, args ->
        begin match m, args with
          | `b, [] -> (* TODO: show list *) acc, args
          | `b, _ :: args -> (* TODO: add/rem from list *) acc, args
        end

      (* Type B: Modes that change a setting on a channel. These modes MUST always have a
         parameter. *)
      | #chan_b as m, dir, args ->
        let arg, args = match args with
          | [] -> raise Missing_args
          | arg :: args -> arg, args
        in
        begin
          try match m, dir, arg with
            | `k, `add, k ->
              if k = "" then invalid_arg "empty key";
              { acc with chan_key = Some (`set k) }, args
            | `k, `rem, _k ->
              { acc with chan_key = Some `unset }, args
            | (`o | `v) as priv, dir, nick ->
              let chan_privs = (dir, priv, nick) :: acc.chan_privs in
              { acc with chan_privs }, args
          with Invalid_argument _ ->
            (* ignore invalid args (+k) *)
            acc, args
        end

      (* Type C: Modes that change a setting on a channel. These modes MUST have a
         parameter when being set, and MUST NOT have a parameter when being unset. *)
      | #chan_c as m, dir, args ->
        let arg, args = match dir, args with
          | `add, [] -> raise Missing_args
          | `add, arg :: args -> `set arg, args
          | `rem, args -> `unset, args
        in
        begin
          try match m, arg with
            | `l, `set n ->
              let n = Option.get (int_of_string_opt n) in
              if n <= 0 then invalid_arg "limit <= 0";
              { acc with chan_limit = Some (`set n) }, args
            | `l, `unset ->
              { acc with chan_limit = Some `unset }, args
          with Invalid_argument _ ->
            (* ignore invalid args (+l) *)
            acc, args
        end

      (* Type D: Modes that change a setting on a channel. These modes MUST NOT have a
         parameter. *)
      | #chan_d as m, dir, args ->
        let chan_modes = Set.change dir m acc.chan_modes in
        { acc with chan_modes }, args
    in        

    let modes, _args = 
      List.fold_left
        parse
        (no_chan_modes, args)
        (parse_mode_flags of_char_chan str)
    in
    (* chan_privs is built in reverse order *)
    { modes with chan_privs = List.rev modes.chan_privs }
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_list [`i]); [%expect {| [i] |}];
  print_set_nl Set.(of_list [`n;`o]); [%expect {| [no] |}];
  print_set_nl Set.(of_list [`s;`m]); [%expect {| [ms] |}];
  print_set_nl Set.(of_list [`w;`i]); [%expect {| [iw] |}];
  print_bool_nl Set.(mem `i (of_list [`i;`n;`s])); [%expect "true"];
  print_bool_nl Set.(mem `w (of_list [`i;`n;`s])); [%expect "false"];
  print_bool_nl Set.(mem `w (of_list [`w;`w;`w;`w])); [%expect "true"];
  print_bool_nl Set.(mem `t (of_list [`i;`m;`n;`o;`s;`w])); [%expect "false"];

  let print_parse_error f =
    try f () |> ignore; print_endline "()"
    with Parse.Unknown_mode c -> Printf.printf "unknown mode %c\n" c
       | Parse.Missing_args -> Printf.printf "missing args\n"
  in

  print_change_nl (Parse.user_modes "+iw"); [%expect {| +iw |}];
  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 |}];
  print_parse_error (fun () -> Parse.user_modes "+I"); [%expect {| unknown mode I |}];

  let m = Set.of_list [`i;`w] in
  let m, c = Set.normalize m (Parse.user_modes "-w+io") in
  Format.printf "%a -> [%a]\n@." 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;
    List.iter
      (fun (add_rem, mode, nick) ->
         Format.printf " %c%c:%s"
           (match add_rem with `add -> '+' | `rem -> '-')
           (to_char mode)
           nick)
      m.chan_privs;
    Format.printf "]\n@."
  in

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

  print_chan_modes (Parse.chan_modes "+ls" ["100"]);
  [%expect {| [+s +l:100] |}];

  print_chan_modes (Parse.chan_modes "+ls" ["-1"]);
  [%expect {| [+s] |}];

  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"] |}];

  print_parse_error (fun () -> Parse.chan_modes "+k" []);
  [%expect {| missing args |}];
  print_parse_error (fun () -> Parse.chan_modes "+l" []);
  [%expect {| missing args |}];

  print_chan_modes (Parse.chan_modes "+o+v" ["aaa"; "bbb"]);
  [%expect {| [+ +o:aaa +v:bbb] |}];

  print_chan_modes (Parse.chan_modes "+o-v" ["aaa"; "bbb"]);
  [%expect {| [+ +o:aaa -v:bbb] |}];