(* Kitty Graphics Protocol - Command *)

type action =
  [ `Transmit
  | `Transmit_and_display
  | `Query
  | `Display
  | `Delete
  | `Frame
  | `Animate
  | `Compose ]

type t = {
  action : action;
  format : Kgp_types.format option;
  transmission : Kgp_types.transmission option;
  compression : Kgp_types.compression option;
  width : int option;
  height : int option;
  size : int option;
  offset : int option;
  quiet : Kgp_types.quiet option;
  image_id : int option;
  image_number : int option;
  placement : Kgp_placement.t option;
  delete : Kgp_types.delete option;
  frame : Kgp_frame.t option;
  animation : Kgp_animation.t option;
  compose : Kgp_compose.t option;
}

let make action =
  {
    action;
    format = None;
    transmission = None;
    compression = None;
    width = None;
    height = None;
    size = None;
    offset = None;
    quiet = None;
    image_id = None;
    image_number = None;
    placement = None;
    delete = None;
    frame = None;
    animation = None;
    compose = None;
  }

let transmit ?image_id ?image_number ?format ?transmission ?compression ?width
    ?height ?size ?offset ?quiet () =
  {
    (make `Transmit) with
    image_id;
    image_number;
    format;
    transmission;
    compression;
    width;
    height;
    size;
    offset;
    quiet;
  }

let transmit_and_display ?image_id ?image_number ?format ?transmission
    ?compression ?width ?height ?size ?offset ?quiet ?placement () =
  {
    (make `Transmit_and_display) with
    image_id;
    image_number;
    format;
    transmission;
    compression;
    width;
    height;
    size;
    offset;
    quiet;
    placement;
  }

let query ?format ?transmission ?width ?height ?quiet () =
  { (make `Query) with format; transmission; width; height; quiet }

let display ?image_id ?image_number ?placement ?quiet () =
  { (make `Display) with image_id; image_number; placement; quiet }

let delete ?quiet del = { (make `Delete) with quiet; delete = Some del }

let frame ?image_id ?image_number ?format ?transmission ?compression ?width
    ?height ?quiet ~frame () =
  {
    (make `Frame) with
    image_id;
    image_number;
    format;
    transmission;
    compression;
    width;
    height;
    quiet;
    frame = Some frame;
  }

let animate ?image_id ?image_number ?quiet anim =
  { (make `Animate) with image_id; image_number; quiet; animation = Some anim }

let compose ?image_id ?image_number ?quiet comp =
  { (make `Compose) with image_id; image_number; quiet; compose = Some comp }

(* Serialization helpers *)
let apc_start = "\027_G"
let apc_end = "\027\\"

(* Key-value writer with separator handling *)
type kv_writer = { mutable first : bool; buf : Buffer.t }

let kv_writer buf = { first = true; buf }

let kv w key value =
  if not w.first then Buffer.add_char w.buf ',';
  w.first <- false;
  Buffer.add_char w.buf key;
  Buffer.add_char w.buf '=';
  Buffer.add_string w.buf value

let kv_int w key value = kv w key (string_of_int value)
let kv_int32 w key value = kv w key (Int32.to_string value)
let kv_char w key value = kv w key (String.make 1 value)

(* Conditional writers using Option.iter *)
let kv_int_opt w key = Option.iter (kv_int w key)
let kv_int32_opt w key = Option.iter (kv_int32 w key)

let kv_int_if w key ~default opt =
  Option.iter (fun v -> if v <> default then kv_int w key v) opt

let action_char : action -> char = function
  | `Transmit -> 't'
  | `Transmit_and_display -> 'T'
  | `Query -> 'q'
  | `Display -> 'p'
  | `Delete -> 'd'
  | `Frame -> 'f'
  | `Animate -> 'a'
  | `Compose -> 'c'

