open Lwt.Infix
open Cohttp
open Cohttp_lwt_unix

(* Chunk size: 256KB *)
let chunk_size = 262144

type download_info =
  { url : string
  ; total_size : int64
  ; supports_ranges : bool
  ; filename : string
  }

let state_file_path output_file =
  output_file ^ ".http-state"
;;

let save_http_state output_file downloaded_bytes total_bytes =
  let state = Printf.sprintf "%Ld/%Ld" downloaded_bytes total_bytes in
  Lwt_io.with_file
    ~mode:Lwt_io.Output
    (state_file_path output_file)
    (fun oc -> Lwt_io.write oc state)
;;

let load_http_state output_file total_bytes =
  let path = state_file_path output_file in
  if Sys.file_exists path then
    Lwt.catch
      (fun () ->
        Lwt_io.with_file ~mode:Lwt_io.Input path (fun ic ->
          Lwt_io.read_line ic >>= fun line ->
          match String.split_on_char '/' line with
          | [downloaded; total] ->
            let downloaded_bytes = Int64.of_string downloaded in
            let total_in_state = Int64.of_string total in
            if total_in_state = total_bytes then (
              Logs.info (fun m -> m "Resuming download from %Ld bytes" downloaded_bytes);
              Lwt.return downloaded_bytes
            ) else (
              Logs.warn (fun m -> m "File size changed, restarting download");
              Lwt.return 0L
            )
          | _ -> Lwt.return 0L
        ))
      (fun _ -> Lwt.return 0L)
  else
    Lwt.return 0L
;;

let delete_http_state output_file =
  let path = state_file_path output_file in
  if Sys.file_exists path then
    Unix.unlink path
;;

let extract_filename_from_url url =
  let uri = Uri.of_string url in
  match Uri.path uri with
  | "" | "/" -> "download"
  | path ->
    let segments = String.split_on_char '/' path in
    let last_segment = List.nth segments (List.length segments - 1) in
    if last_segment = "" then "download" else last_segment
;;

let get_download_info url =
  Logs.info (fun m -> m "Fetching file information from %s" url);
  Client.head (Uri.of_string url) >>= fun resp ->
  let status = Response.status resp in
  let headers = Response.headers resp in

  match status with
  | `OK ->
    let content_length = Header.get headers "content-length" in
    let accept_ranges = Header.get headers "accept-ranges" in
    let supports_ranges = match accept_ranges with
      | Some "bytes" -> true
      | _ -> false
    in

    (match content_length with
     | Some len ->
       let total_size = Int64.of_string len in
       let filename = extract_filename_from_url url in
       Logs.info (fun m -> m "File size: %Ld bytes, Range support: %b" total_size supports_ranges);
       Lwt.return (Some { url; total_size; supports_ranges; filename })
     | None ->
       Logs.err (fun m -> m "Server did not provide Content-Length");
       Lwt.return None)
  | code ->
    Logs.err (fun m -> m "HEAD request failed with status: %s" (Code.string_of_status code));
    Lwt.return None
;;

let download_range_to_channel url start_byte end_byte oc =
  let range_header = Printf.sprintf "bytes=%Ld-%Ld" start_byte end_byte in
  let headers = Header.init_with "Range" range_header in

  Client.get ~headers (Uri.of_string url) >>= fun (resp, body) ->
  let status = Response.status resp in

  match status with
  | `Partial_content | `OK ->
    (* Stream body directly to the output channel *)
    let stream = Cohttp_lwt.Body.to_stream body in
    let bytes_written = ref 0 in
    Lwt_stream.iter_s (fun chunk ->
      bytes_written := !bytes_written + String.length chunk;
      Lwt_io.write oc chunk
    ) stream >>= fun () ->
    Lwt.return (Some !bytes_written)
  | code ->
    Logs.err (fun m -> m "Range request failed with status: %s" (Code.string_of_status code));
    Lwt.return None
;;

let download_full_to_channel url oc =
  Logs.info (fun m -> m "Downloading entire file (server doesn't support ranges)");
  Client.get (Uri.of_string url) >>= fun (resp, body) ->
  let status = Response.status resp in

  match status with
  | `OK ->
    (* Stream body directly to the output channel *)
    let stream = Cohttp_lwt.Body.to_stream body in
    Lwt_stream.iter_s (fun chunk ->
      Lwt_io.write oc chunk
    ) stream >>= fun () ->
    Lwt.return (Some ())
  | code ->
    Logs.err (fun m -> m "Download failed with status: %s" (Code.string_of_status code));
    Lwt.return None
;;

