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

(** {1 Progress Bar State} *)

type state = { total : int; current : int; message : string; phase : string }

let state ?(phase = "") ?(current = 0) ~total message =
  { total; current; message; phase }

let with_current current s = { s with current }
let with_message message s = { s with message }
let with_phase phase s = { s with phase }
let incr s = { s with current = s.current + 1 }

let update ?(phase = "") ~message s =
  { s with current = s.current + 1; message; phase }

(** {1 Configuration} *)

type config = {
  ppf : Format.formatter;
  width : int;
  style : [ `ASCII | `UTF8 | `Plain ];
  bar_width : int;
}

let config ?(ppf = Format.std_formatter) ?width ?(style = `UTF8)
    ?(bar_width = 25) () =
  let width =
    match width with Some w -> w | None -> Width.terminal_width ()
  in
  { ppf; width; style; bar_width }

(** {1 Rendering} *)

(* Braille spinner frames *)
let spinner_utf8 = [| "⠋"; "⠙"; "⠹"; "⠸"; "⠼"; "⠴"; "⠦"; "⠧"; "⠇"; "⠏" |]

let spinner_ascii = [| "|"; "/"; "-"; "\\" |]

(* Bar characters for smooth animation *)
let bar_chars = [| " "; "▏"; "▎"; "▍"; "▌"; "▋"; "▊"; "▉"; "█" |]

(* ANSI codes *)
let cyan = "\027[36m"
let green = "\027[32m"
let dim = "\027[2m"
let reset_code = "\027[0m"
let hide_cursor = "\027[?25l"
let show_cursor = "\027[?25h"

(* Ensure cursor is restored on exit (handles Ctrl-C, uncaught exceptions) *)
let cursor_hidden = ref false

let restore_cursor () =
  if !cursor_hidden then begin
    output_string stdout show_cursor;
    flush stdout;
    cursor_hidden := false
  end

let () = at_exit restore_cursor

let render_bar_utf8 ~bar_width ~pct =
  let total_eighths = bar_width * 8 in
  let filled_eighths = total_eighths * pct / 100 in
  let full_blocks = filled_eighths / 8 in
  let partial = filled_eighths mod 8 in
  let empty_blocks = bar_width - full_blocks - if partial > 0 then 1 else 0 in
  let full =
    String.concat "" (List.init full_blocks (fun _ -> bar_chars.(8)))
  in
  let part = if partial > 0 then bar_chars.(partial) else "" in
  let empty =
    String.concat "" (List.init empty_blocks (fun _ -> bar_chars.(0)))
  in
  full ^ part ^ empty

let render_bar_ascii ~bar_width ~pct =
  let filled = bar_width * pct / 100 in
  let empty = max 0 (bar_width - filled - 1) in
  String.make filled '='
  ^ (if filled < bar_width then ">" else "")
  ^ String.make empty ' '

let pad_to_width width s =
  let w = Width.string_width s in
  if w < width then s ^ String.make (width - w) ' ' else s

let render ~frame cfg s =
  let pct = if s.total > 0 then min 100 (s.current * 100 / s.total) else 0 in
  let counter = Fmt.str "%d/%d" s.current s.total in
  let full_msg =
    if s.phase <> "" then s.phase ^ ": " ^ s.message else s.message
  in
  match cfg.style with
  | `UTF8 ->
      let bar = render_bar_utf8 ~bar_width:cfg.bar_width ~pct in
      let spinner = spinner_utf8.(frame mod Array.length spinner_utf8) in
      (* Calculate prefix width to determine how much space is left for message *)
      let prefix = Fmt.str "%s ▕%s▏ %3d%% %s " spinner bar pct counter in
      let prefix_width = Width.string_width prefix in
      let max_msg_len = max 0 (cfg.width - prefix_width) in
      let display_msg =
        if String.length full_msg > max_msg_len then
          String.sub full_msg 0 (max 0 (max_msg_len - 1)) ^ "…"
        else full_msg
      in
      let visible_width = prefix_width + Width.string_width display_msg in
      let padding =
        if visible_width < cfg.width then
          String.make (cfg.width - visible_width) ' '
        else ""
      in
      Fmt.str "%s%s%s ▕%s%s%s▏ %3d%% %s%s%s %s%s" cyan spinner reset_code green
        bar reset_code pct dim counter reset_code display_msg padding
  | `ASCII ->
      let bar = render_bar_ascii ~bar_width:cfg.bar_width ~pct in
      let spinner = spinner_ascii.(frame mod Array.length spinner_ascii) in
      let prefix = Fmt.str "%s [%s] %3d%% %s " spinner bar pct counter in
      let prefix_width = Width.string_width prefix in
      let max_msg_len = max 0 (cfg.width - prefix_width) in
      let display_msg =
        if String.length full_msg > max_msg_len then
          String.sub full_msg 0 (max 0 (max_msg_len - 3)) ^ "..."
        else full_msg
      in
      pad_to_width cfg.width (prefix ^ display_msg)
  | `Plain ->
      (* Minimal: no spinner, no colors, just counter and message *)
      pad_to_width cfg.width (Fmt.str "[%3d%%] %s %s" pct counter full_msg)

(** {1 Imperative Wrapper} *)

type t = {
  cfg : config;
  mutable s : state;
  mutable frame : int;
  mutable finished : bool;
  enabled : bool;
  start_time : float;
}

let debug =
  match Sys.getenv_opt "TTY_PROGRESS_DEBUG" with
  | Some _ -> true
  | None -> false

let output_line ppf line =
  (* Use direct stdout when ppf is std_formatter to bypass Format line-breaking *)
  let use_direct = ppf == Format.std_formatter in
  if debug then
    Fmt.epr "[PROGRESS] use_direct=%b is_tty=%b line_bytes=%d line_width=%d\n"
      use_direct (Width.is_tty ()) (String.length line)
      (Width.string_width line);
  (* Use \r\027[K to: move to start of line, clear to end of line *)
  let cr_clear = "\r\027[K" in
  if use_direct then begin
    output_string stdout cr_clear;
    output_string stdout line;
    flush stdout
  end
  else begin
    (* For buffers/custom formatters, use Format but disable line-breaking *)
    let old_margin = Format.pp_get_margin ppf () in
    Format.pp_set_margin ppf max_int;
    Fmt.string ppf cr_clear;
    Fmt.string ppf line;
    Format.pp_print_flush ppf ();
    Format.pp_set_margin ppf old_margin
  end

let v ?(ppf = Format.std_formatter) ?width ?enabled ?(style = `UTF8) ~total msg
    =
  let cfg = config ~ppf ?width ~style () in
  let enabled = match enabled with Some e -> e | None -> Width.is_tty () in
  let s = state ~total msg in
  let start_time = Unix.gettimeofday () in
  let t = { cfg; s; frame = 0; finished = false; enabled; start_time } in
  (* Hide cursor and render immediately *)
  if enabled && total > 0 then begin
    if ppf == Format.std_formatter then begin
      cursor_hidden := true;
      output_string stdout hide_cursor;
      flush stdout
    end;
    let line = render ~frame:0 cfg s in
    output_line ppf line
  end;
  t