let delete_char : Kgp_types.delete -> char = function
  | `All_visible -> 'a'
  | `All_visible_and_free -> 'A'
  | `By_id _ -> 'i'
  | `By_id_and_free _ -> 'I'
  | `By_number _ -> 'n'
  | `By_number_and_free _ -> 'N'
  | `At_cursor -> 'c'
  | `At_cursor_and_free -> 'C'
  | `At_cell _ -> 'p'
  | `At_cell_and_free _ -> 'P'
  | `At_cell_z _ -> 'q'
  | `At_cell_z_and_free _ -> 'Q'
  | `By_column _ -> 'x'
  | `By_column_and_free _ -> 'X'
  | `By_row _ -> 'y'
  | `By_row_and_free _ -> 'Y'
  | `By_z_index _ -> 'z'
  | `By_z_index_and_free _ -> 'Z'
  | `By_id_range _ -> 'r'
  | `By_id_range_and_free _ -> 'R'
  | `Frames -> 'f'
  | `Frames_and_free -> 'F'

let write_placement w (p : Kgp_placement.t) =
  kv_int_opt w 'x' p.source_x;
  kv_int_opt w 'y' p.source_y;
  kv_int_opt w 'w' p.source_width;
  kv_int_opt w 'h' p.source_height;
  kv_int_opt w 'X' p.cell_x_offset;
  kv_int_opt w 'Y' p.cell_y_offset;
  kv_int_opt w 'c' p.columns;
  kv_int_opt w 'r' p.rows;
  kv_int_opt w 'z' p.z_index;
  kv_int_opt w 'p' p.placement_id;
  p.cursor
  |> Option.iter (fun c ->
         kv_int_if w 'C' ~default:0 (Some (Kgp_types.Cursor.to_int c)));
  if p.unicode_placeholder then kv_int w 'U' 1

let write_delete w (d : Kgp_types.delete) =
  kv_char w 'd' (delete_char d);
  match d with
  | `By_id (id, pid) | `By_id_and_free (id, pid) ->
      kv_int w 'i' id;
      kv_int_opt w 'p' pid
  | `By_number (n, pid) | `By_number_and_free (n, pid) ->
      kv_int w 'I' n;
      kv_int_opt w 'p' pid
  | `At_cell (x, y) | `At_cell_and_free (x, y) ->
      kv_int w 'x' x;
      kv_int w 'y' y
  | `At_cell_z (x, y, z) | `At_cell_z_and_free (x, y, z) ->
      kv_int w 'x' x;
      kv_int w 'y' y;
      kv_int w 'z' z
  | `By_column c | `By_column_and_free c -> kv_int w 'x' c
  | `By_row r | `By_row_and_free r -> kv_int w 'y' r
  | `By_z_index z | `By_z_index_and_free z -> kv_int w 'z' z
  | `By_id_range (min_id, max_id) | `By_id_range_and_free (min_id, max_id) ->
      kv_int w 'x' min_id;
      kv_int w 'y' max_id
  | `All_visible | `All_visible_and_free | `At_cursor | `At_cursor_and_free
  | `Frames | `Frames_and_free ->
      ()

let write_frame w (f : Kgp_frame.t) =
  kv_int_opt w 'x' f.x;
  kv_int_opt w 'y' f.y;
  kv_int_opt w 'c' f.base_frame;
  kv_int_opt w 'r' f.edit_frame;
  kv_int_opt w 'z' f.gap_ms;
  f.composition
  |> Option.iter (fun c ->
         kv_int_if w 'X' ~default:0 (Some (Kgp_types.Composition.to_int c)));
  kv_int32_opt w 'Y' f.background_color

let write_animation w : Kgp_animation.t -> unit = function
  | `Set_state (state, loops) ->
      let s = match state with `Stop -> 1 | `Loading -> 2 | `Run -> 3 in
      kv_int w 's' s;
      kv_int_opt w 'v' loops
  | `Set_gap (frame, gap_ms) ->
      kv_int w 'r' frame;
      kv_int w 'z' gap_ms
  | `Set_current frame -> kv_int w 'c' frame