let download_http_file url output_file use_tui =
  get_download_info url >>= function
  | None ->
    Logs.err (fun m -> m "Failed to get file information");
    Lwt.return_unit
  | Some info ->
    let final_output = match output_file with
      | Some name -> name
      | None -> info.filename
    in

    Logs.info (fun m -> m "Downloading to: %s" final_output);

    if info.supports_ranges then (
      (* Resumable download with ranges *)
      load_http_state final_output info.total_size >>= fun start_pos ->

      (* Setup TUI if enabled *)
      let total_chunks = Int64.div (Int64.add info.total_size (Int64.of_int chunk_size)
                                    |> Int64.sub (Int64.of_int 1)) (Int64.of_int chunk_size) |> Int64.to_int in
      let completed_chunks = Int64.div start_pos (Int64.of_int chunk_size) |> Int64.to_int in

      (if use_tui then Tui.create_tui final_output total_chunks
       else Lwt.return (None, ref { Tui.name = final_output; total_pieces = total_chunks;
                                     completed_pieces = completed_chunks; peers_count = 1; download_speed = 0.0 }))
      >>= fun (term_opt, stats_ref) ->

      let downloaded = ref start_pos in
      let start_time = Unix.gettimeofday () in
      let should_quit = ref false in

      (* Setup Ctrl+C handler *)
      let sigint_handler = Lwt_unix.on_signal Sys.sigint (fun _ ->
        should_quit := true;
      ) in

      (* Update TUI periodically *)
      (if use_tui && term_opt <> None then
        let rec update_loop () =
          Lwt_unix.sleep 0.5 >>= fun () ->
          if !downloaded < info.total_size && not !should_quit then (
            let elapsed = Unix.gettimeofday () -. start_time in
            let bytes_dl = Int64.sub !downloaded start_pos |> Int64.to_float in
            let speed = if elapsed > 0.0 then bytes_dl /. elapsed else 0.0 in
            let chunks_done = Int64.div !downloaded (Int64.of_int chunk_size) |> Int64.to_int in
            Tui.update_stats stats_ref ~completed:chunks_done
              ~peers:1 ~speed;
            (match term_opt with
             | Some term -> Tui.update_display term !stats_ref
             | None -> Lwt.return_unit) >>= fun () ->
            update_loop ()
          ) else
            Lwt.return_unit
        in
        Lwt.async update_loop
      );

      (* Open file for writing *)
      Lwt_io.with_file
        ~mode:Lwt_io.Output
        ~flags:[Unix.O_WRONLY; Unix.O_CREAT; Unix.O_APPEND]
        final_output
        (fun oc ->
          let rec download_chunks current_pos =
            if !should_quit then (
              Lwt_unix.disable_signal_handler sigint_handler;
              Tui.close_tui term_opt >>= fun () ->
              Logs.info (fun m -> m "Download cancelled by user");
              exit 130
            ) else if current_pos >= info.total_size then
              Lwt.return_unit
            else
              let end_pos = Int64.add current_pos (Int64.of_int chunk_size)
                            |> Int64.sub (Int64.of_int 1)
                            |> Int64.min (Int64.sub info.total_size 1L) in

              download_range_to_channel url current_pos end_pos oc >>= function
              | None ->
                Logs.err (fun m -> m "Failed to download chunk at offset %Ld" current_pos);
                Lwt.return_unit
              | Some bytes_written ->
                let new_pos = Int64.add current_pos (Int64.of_int bytes_written) in
                downloaded := new_pos;
                save_http_state final_output new_pos info.total_size >>= fun () ->

                let percent = Int64.to_float new_pos /. Int64.to_float info.total_size *. 100.0 in
                if not use_tui then
                  Logs.app (fun m -> m "Progress: %.1f%% (%Ld / %Ld bytes)"
                    percent new_pos info.total_size);

                download_chunks new_pos
          in

          download_chunks start_pos
        ) >>= fun () ->

      Lwt_unix.disable_signal_handler sigint_handler;
      Tui.close_tui term_opt >>= fun () ->
      delete_http_state final_output;
      Logs.info (fun m -> m "Download completed: %s" final_output);
      Lwt.return_unit
    ) else (
      (* Non-resumable download *)
      Logs.warn (fun m -> m "Server doesn't support ranges, downloading entire file at once");
      Lwt_io.with_file ~mode:Lwt_io.Output final_output (fun oc ->
        download_full_to_channel url oc
      ) >>= function
      | None ->
        Logs.err (fun m -> m "Download failed");
        Lwt.return_unit
      | Some () ->
        Logs.info (fun m -> m "Download completed: %s" final_output);
        Lwt.return_unit
    )
;;