let position t = t.s.current

let pp ppf t =
  let status = if t.finished then "finished" else "active" in
  Fmt.pf ppf "progress(%s %d/%d %S)" status t.s.current t.s.total t.s.message

let display t =
  if t.enabled && not t.finished then begin
    (* Time-based frame for faster spinner: ~10 frames per second *)
    let elapsed = Unix.gettimeofday () -. t.start_time in
    let frame = int_of_float (elapsed *. 10.) in
    t.frame <- frame;
    let line = render ~frame t.cfg t.s in
    output_line t.cfg.ppf line
  end

let set t pos =
  t.s <- with_current pos t.s;
  display t

let tick t =
  t.s <- incr t.s;
  display t

let message t msg =
  t.s <- with_message msg t.s;
  display t

let set_phase t phase =
  t.s <- with_phase phase t.s;
  display t

let update t ~phase ~msg =
  t.s <- update ~phase ~message:msg t.s;
  display t

let reset t ~total msg =
  t.s <- state ~total msg;
  t.frame <- 0;
  t.finished <- false;
  display t

let finish t =
  if not t.finished then begin
    t.s <- with_current t.s.total t.s;
    display t;
    t.finished <- true;
    if t.enabled then begin
      if t.cfg.ppf == Format.std_formatter then begin
        cursor_hidden := false;
        output_string stdout show_cursor;
        output_char stdout '\n';
        flush stdout
      end
      else begin
        Format.pp_print_newline t.cfg.ppf ();
        Format.pp_print_flush t.cfg.ppf ()
      end
    end
  end

let clear t =
  if not t.finished then begin
    t.finished <- true;
    if t.enabled then begin
      if t.cfg.ppf == Format.std_formatter then begin
        cursor_hidden := false;
        output_string stdout "\r\027[K";
        output_string stdout show_cursor;
        flush stdout
      end
      else begin
        Fmt.string t.cfg.ppf "\r\027[K";
        Format.pp_print_flush t.cfg.ppf ()
      end
    end
  end