let write_compose w (c : Kgp_compose.t) =
  kv_int w 'r' c.source_frame;
  kv_int w 'c' c.dest_frame;
  kv_int_opt w 'w' c.width;
  kv_int_opt w 'h' c.height;
  kv_int_opt w 'x' c.dest_x;
  kv_int_opt w 'y' c.dest_y;
  kv_int_opt w 'X' c.source_x;
  kv_int_opt w 'Y' c.source_y;
  c.composition
  |> Option.iter (fun comp ->
         kv_int_if w 'C' ~default:0 (Some (Kgp_types.Composition.to_int comp)))

let write_control_data buf cmd =
  let w = kv_writer buf in
  (* Action *)
  kv_char w 'a' (action_char cmd.action);
  (* Quiet - only if non-default *)
  cmd.quiet
  |> Option.iter (fun q ->
         kv_int_if w 'q' ~default:0 (Some (Kgp_types.Quiet.to_int q)));
  (* Format *)
  cmd.format
  |> Option.iter (fun f -> kv_int w 'f' (Kgp_types.Format.to_int f));
  (* Transmission - only for transmit/frame actions, always include t=d for compatibility *)
  (match cmd.action with
  | `Transmit | `Transmit_and_display | `Frame -> (
      match cmd.transmission with
      | Some t -> kv_char w 't' (Kgp_types.Transmission.to_char t)
      | None -> kv_char w 't' 'd')
  | _ -> ());
  (* Compression *)
  cmd.compression
  |> Option.iter (fun c ->
         Kgp_types.Compression.to_char c |> Option.iter (kv_char w 'o'));
  (* Dimensions *)
  kv_int_opt w 's' cmd.width;
  kv_int_opt w 'v' cmd.height;
  (* File size/offset *)
  kv_int_opt w 'S' cmd.size;
  kv_int_opt w 'O' cmd.offset;
  (* Image ID/number *)
  kv_int_opt w 'i' cmd.image_id;
  kv_int_opt w 'I' cmd.image_number;
  (* Complex options *)
  cmd.placement |> Option.iter (write_placement w);
  cmd.delete |> Option.iter (write_delete w);
  cmd.frame |> Option.iter (write_frame w);
  cmd.animation |> Option.iter (write_animation w);
  cmd.compose |> Option.iter (write_compose w);
  w

(* Use large chunk size to avoid chunking - Kitty animation doesn't handle chunks well *)
let chunk_size = 1024 * 1024 (* 1MB - effectively no chunking *)

let write buf cmd ~data =
  Buffer.add_string buf apc_start;
  let w = write_control_data buf cmd in
  if String.length data > 0 then begin
    let encoded = Base64.encode_string data in
    let len = String.length encoded in
    if len <= chunk_size then (
      Buffer.add_char buf ';';
      Buffer.add_string buf encoded;
      Buffer.add_string buf apc_end)
    else begin
      (* Multiple chunks *)
      let rec write_chunks pos first =
        if pos < len then begin
          let remaining = len - pos in
          let this_chunk = min chunk_size remaining in
          let is_last = pos + this_chunk >= len in
          if first then (
            kv_int w 'm' 1;
            Buffer.add_char buf ';';
            Buffer.add_substring buf encoded pos this_chunk;
            Buffer.add_string buf apc_end)
          else (
            Buffer.add_string buf apc_start;
            Buffer.add_string buf (if is_last then "m=0" else "m=1");
            Buffer.add_char buf ';';
            Buffer.add_substring buf encoded pos this_chunk;
            Buffer.add_string buf apc_end);
          write_chunks (pos + this_chunk) false
        end
      in
      write_chunks 0 true
    end
  end
  else Buffer.add_string buf apc_end

let to_string cmd ~data =
  let buf = Buffer.create 1024 in
  write buf cmd ~data;
  Buffer.contents buf