Source file matrix_ansi_parser.ml

(******************************************************************************)
(*                                                                            *)
(* SPDX-License-Identifier: MIT                                               *)
(* Copyright (c) 2025 Nomadic Labs <contact@nomadic-labs.com>                 *)
(*                                                                            *)
(******************************************************************************)

type parse_state =
  | Normal
  | EscapeStart
  | CSI of int list * int (* accumulated params, current param *)
  | OSC (* inside an OSC sequence — skip until ST *)
  | DCS (* inside a DCS sequence (e.g. Sixel) — accumulate until ST *)

type t = {
  mutable style : Matrix_cell.style;
  mutable state : parse_state;
  osc_buf : Buffer.t;
  dcs_buf : Buffer.t;
}

let create () =
  {
    style = Matrix_cell.default_style;
    state = Normal;
    osc_buf = Buffer.create 64;
    dcs_buf = Buffer.create 1024;
  }

let reset t =
  t.style <- Matrix_cell.default_style ;
  t.state <- Normal ;
  Buffer.clear t.osc_buf ;
  Buffer.clear t.dcs_buf

let current_style t = t.style

(* Apply SGR (Select Graphic Rendition) parameters to style *)
let apply_sgr params style =
  let rec process params style =
    match params with
    | [] -> style
    | 0 :: rest ->
        (* Reset all SGR attributes — preserves URL (OSC 8 is not SGR) *)
        process
          rest
          {
            Matrix_cell.default_style with
            Matrix_cell.url = style.Matrix_cell.url;
          }
    | 1 :: rest ->
        (* Bold *)
        process rest {style with bold = true}
    | 2 :: rest ->
        (* Dim *)
        process rest {style with dim = true}
    | 4 :: rest ->
        (* Underline *)
        process rest {style with underline = true}
    | 7 :: rest ->
        (* Reverse *)
        process rest {style with reverse = true}
    | 22 :: rest ->
        (* Normal intensity (reset bold/dim) *)
        process rest {style with bold = false; dim = false}
    | 24 :: rest ->
        (* Underline off *)
        process rest {style with underline = false}
    | 27 :: rest ->
        (* Reverse off *)
        process rest {style with reverse = false}
    (* Basic foreground colors 30-37 *)
    | n :: rest when n >= 30 && n <= 37 -> process rest {style with fg = n - 30}
    | 39 :: rest ->
        (* Default foreground *)
        process rest {style with fg = -1}
    (* Basic background colors 40-47 *)
    | n :: rest when n >= 40 && n <= 47 -> process rest {style with bg = n - 40}
    | 49 :: rest ->
        (* Default background *)
        process rest {style with bg = -1}
    (* Bright foreground colors 90-97 *)
    | n :: rest when n >= 90 && n <= 97 ->
        process rest {style with fg = n - 90 + 8}
    (* Bright background colors 100-107 *)
    | n :: rest when n >= 100 && n <= 107 ->
        process rest {style with bg = n - 100 + 8}
    (* 256-color foreground: 38;5;N *)
    | 38 :: 5 :: n :: rest ->
        let fg = if n >= 0 && n <= 255 then n else -1 in
        process rest {style with fg}
    (* 256-color background: 48;5;N *)
    | 48 :: 5 :: n :: rest ->
        let bg = if n >= 0 && n <= 255 then n else -1 in
        process rest {style with bg}
    (* Unknown - skip *)
    | _ :: rest -> process rest style
  in
  process params style

(* Get UTF-8 character length in bytes *)
let utf8_char_length c =
  let code = Char.code c in
  if code land 0x80 = 0 then 1
  else if code land 0xE0 = 0xC0 then 2
  else if code land 0xF0 = 0xE0 then 3
  else if code land 0xF8 = 0xF0 then 4
  else 1 (* Invalid, treat as single byte *)

