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

let rec is_params_trailing = function
  | [] -> false
  | [tr] -> String.contains tr ' '
  | _ :: tl -> is_params_trailing tl

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

let write buf t =
  Option.iter
    (fun pre ->
       Buffer.add_char buf ':';
       Buffer.add_string buf pre;
       Buffer.add_char buf ' ')
    t.prefix;
  Buffer.add_string buf t.command;
  let rec add_params = function
    | [] -> ()
    | [tr] when t.trailing ->
      Buffer.add_string buf " :";
      Buffer.add_string buf tr
    | hd::tl ->
      Buffer.add_char buf ' ';
      Buffer.add_string buf hd;
      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' -> true | _ -> false
let cl = function ':' -> true | _ -> false
let sp = function ' ' -> true | _ -> false
let nocrlf = function '\r' | '\n' -> false | _ -> true
let nospcrlf = function ' ' | '\r' | '\n' -> false | _ -> true

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
      i, Some pfx
    else
      i, None
  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

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

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

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

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

  parse "PRIVMSG #lol :Hello world\r\nPRIVMSG #lol :meow\r\nPRIVM" |> print_parsed_msgs;
  [%expect {|
    { prefix = None; command = "PRIVMSG"; params = ["#lol"; "Hello world"];
      trailing = true }
    { prefix = None; 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 = None; command = "FOO"; params = []; trailing = false }
    ""
  |}];