(*---------------------------------------------------------------------------
  Copyright (c) 2025 Anil Madhavapeddy <anil@recoil.org>. All rights reserved.
  SPDX-License-Identifier: ISC
 ---------------------------------------------------------------------------*)

(** IMAP Command Serialization

    Serializes IMAP commands to the wire format using Eio.Buf_write. *)

module W = Eio.Buf_write

(** {1 Abstract Type} *)

type t = W.t
(** A command writer backed by Eio.Buf_write. *)

let pp ppf _ = Fmt.string ppf "<Imap.Write.t>"
let to_string _ = "<Imap.Write.t>"

(** {1 Low-level Writers} *)

let sp w = W.char w ' '
let crlf w = W.string w "\r\n"

let is_atom_char = function
  | '(' | ')' | '{' | ' ' | '\x00' .. '\x1f' | '\x7f' | '%' | '*' | '"' | '\\'
  | ']' ->
      false
  | c -> c >= '\x21' && c <= '\x7e'

let is_atom s = String.length s > 0 && String.for_all is_atom_char s

let atom w s = W.string w s

let quoted_string w s =
  W.char w '"';
  String.iter
    (fun c ->
      match c with
      | '"' | '\\' ->
          W.char w '\\';
          W.char w c
      | _ -> W.char w c)
    s;
  W.char w '"'

let literal w s =
  W.char w '{';
  W.string w (string_of_int (String.length s));
  W.string w "}\r\n";
  W.string w s

let literal_plus w s =
  W.char w '{';
  W.string w (string_of_int (String.length s));
  W.string w "+}\r\n";
  W.string w s

let needs_literal s =
  String.exists (function '\r' | '\n' | '\x00' -> true | _ -> false) s

let astring w s =
  if is_atom s then atom w s
  else if needs_literal s then literal w s
  else quoted_string w s

let nstring w = function
  | None -> W.string w "NIL"
  | Some s -> if needs_literal s then literal w s else quoted_string w s

let number w n = W.string w (string_of_int n)
let number32 w n = W.string w (Int32.to_string n)
let number64 w n = W.string w (Int64.to_string n)

(** {1 Sequence Sets} *)

let sequence_range w = function
  | Seq.Single n -> number w n
  | Seq.Range (a, b) ->
      number w a;
      W.char w ':';
      number w b
  | Seq.From n ->
      number w n;
      W.string w ":*"
  | Seq.All -> W.char w '*'

let sequence_set w set =
  List.iteri
    (fun i r ->
      if i > 0 then W.char w ',';
      sequence_range w r)
    set

(** {1 Flags} *)