(* Extract UTF-8 character from string at position *)
let extract_utf8_char s pos =
  if pos >= String.length s then ("", 0)
  else
    let len = utf8_char_length s.[pos] in
    let len = min len (String.length s - pos) in
    (String.sub s pos len, len)

(* Process a completed OSC sequence payload.
   For OSC 8 (hyperlinks): "8;params;uri" — extract the URI and update style.
   Empty URI closes the current hyperlink. *)
let process_osc t =
  let payload = Buffer.contents t.osc_buf in
  let plen = String.length payload in
  (* Check for OSC 8: starts with "8;" *)
  if plen >= 2 && payload.[0] = '8' && payload.[1] = ';' then
    (* Find the second semicolon that separates params from URI *)
    match String.index_from_opt payload 2 ';' with
    | Some idx ->
        let url = String.sub payload (idx + 1) (plen - idx - 1) in
        t.style <- {t.style with url}
    | None ->
        (* Malformed OSC 8 — no second semicolon, ignore *)
        ()

(* Core parsing function - shared state machine logic.
   Takes a callback ~emit_char that receives (char, style) for each visible character.
   Returns the number of visible characters parsed. *)
let parse_core t ~emit_char input =
  let len = String.length input in
  let count = ref 0 in
  let pos = ref 0 in

  while !pos < len do
    match t.state with
    | Normal -> (
        let c = input.[!pos] in
        match c with
        | '\027' ->
            t.state <- EscapeStart ;
            incr pos
        | '\n' | '\r' ->
            (* Skip newlines in single-line parsing *)
            incr pos
        | _ ->
            (* Regular character - extract full UTF-8 *)
            let char, char_len = extract_utf8_char input !pos in
            if char_len > 0 then begin
              emit_char char t.style ;
              incr count ;
              pos := !pos + char_len
            end
            else incr pos)
    | EscapeStart ->
        if !pos < len then (
          match input.[!pos] with
          | '[' ->
              t.state <- CSI ([], 0) ;
              incr pos
          | ']' ->
              (* OSC sequence (e.g. OSC 8 hyperlinks) *)
              Buffer.clear t.osc_buf ;
              t.state <- OSC ;
              incr pos
          | 'P' ->
              (* DCS sequence (e.g. Sixel graphics) — accumulate and emit as cell *)
              Buffer.clear t.dcs_buf ;
              Buffer.add_string t.dcs_buf "\027P" ;
              t.state <- DCS ;
              incr pos
          | _ ->
              (* Not a CSI/OSC/DCS sequence, back to normal *)
              t.state <- Normal ;
              incr pos)
        else t.state <- Normal
    | OSC ->
        (* Collect bytes until String Terminator: ESC \ or BEL (\007) *)
        if !pos < len then
          let c = input.[!pos] in
          if c = '\007' then (
            (* BEL terminates OSC — process payload *)
            process_osc t ;
            t.state <- Normal ;
            incr pos)
          else if c = '\027' && !pos + 1 < len && input.[!pos + 1] = '\\' then (
            (* ESC \ terminates OSC — process payload *)
            process_osc t ;
            t.state <- Normal ;
            pos := !pos + 2)
          else (
            Buffer.add_char t.osc_buf c ;
            incr pos)
        else t.state <- Normal
    | DCS ->
        (* Accumulate DCS payload until ESC \ (String Terminator) *)
        if !pos < len then
          let c = input.[!pos] in
          if c = '\027' && !pos + 1 < len && input.[!pos + 1] = '\\' then begin
            Buffer.add_string t.dcs_buf "\027\\" ;
            (* Emit the entire DCS sequence as a single passthrough cell *)
            emit_char (Buffer.contents t.dcs_buf) t.style ;
            t.state <- Normal ;
            pos := !pos + 2
          end
          else begin
            Buffer.add_char t.dcs_buf c ;
            incr pos
          end
        else t.state <- Normal
    | CSI (params, current) ->
        if !pos < len then (
          let c = input.[!pos] in
          match c with
          | '0' .. '9' ->
              let digit = Char.code c - Char.code '0' in
              t.state <- CSI (params, (current * 10) + digit) ;
              incr pos
          | ';' ->
              t.state <- CSI (params @ [current], 0) ;
              incr pos
          | 'm' ->
              (* SGR complete *)
              let all_params = params @ [current] in
              t.style <- apply_sgr all_params t.style ;
              t.state <- Normal ;
              incr pos
          | 'A' .. 'Z' | 'a' .. 'l' | 'n' .. 'z' ->
              (* Other CSI sequence terminator - ignore *)
              t.state <- Normal ;
              incr pos
          | _ ->
              (* Unknown, abort CSI parsing *)
              t.state <- Normal ;
              incr pos)
        else t.state <- Normal
  done ;
  !count

