open Lwt.Infix
open Notty
open Notty_lwt

type stats =
  { name : string
  ; total_pieces : int
  ; completed_pieces : int
  ; peers_count : int
  ; download_speed : float
  }

let progress_bar ~width pct =
  let filled = (pct * width) / 100 in
  let filled = max 0 (min width filled) in
  let bar_filled = String.make filled '#' in
  let bar_empty = String.make (width - filled) '-' in
  I.hcat
    [ I.string A.(fg green) bar_filled
    ; I.string A.(fg (gray 8)) bar_empty
    ]
;;

let format_speed bytes_per_sec =
  if bytes_per_sec < 1024.0 then
    Printf.sprintf "%.0f B/s" bytes_per_sec
  else if bytes_per_sec < 1024.0 *. 1024.0 then
    Printf.sprintf "%.1f KB/s" (bytes_per_sec /. 1024.0)
  else
    Printf.sprintf "%.2f MB/s" (bytes_per_sec /. 1024.0 /. 1024.0)
;;

let render stats =
  let pct = if stats.total_pieces = 0 then 0 else
    (stats.completed_pieces * 100) / stats.total_pieces in

  let header = I.string A.(st bold) "Tornado BitTorrent Client" in
  let separator = I.string A.(fg (gray 12)) (String.make 60 '=') in

  let name_line = I.hcat
    [ I.string A.(fg cyan) "Torrent: "
    ; I.string A.empty stats.name
    ] in

  let progress_line = I.hcat
    [ I.string A.(fg cyan) "Progress: "
    ; progress_bar ~width:30 pct
    ; I.strf " %d%%" pct
    ] in

  let pieces_line = I.hcat
    [ I.string A.(fg cyan) "Pieces: "
    ; I.strf "%d / %d" stats.completed_pieces stats.total_pieces
    ] in

  let peers_line = I.hcat
    [ I.string A.(fg cyan) "Peers: "
    ; I.strf "%d" stats.peers_count
    ] in

  let speed_line = I.hcat
    [ I.string A.(fg cyan) "Speed: "
    ; I.string A.empty (format_speed stats.download_speed)
    ] in

  let footer = I.string A.(fg (gray 12)) "Press Ctrl-C to stop" in

  I.vcat
    [ I.void 0 1
    ; header
    ; I.void 0 1
    ; separator
    ; I.void 0 1
    ; name_line
    ; I.void 0 1
    ; progress_line
    ; I.void 0 1
    ; pieces_line
    ; peers_line
    ; speed_line
    ; I.void 0 1
    ; separator
    ; I.void 0 1
    ; footer
    ; I.void 0 1
    ]
;;

let update_display term stats =
  let image = render stats in
  Term.image term image
;;

let create_tui name total_pieces =
  let stats = ref {
    name;
    total_pieces;
    completed_pieces = 0;
    peers_count = 0;
    download_speed = 0.0;
  } in

  let term = Term.create () in
  Term.cursor term None >>= fun () ->
  update_display term !stats >>= fun () ->
  Lwt.return (Some term, stats)
;;

let update_stats stats_ref ~completed ~peers ~speed =
  stats_ref := { !stats_ref with
    completed_pieces = completed;
    peers_count = peers;
    download_speed = speed;
  }
;;

let close_tui term_opt =
  match term_opt with
  | None -> Lwt.return_unit
  | Some term ->
    Term.cursor term (Some (0, 0)) >>= fun () ->
    Term.release term >>= fun () ->
    Lwt.return_unit
;;