let system_flag w = function
  | `Seen -> W.string w "\\Seen"
  | `Answered -> W.string w "\\Answered"
  | `Flagged -> W.string w "\\Flagged"
  | `Deleted -> W.string w "\\Deleted"
  | `Draft -> W.string w "\\Draft"

let flag w = function
  | Flag.System f -> system_flag w f
  | Flag.Keyword k -> W.string w (Mail_flag.Keyword.to_imap_string k)

let flag_list w flags =
  W.char w '(';
  List.iteri
    (fun i f ->
      if i > 0 then sp w;
      flag w f)
    flags;
  W.char w ')'

(** {1 Search Return Options (RFC 4731 ESEARCH)} *)

let search_return_opt w = function
  | Command.Return_min -> W.string w "MIN"
  | Command.Return_max -> W.string w "MAX"
  | Command.Return_all -> W.string w "ALL"
  | Command.Return_count -> W.string w "COUNT"

let search_return_opts w opts =
  W.string w "RETURN (";
  List.iteri (fun i opt ->
    if i > 0 then sp w;
    search_return_opt w opt
  ) opts;
  W.char w ')'

(** {1 Search Keys} *)

let rec search_key w = function
  | Search.All -> W.string w "ALL"
  | Search.Answered -> W.string w "ANSWERED"
  | Search.Bcc s ->
      W.string w "BCC ";
      astring w s
  | Search.Before s ->
      W.string w "BEFORE ";
      atom w s
  | Search.Body s ->
      W.string w "BODY ";
      astring w s
  | Search.Cc s ->
      W.string w "CC ";
      astring w s
  | Search.Deleted -> W.string w "DELETED"
  | Search.Flagged -> W.string w "FLAGGED"
  | Search.From s ->
      W.string w "FROM ";
      astring w s
  | Search.Keyword s ->
      W.string w "KEYWORD ";
      atom w s
  | Search.New -> W.string w "NEW"
  | Search.Not k ->
      W.string w "NOT ";
      search_key w k
  | Search.Old -> W.string w "OLD"
  | Search.On s ->
      W.string w "ON ";
      atom w s
  | Search.Or (k1, k2) ->
      W.string w "OR ";
      search_key w k1;
      sp w;
      search_key w k2
  | Search.Seen -> W.string w "SEEN"
  | Search.Since s ->
      W.string w "SINCE ";
      atom w s
  | Search.Subject s ->
      W.string w "SUBJECT ";
      astring w s
  | Search.Text s ->
      W.string w "TEXT ";
      astring w s
  | Search.To s ->
      W.string w "TO ";
      astring w s
  | Search.Unanswered -> W.string w "UNANSWERED"
  | Search.Undeleted -> W.string w "UNDELETED"
  | Search.Unflagged -> W.string w "UNFLAGGED"
  | Search.Unkeyword s ->
      W.string w "UNKEYWORD ";
      atom w s
  | Search.Unseen -> W.string w "UNSEEN"
  | Search.Draft -> W.string w "DRAFT"
  | Search.Undraft -> W.string w "UNDRAFT"
  | Search.Header (field, value) ->
      W.string w "HEADER ";
      astring w field;
      sp w;
      astring w value
  | Search.Larger n ->
      W.string w "LARGER ";
      number64 w n
  | Search.Smaller n ->
      W.string w "SMALLER ";
      number64 w n
  | Search.Uid set ->
      W.string w "UID ";
      sequence_set w set
  | Search.Sequence_set set -> sequence_set w set
  | Search.And keys ->
      W.char w '(';
      List.iteri
        (fun i k ->
          if i > 0 then sp w;
          search_key w k)
        keys;
      W.char w ')'
  | Search.Sentbefore s ->
      W.string w "SENTBEFORE ";
      atom w s
  | Search.Senton s ->
      W.string w "SENTON ";
      atom w s
  | Search.Sentsince s ->
      W.string w "SENTSINCE ";
      atom w s

(** {1 Fetch Items} *)

let write_partial w = function
  | Some (offset, len) ->
      W.char w '<';
      number w offset;
      W.char w '.';
      number w len;
      W.char w '>'
  | None -> ()

let fetch_item w = function
  | Fetch.Envelope -> W.string w "ENVELOPE"
  | Fetch.Flags -> W.string w "FLAGS"
  | Fetch.Internaldate -> W.string w "INTERNALDATE"
  | Fetch.Rfc822 -> W.string w "RFC822"
  | Fetch.Rfc822_size -> W.string w "RFC822.SIZE"
  | Fetch.Rfc822_header -> W.string w "RFC822.HEADER"
  | Fetch.Rfc822_text -> W.string w "RFC822.TEXT"
  | Fetch.Uid -> W.string w "UID"
  | Fetch.Body -> W.string w "BODY"
  | Fetch.Bodystructure -> W.string w "BODYSTRUCTURE"
  | Fetch.Body_section (section, partial) ->
      W.string w "BODY["; W.string w section; W.char w ']';
      write_partial w partial
  | Fetch.Body_peek (section, partial) ->
      W.string w "BODY.PEEK["; W.string w section; W.char w ']';
      write_partial w partial
  | Fetch.Binary (section, partial) ->
      W.string w "BINARY["; W.string w section; W.char w ']';
      write_partial w partial
  | Fetch.Binary_peek (section, partial) ->
      W.string w "BINARY.PEEK["; W.string w section; W.char w ']';
      write_partial w partial
  | Fetch.Binary_size section ->
      W.string w "BINARY.SIZE["; W.string w section; W.char w ']'
  | Fetch.Modseq ->
      (* RFC 7162 Section 3.1.5: MODSEQ fetch data item *)
      W.string w "MODSEQ"

let fetch_items w = function
  | [ item ] -> fetch_item w item
  | items ->
      W.char w '(';
      List.iteri
        (fun i item ->
          if i > 0 then sp w;
          fetch_item w item)
        items;
      W.char w ')'

(** {1 Status Items} *)

let status_item w = function
  | Status.Messages -> W.string w "MESSAGES"
  | Status.Uidnext -> W.string w "UIDNEXT"
  | Status.Uidvalidity -> W.string w "UIDVALIDITY"
  | Status.Unseen -> W.string w "UNSEEN"
  | Status.Deleted -> W.string w "DELETED"
  | Status.Size -> W.string w "SIZE"
  | Status.Highestmodseq -> W.string w "HIGHESTMODSEQ"  (* RFC 7162 CONDSTORE *)

let status_items w items =
  W.char w '(';
  List.iteri
    (fun i item ->
      if i > 0 then sp w;
      status_item w item)
    items;
  W.char w ')'

(** {1 Store Actions} *)

let store_action w = function
  | Store.Set -> W.string w "FLAGS"
  | Store.Add -> W.string w "+FLAGS"
  | Store.Remove -> W.string w "-FLAGS"

(** {1 Sort Criteria} *)

let sort_key w = function
  | Sort.Arrival -> W.string w "ARRIVAL"
  | Sort.Cc -> W.string w "CC"
  | Sort.Date -> W.string w "DATE"
  | Sort.From -> W.string w "FROM"
  | Sort.Size -> W.string w "SIZE"
  | Sort.Subject -> W.string w "SUBJECT"
  | Sort.To -> W.string w "TO"

let sort_criterion w c =
  if c.Sort.reverse then W.string w "REVERSE ";
  sort_key w c.Sort.key

let sort_criteria w criteria =
  W.char w '(';
  List.iteri
    (fun i c ->
      if i > 0 then sp w;
      sort_criterion w c)
    criteria;
  W.char w ')'

(** {1 Thread Algorithm} *)

let thread_algorithm w = function
  | Thread.Orderedsubject -> W.string w "ORDEREDSUBJECT"
  | Thread.References -> W.string w "REFERENCES"
  | Thread.Extension s -> W.string w (String.uppercase_ascii s)

(** {1 ID Parameters} *)

let id_params w = function
  | None -> W.string w "NIL"
  | Some pairs ->
      W.char w '(';
      List.iteri
        (fun i (k, v) ->
          if i > 0 then sp w;
          quoted_string w k;
          sp w;
          quoted_string w v)
        pairs;
      W.char w ')'

(** {1 Commands} *)

let write_search w charset criteria return_opts =
  W.string w "SEARCH";
  Option.iter (fun opts -> sp w; search_return_opts w opts) return_opts;
  Option.iter (fun cs -> W.string w " CHARSET "; astring w cs) charset;
  sp w;
  search_key w criteria

let write_sort w charset criteria search =
  W.string w "SORT ";
  sort_criteria w criteria;
  sp w;
  astring w charset;
  sp w;
  search_key w search

let write_thread w algorithm charset search =
  W.string w "THREAD ";
  thread_algorithm w algorithm;
  sp w;
  astring w charset;
  sp w;
  search_key w search

let command_body w = function
  | Command.Capability -> W.string w "CAPABILITY"
  | Command.Noop -> W.string w "NOOP"
  | Command.Logout -> W.string w "LOGOUT"
  | Command.Starttls -> W.string w "STARTTLS"
  | Command.Login { username; password } ->
      W.string w "LOGIN ";
      astring w username;
      sp w;
      astring w password
  | Command.Authenticate { mechanism; initial_response } ->
      W.string w "AUTHENTICATE ";
      atom w mechanism;
      Option.iter (fun r -> sp w; W.string w r) initial_response
  | Command.Enable caps ->
      W.string w "ENABLE";
      List.iter (fun c -> sp w; atom w c) caps
  | Command.Select mailbox ->
      W.string w "SELECT ";
      astring w mailbox
  | Command.Examine mailbox ->
      W.string w "EXAMINE ";
      astring w mailbox
  | Command.Create mailbox ->
      W.string w "CREATE ";
      astring w mailbox
  | Command.Delete mailbox ->
      W.string w "DELETE ";
      astring w mailbox
  | Command.Rename { old_name; new_name } ->
      W.string w "RENAME ";
      astring w old_name;
      sp w;
      astring w new_name
  | Command.Subscribe mailbox ->
      W.string w "SUBSCRIBE ";
      astring w mailbox
  | Command.Unsubscribe mailbox ->
      W.string w "UNSUBSCRIBE ";
      astring w mailbox
  | Command.List { reference; pattern } ->
      W.string w "LIST ";
      astring w reference;
      sp w;
      astring w pattern
  | Command.Namespace -> W.string w "NAMESPACE"
  | Command.Status { mailbox; items } ->
      W.string w "STATUS ";
      astring w mailbox;
      sp w;
      status_items w items
  | Command.Append { mailbox; flags; date; message } ->
      W.string w "APPEND ";
      astring w mailbox;
      (match flags with
      | [] -> ()
      | flags -> sp w; flag_list w flags);
      Option.iter (fun d -> sp w; quoted_string w d) date;
      sp w;
      (* Use LITERAL+ to avoid synchronization issues *)
      literal_plus w message
  | Command.Idle -> W.string w "IDLE"
  | Command.Close -> W.string w "CLOSE"
  | Command.Unselect -> W.string w "UNSELECT"
  | Command.Expunge -> W.string w "EXPUNGE"
  | Command.Search { charset; criteria; return_opts } ->
      write_search w charset criteria return_opts
  | Command.Sort { charset; criteria; search } ->
      write_sort w charset criteria search
  | Command.Thread { algorithm; charset; search } ->
      write_thread w algorithm charset search
  | Command.Fetch { sequence; items; changedsince } ->
      W.string w "FETCH ";
      sequence_set w sequence;
      sp w;
      fetch_items w items;
      Option.iter (fun modseq ->
        W.string w " (CHANGEDSINCE ";
        number64 w modseq;
        W.char w ')') changedsince
  | Command.Store { sequence; silent; action; flags; unchangedsince } ->
      W.string w "STORE ";
      sequence_set w sequence;
      sp w;
      (match unchangedsince with
       | Some modseq ->
           W.string w "(UNCHANGEDSINCE ";
           number64 w modseq;
           W.string w ") "
       | None -> ());
      store_action w action;
      if silent then W.string w ".SILENT";
      sp w;
      flag_list w flags
  | Command.Copy { sequence; mailbox } ->
      W.string w "COPY ";
      sequence_set w sequence;
      sp w;
      astring w mailbox
  | Command.Move { sequence; mailbox } ->
      W.string w "MOVE ";
      sequence_set w sequence;
      sp w;
      astring w mailbox
  | Command.Uid cmd -> (
      W.string w "UID ";
      match cmd with
      | Command.Uid_fetch { sequence; items; changedsince } ->
          W.string w "FETCH ";
          sequence_set w sequence;
          sp w;
          fetch_items w items;
          Option.iter (fun modseq ->
            W.string w " (CHANGEDSINCE ";
            number64 w modseq;
            W.char w ')') changedsince
      | Command.Uid_store { sequence; silent; action; flags; unchangedsince } ->
          W.string w "STORE ";
          sequence_set w sequence;
          sp w;
          (match unchangedsince with
           | Some modseq ->
               W.string w "(UNCHANGEDSINCE ";
               number64 w modseq;
               W.string w ") "
           | None -> ());
          store_action w action;
          if silent then W.string w ".SILENT";
          sp w;
          flag_list w flags
      | Command.Uid_copy { sequence; mailbox } ->
          W.string w "COPY ";
          sequence_set w sequence;
          sp w;
          astring w mailbox
      | Command.Uid_move { sequence; mailbox } ->
          W.string w "MOVE ";
          sequence_set w sequence;
          sp w;
          astring w mailbox
      | Command.Uid_search { charset; criteria; return_opts } ->
          write_search w charset criteria return_opts
      | Command.Uid_sort { charset; criteria; search } ->
          write_sort w charset criteria search
      | Command.Uid_thread { algorithm; charset; search } ->
          write_thread w algorithm charset search
      | Command.Uid_expunge set ->
          W.string w "EXPUNGE ";
          sequence_set w set)
  | Command.Id params ->
      W.string w "ID ";
      id_params w params

let command w ~tag cmd =
  atom w tag;
  sp w;
  command_body w cmd;
  crlf w

let idle_done w =
  W.string w "DONE";
  crlf w

let authenticate_response w data =
  W.string w data;
  crlf w