(*---------------------------------------------------------------------------
   Copyright (c) 2025 Thomas Gazagnaire. All rights reserved.
   SPDX-License-Identifier: MIT
  ---------------------------------------------------------------------------*)

(* ANSI escape sequence detection *)
type ansi_state = Normal | Escape | Csi

let char_width uchar =
  match Uucp.Break.tty_width_hint uchar with
  | -1 -> 0 (* Control characters *)
  | n -> n

let string_width str =
  let len = String.length str in
  let width = ref 0 in
  let state = ref Normal in
  let decoder = Uutf.decoder ~encoding:`UTF_8 (`String str) in
  let rec loop () =
    match Uutf.decode decoder with
    | `Uchar u -> (
        let c = Uchar.to_int u in
        match !state with
        | Normal ->
            if c = 0x1b then state := Escape else width := !width + char_width u;
            loop ()
        | Escape ->
            if c = 0x5b (* '[' *) then state := Csi else state := Normal;
            loop ()
        | Csi ->
            (* CSI sequence ends with byte in range 0x40-0x7E *)
            if c >= 0x40 && c <= 0x7e then state := Normal;
            loop ())
    | `End -> ()
    | `Await -> assert false
    | `Malformed _ -> loop ()
  in
  if len = 0 then 0
  else (
    loop ();
    !width)

let truncate target_width str =
  if target_width <= 0 then ""
  else
    let buf = Buffer.create (String.length str) in
    let width = ref 0 in
    let state = ref Normal in
    let decoder = Uutf.decoder ~encoding:`UTF_8 (`String str) in
    let rec loop () =
      match Uutf.decode decoder with
      | `Uchar u -> (
          let c = Uchar.to_int u in
          match !state with
          | Normal ->
              if c = 0x1b then (
                state := Escape;
                Uutf.Buffer.add_utf_8 buf u;
                loop ())
              else
                let w = char_width u in
                if !width + w <= target_width then (
                  width := !width + w;
                  Uutf.Buffer.add_utf_8 buf u;
                  loop ())
                else ()
          | Escape ->
              Uutf.Buffer.add_utf_8 buf u;
              if c = 0x5b then state := Csi else state := Normal;
              loop ()
          | Csi ->
              Uutf.Buffer.add_utf_8 buf u;
              if c >= 0x40 && c <= 0x7e then state := Normal;
              loop ())
      | `End -> ()
      | `Await -> assert false
      | `Malformed _ -> loop ()
    in
    loop ();
    Buffer.contents buf

let pad_right target_width str =
  let str_width = string_width str in
  if str_width >= target_width then str
  else str ^ String.make (target_width - str_width) ' '

let pad_left target_width str =
  let str_width = string_width str in
  if str_width >= target_width then str
  else String.make (target_width - str_width) ' ' ^ str

let center target_width str =
  let str_width = string_width str in
  if str_width >= target_width then str
  else
    let total_pad = target_width - str_width in
    let left_pad = total_pad / 2 in
    let right_pad = total_pad - left_pad in
    String.make left_pad ' ' ^ str ^ String.make right_pad ' '

let wrap ?(indent = 0) width text =
  let effective = width - indent in
  if effective <= 0 then text
  else
    let prefix = String.make indent ' ' in
    let normalized =
      text |> String.split_on_char '\n' |> List.map String.trim
      |> String.concat " "
    in
    let words = String.split_on_char ' ' normalized in
    let rec build acc line len = function
      | [] -> if line = "" then acc else line :: acc
      | word :: rest ->
          let wlen = String.length word in
          let space = if line = "" then 0 else 1 in
          if len + space + wlen <= effective then
            let line = if line = "" then word else line ^ " " ^ word in
            build acc line (len + space + wlen) rest
          else if line = "" then
            (* Word longer than width; accept it on its own line *)
            build (word :: acc) "" 0 rest
          else build (line :: acc) word wlen rest
    in
    let lines = List.rev (build [] "" 0 words) in
    String.concat "\n" (List.map (fun l -> prefix ^ l) lines)

(* Terminal queries *)

let cached_terminal_width = ref None

let terminal_width () =
  match !cached_terminal_width with
  | Some cached -> cached
  | None ->
      let width =
        let from_env =
          match Sys.getenv_opt "COLUMNS" with
          | Some cols -> int_of_string_opt cols
          | None -> None
        in
        let from_tput () =
          try
            let ic = Unix.open_process_in "tput cols 2>/dev/null" in
            let result =
              try input_line ic |> int_of_string_opt with End_of_file -> None
            in
            ignore (Unix.close_process_in ic);
            result
          with Unix.Unix_error _ -> None
        in
        (* Use env if valid (> 10), otherwise try tput, fallback to 80 *)
        match from_env with
        | Some w when w > 10 -> w
        | _ -> ( match from_tput () with Some w when w > 10 -> w | _ -> 80)
      in
      cached_terminal_width := Some width;
      width

let is_tty () = Unix.isatty Unix.stdout