open Types

type prefix =
  | No_prefix
  | Server_prefix of string
  | User_prefix of name * userinfo option

let prefix_string = function
  | No_prefix -> ""
  | Server_prefix servername -> servername
  | User_prefix (nick, None) -> nick
  | User_prefix (nick, Some uinfo) -> Format.asprintf "%s%a" nick pp_userinfo uinfo

let pp_prefix ppf p =
  Format.fprintf ppf "%S" (prefix_string p)

type t = {
  prefix : prefix;
  command : string;
  params : string list;
  trailing : bool;
} [@@deriving show { with_path = false }]

let is_param_trailing p =
  String.starts_with p ~prefix:":" || String.contains p ' '

let rec is_params_trailing = function
  | [] -> false
  | [p] -> is_param_trailing p
  | _ :: tl -> is_params_trailing tl

let make ?(prefix = No_prefix) ?(always_trailing = false) command params =
  let trailing = always_trailing || is_params_trailing params in
  { prefix; command; params; trailing }

let write buf t =
  if t.prefix <> No_prefix then begin
    Buffer.add_char buf ':';
    Buffer.add_string buf (prefix_string t.prefix);
    Buffer.add_char buf ' ';
  end;
  Buffer.add_string buf t.command;
  let rec add_params = function
    | [] -> ()
    | [p] when t.trailing ->
      Buffer.add_string buf " :";
      Buffer.add_string buf p
    | p::tl ->
      Buffer.add_char buf ' ';
      Buffer.add_string buf p;
      add_params tl
  in
  add_params t.params;
  Buffer.add_string buf "\r\n"

exception Empty_message of int
exception Incomplete

let crlf = function '\r' | '\n' | '\x00' -> true | _ -> false
let cl = function ':' -> true | _ -> false
let sp = function ' ' | '\t' -> true | _ -> false
let nocrlf c = not (crlf c)
let nospcrlf c = not (sp c || crlf c)

let startswith str f i =
  if i >= String.length str then
    raise_notrace Incomplete;
  f str.[i]

let rec skip str f i =
  if i < String.length str && f str.[i] then
    skip str f (i + 1)
  else
    i

let span str f i =
  let e = skip str f i in
  e, String.sub str i (e - i)

let parse str =
  let parse_prefix i =
    if startswith str cl i then
      let i, pfx = span str nospcrlf (i + 1) in
      (* TODO: parse [User_prefix] *)
      (* not really needed because prefixes from clients are always ignored *)
      i, Server_prefix pfx
    else
      i, No_prefix
  in

  let parse_command i =
    let i = skip str sp i in
    if startswith str crlf i then
      raise_notrace (Empty_message (skip str crlf i));
    let i, cmd = span str nospcrlf i in
    i, String.uppercase_ascii cmd
  in

  let rec parse_params acc i =
    let i = skip str sp i in
    if startswith str crlf i then
      let i = skip str crlf i in
      i, List.rev acc, false
    else if startswith str cl i then
      let i, par = span str nocrlf (i + 1) in
      i, List.rev (par :: acc), true
    else
      let i, par = span str nospcrlf i in
      parse_params (par :: acc) i
  in

  let parse_msg i =
    let i, prefix = parse_prefix i in
    let i, command = parse_command i in
    let i, params, always_trailing = parse_params [] i in
    let msg = make command params ~prefix ~always_trailing in
    i, msg
  in

  let rec parse_all acc i =
    match parse_msg i with
    | i, msg -> parse_all (msg :: acc) i
    | exception Empty_message i -> parse_all acc i
    | exception Incomplete -> i, List.rev acc
  in

  let i, msgs = parse_all [] 0 in
  let rest = if i = 0 then str else String.sub str i (String.length str - i) in
  msgs, rest

;;

let%expect_test _ =
  let print_msg_nl m =
    let buf = Buffer.create 64 in
    write buf m;
    Printf.printf "%S\n" (Buffer.contents buf)
  in
  let print_parsed_msgs (msgs, rest) =
    List.iter
      (fun m -> Format.kasprintf print_string "%a\n" pp m)
      msgs;
    Format.printf "%S\n" rest
  in

  Format.printf "%a" pp_prefix No_prefix;
  [%expect {| "" |}];

  Format.printf "%a" pp_prefix (Server_prefix "localhost");
  [%expect {| "localhost" |}];

  Format.printf "%a" pp_prefix (User_prefix ("tali", None));
  [%expect {| "tali" |}];

  let milo = {
    username = "milo";
    realname = "Milo";
    hostname = "elsewhere";
  } in
  Format.printf "%a" pp_prefix (User_prefix ("tali", Some milo));
  [%expect {| "tali!milo@elsewhere" |}];

  make "NICK" ["tali"] |> print_msg_nl;
  [%expect {| "NICK tali\r\n" |}];

  make "NICK" [":)"] |> print_msg_nl;
  [%expect {| "NICK ::)\r\n" |}];

  make "USER" ["milo"; "0"; "*"; "milo"] ~always_trailing:true |> print_msg_nl;
  [%expect {| "USER milo 0 * :milo\r\n" |}];

  let lo = Server_prefix "localhost" in
  make "001" ["tali"; "Welcome to the IRC Network"] ~prefix:lo |> print_msg_nl;
  [%expect {| ":localhost 001 tali :Welcome to the IRC Network\r\n" |}];

  parse ":server.com  usEr  tali  0  *  iitalics\r\n" |> print_parsed_msgs;
  [%expect {|
    { prefix = "server.com"; command = "USER";
      params = ["tali"; "0"; "*"; "iitalics"]; trailing = false }
    ""
  |}];

  parse "PRIVMSG #lol :Hello world\r\nPRIVMSG #lol :meow\r\nPRIVM" |> print_parsed_msgs;
  [%expect {|
    { prefix = ""; command = "PRIVMSG"; params = ["#lol"; "Hello world"];
      trailing = true }
    { prefix = ""; command = "PRIVMSG"; params = ["#lol"; "meow"];
      trailing = true }
    "PRIVM"
  |}];

  parse "    \r\n" |> print_parsed_msgs;
  [%expect {| "" |}];

  parse ":ignore-me \r\nFOO\r\n" |> print_parsed_msgs;
  [%expect {|
    { prefix = ""; command = "FOO"; params = []; trailing = false }
    ""
  |}];