(* Parse a single line into a buffer at given row/col *)
let parse_line t buf ~row ~col input =
  let col = ref col in
  let emit_char char style =
    Matrix_buffer.set_char buf ~row ~col:!col ~char ~style ;
    incr col
  in
  let _ = parse_core t ~emit_char input in
  !col

(* Parse into buffer, handling newlines *)
let parse_into t buf ~row ~col input =
  let lines = String.split_on_char '\n' input in
  let row = ref row in
  let col = ref col in
  List.iteri
    (fun i line ->
      if i > 0 then begin
        incr row ;
        col := 0
      end ;
      col := parse_line t buf ~row:!row ~col:!col line)
    lines ;
  (!row, !col)

(* Batch version using batch_ops for thread-safe access *)
let parse_line_batch t (ops : Matrix_buffer.batch_ops) ~row ~col input =
  let col = ref col in
  let emit_char char style =
    ops.set_char ~row ~col:!col ~char ~style ;
    incr col
  in
  let _ = parse_core t ~emit_char input in
  !col

let parse_into_batch t ops ~row ~col input =
  let lines = String.split_on_char '\n' input in
  let row = ref row in
  let col = ref col in
  List.iteri
    (fun i line ->
      if i > 0 then begin
        incr row ;
        col := 0
      end ;
      col := parse_line_batch t ops ~row:!row ~col:!col line)
    lines ;
  (!row, !col)

(* Parse to a list of (char, style) pairs *)
let parse_to_cells t input =
  let results = ref [] in
  let emit_char char style = results := (char, style) :: !results in
  let _ = parse_core t ~emit_char input in
  List.rev !results

(* Calculate visible length (chars excluding ANSI codes and OSC sequences) *)
let visible_length input =
  let len = String.length input in
  let count = ref 0 in
  let pos = ref 0 in
  let in_escape = ref false in
  let in_osc = ref false in
  let in_dcs = ref false in

  while !pos < len do
    let c = input.[!pos] in
    if !in_dcs then begin
      (* Skip DCS payload until ESC \ *)
      if c = '\027' && !pos + 1 < len && input.[!pos + 1] = '\\' then (
        in_dcs := false ;
        pos := !pos + 2)
      else incr pos
    end
    else if !in_osc then begin
      if c = '\007' then (
        in_osc := false ;
        incr pos)
      else if c = '\027' && !pos + 1 < len && input.[!pos + 1] = '\\' then (
        in_osc := false ;
        pos := !pos + 2)
      else incr pos
    end
    else if !in_escape then begin
      if c = ']' then (
        in_escape := false ;
        in_osc := true ;
        incr pos)
      else if c = 'P' then (
        (* DCS sequence *)
        in_escape := false ;
        in_dcs := true ;
        incr pos)
      else (
        if c = 'm' || (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') then
          in_escape := false ;
        incr pos)
    end
    else if c = '\027' then begin
      in_escape := true ;
      incr pos
    end
    else if c = '\n' || c = '\r' then incr pos
    else begin
      let char_len = utf8_char_length c in
      incr count ;
      pos := !pos + char_len
    end
  done ;
  !count