open Lwt.Infix

type t =
  { peers : Peers.t list
  ; peer_id : bytes
  ; info_hash : bytes
  ; piece_hashes : bytes array
  ; piece_length : int
  ; length : int
  }

let create_torrent peers peer_id info_hash piece_hashes piece_length length =
  { peers; peer_id; info_hash; piece_hashes; piece_length; length }
;;

let calculate_bounds_for_piece torrent index =
  let start = index * torrent.piece_length in
  let final = start + torrent.piece_length in
  let torrent_length = torrent.length in
  if final > torrent_length then start, torrent_length else start, final
;;

let calculate_piece_size torrent index =
  let start, final = calculate_bounds_for_piece torrent index in
  final - start
;;

type piece_work =
  { index : int
  ; length : int
  ; hash : bytes
  }

type piece_result =
  { index : int
  ; length : int
  ; buf : bytes
  }

type state =
  { requested : int ref
  ; downloaded : int ref
  ; backlog : int ref
  ; buf : bytes
  }

(* Last block might be shorter than the typical block *)
let calculate_block_size piece_length requested =
  if piece_length - requested < Constants.block_len
  then piece_length - requested
  else Constants.block_len
;;

let read_message (client : Client.t) (pw : piece_work) _torrent state =
  Client.read client.flow >>= fun message ->
  match message with
  | None -> Lwt.return_unit
  | Some msg ->
    (* let () = *)
    (*   Logs.debug (fun m -> m "There is msg %s" (Message.to_string message)) *)
    (* in *)
    (match Stdint.Uint8.to_int msg.id with
     | 0 ->
       client.choked := true;
       Lwt.return_unit
     | 1 ->
       client.choked := false;
       Lwt.return_unit
     | 4 ->
       (match Message.parse_have msg with
        | Error (`Error _e) -> Lwt.return_unit
        | Ok index ->
          Bitfield.set_piece client.bitfield index;
          Lwt.return_unit)
     | 7 ->
       (match Message.parse_piece pw.index state.buf msg with
        | Error (`Error _e) -> Lwt.return_unit
        | Ok buf_len ->
          state.backlog := !(state.backlog) - 1;
          state.downloaded := !(state.downloaded) + buf_len;
          Lwt.return_unit)
     | _ -> Lwt.return_unit)
;;

let rec request (client : Client.t) state (pw : piece_work) =
  if !(state.backlog) < 5 && !(state.requested) < pw.length
  then (
    let block_size = calculate_block_size pw.length !(state.requested) in
    Client.send_request client pw.index !(state.requested) block_size
    >>= fun () ->
    state.requested := !(state.requested) + block_size;
    state.backlog := !(state.backlog) + 1;
    request client state pw)
  else Lwt.return_unit
;;

let rec download_piece (client : Client.t) (pw : piece_work) torrent state =
  if !(state.downloaded) < pw.length
  then (
    (if not !(client.choked) then request client state pw else Lwt.return_unit)
    >>= fun () ->
    read_message client pw torrent state >>= fun () ->
    download_piece client pw torrent state)
  else Lwt.return state.buf
;;

let with_timeout timeout f =
  Lwt.pick
    [ f ()
    ; (Lwt_unix.sleep timeout >>= fun () -> Lwt.fail (Failure "Timeout"))
    ]
;;

let try_download_piece timeout (client : Client.t) (pw : piece_work) torrent =
  let state =
    { requested = ref 0
    ; downloaded = ref 0
    ; backlog = ref 0
    ; buf = Bytes.create pw.length
    }
  in
  Logs.info (fun m ->
    m
      "Try to download piece %d of %d"
      (pw.index + 1)
      (Array.length torrent.piece_hashes));
  with_timeout timeout (fun () ->
    download_piece client pw torrent state >>= fun piece_buf ->
    Logs.info (fun m ->
      m
        "Downloaded piece (%d of %d) with size %d"
        (pw.index + 1)
        (Array.length torrent.piece_hashes)
        (Bytes.length piece_buf));
    Lwt.return piece_buf)
;;

let check_integrity (pw : piece_work) buf =
  let hash = Sha1.digest buf in
  if hash = pw.hash
  then Result.ok ()
  else Result.Error (`Error "Hash mismatch")
;;

let download_piece_task
  timeout
  client
  torrent
  (push_work : piece_work option -> unit)
  (push_result : piece_result option -> unit)
  pw
  =
  try_download_piece timeout client pw torrent >>= fun piece_buf ->
  let integrity_result = check_integrity pw piece_buf in
  match integrity_result with
  | Error (`Error e) ->
    Logs.debug (fun m -> m "Try send pw \n");
    push_work (Some pw);
    Logs.debug (fun m -> m "sent pw \n");
    Logs.err (fun m -> m "Error: %s" e);
    Lwt.return_unit
  | Ok () ->
    Log.debug "Integrity of piece %d is ok." pw.index;
    let piece_result =
      { buf = piece_buf; length = pw.length; index = pw.index }
    in
    Log.debug "Try send result %d." pw.index;
    push_result (Some piece_result);
    Log.debug "Sent result %d." pw.index;
    Client.send_have client pw.index
;;

let rec download_torrent_worker
  timeout
  (torrent : t)
  (client : Client.t)
  (work_stream : piece_work Lwt_stream.t)
  (push_work : piece_work option -> unit)
  (push_result : piece_result option -> unit)
  =
  Logs.debug (fun m -> m "Try receive pw \n");
  Lwt_stream.get work_stream >>= fun pw_opt ->
  match pw_opt with
  | None -> Lwt.return_unit
  | Some pw ->
    Lwt.catch
      (fun () ->
        Logs.debug (fun m ->
          m "With peer_id %s\n" (client.peer.ip |> Ipaddr.V4.to_string));
        if not (Bitfield.has_piece client.bitfield pw.index)
        then (
          Logs.debug (fun m -> m "Try send pw \n");
          push_work (Some pw);
          Logs.debug (fun m -> m "sent pw \n");
          Logs.debug (fun m -> m "This client does NOT have piece %d\n" pw.index);
          Lwt.return_unit)
        else
          download_piece_task
            timeout
            client
            torrent
            push_work
            push_result
            pw)
      (fun ex ->
        Logs.debug (fun m -> m "Try send pw \n");
        push_work (Some pw);
        Logs.debug (fun m -> m "sent pw \n");
        Logs.err (fun m -> m "%a" Fmt.exn ex);
        Lwt.return_unit)
    >>= fun () ->
    download_torrent_worker timeout torrent client work_stream push_work push_result
;;

let connect_and_download_torrent
  timeout
  connect_timeout
  torrent
  peer
  (work_stream : piece_work Lwt_stream.t)
  (push_work : piece_work option -> unit)
  (push_result : piece_result option -> unit)
  =
  with_timeout connect_timeout (fun () -> Client.connect peer torrent.info_hash torrent.peer_id)
  >>= fun client_result ->
  let client = Result.get_ok client_result in
  Logs.debug (fun m ->
    m "Completed handshake with %s\n" (Ipaddr.V4.to_string peer.ip));
  Client.send_unchoke client >>= fun () ->
  Client.send_interested client >>= fun () ->
  download_torrent_worker timeout torrent client work_stream push_work push_result
;;

let rec download_torrent_by_peers
  timeout
  connect_timeout
  torrent
  (peers : Peers.t list)
  (work_stream : piece_work Lwt_stream.t)
  (push_work : piece_work option -> unit)
  (push_result : piece_result option -> unit)
  =
  match peers with
  | [] -> Lwt.return_unit
  | peers_head :: peers_tail ->
    Lwt.async (fun () ->
      Lwt.catch
        (fun () ->
          connect_and_download_torrent
            timeout
            connect_timeout
            torrent
            peers_head
            work_stream
            push_work
            push_result)
        (fun _ex -> Lwt.return_unit));
    download_torrent_by_peers
      timeout
      connect_timeout
      torrent
      peers_tail
      work_stream
      push_work
      push_result
;;

(* State management for resume/pause functionality *)

let state_file_path output_file = output_file ^ ".tornado-state"

let save_state output_file info_hash completed_pieces =
  let state_dict =
    Bencode.Dict
      [ "info_hash", Bencode.String (Bytes.to_string info_hash)
      ; "completed_pieces", Bencode.String (Bytes.to_string completed_pieces)
      ; "total_pieces", Bencode.Integer (Int64.of_int (Bytes.length completed_pieces))
      ]
  in
  let encoded_bytes = Bencode_streaming.Encode.to_bytes state_dict in
  let encoded = Bytes.to_string encoded_bytes in
  let path = state_file_path output_file in
  Lwt_io.with_file
    ~mode:Lwt_io.Output
    ~flags:[ Unix.O_WRONLY; Unix.O_CREAT; Unix.O_TRUNC ]
    ~perm:0o644
    path
    (fun oc -> Lwt_io.write oc encoded)
;;

let load_state output_file info_hash total_pieces =
  let path = state_file_path output_file in
  if not (Sys.file_exists path) then
    Lwt.return None
  else
    Lwt.catch
      (fun () ->
        Lwt_io.with_file ~mode:Lwt_io.Input path (fun ic ->
          Lwt_io.read ic >>= fun content ->
          let state = Bencode.decode (`String content) in
          let saved_hash = Bencode_utils.bencode_to_string state "info_hash" in
          let completed_str = Bencode_utils.bencode_to_string state "completed_pieces" in
          let saved_total = Bencode_utils.bencode_to_int state "total_pieces" in
          match saved_hash, completed_str, saved_total with
          | Some h, Some c, Some t when
              Bytes.to_string info_hash = h &&
              Int64.to_int t = total_pieces ->
            let completed = Bytes.of_string c in
            if Bytes.length completed = total_pieces then
              Lwt.return (Some completed)
            else
              Lwt.return None
          | _ -> Lwt.return None))
      (fun _ex -> Lwt.return None)
;;

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

let download_torrent timeout connect_timeout use_tui (torrent : t) file_name completed_pieces =
  let open Lwt.Infix in
  let final_buf = Bytes.create torrent.length in
  let pieces_hashes_len = Array.length torrent.piece_hashes in

  (* Create work and result streams with explicit types *)
  let work_stream, work_push =
    (Lwt_stream.create () : piece_work Lwt_stream.t * (piece_work option -> unit)) in
  let result_stream, result_push =
    (Lwt_stream.create () : piece_result Lwt_stream.t * (piece_result option -> unit)) in

  (* Initialize work queue - skip completed pieces *)
  Lwt.async (fun () ->
    Lwt_list.iter_s
      (fun index ->
        if Bytes.get completed_pieces index = '\000' then (
          let hash = torrent.piece_hashes.(index) in
          let length = calculate_piece_size torrent index in
          let pw = { index; length; hash } in
          work_push (Some pw);
          Lwt.return_unit)
        else
          Lwt.return_unit)
      (List.init pieces_hashes_len (fun i -> i)));

  (* Start peer download tasks *)
  Lwt.async (fun () ->
    download_torrent_by_peers
      timeout
      connect_timeout
      torrent
      torrent.peers
      work_stream
      work_push
      result_push);

  (* Count already completed pieces *)
  let initial_completed = ref 0 in
  for i = 0 to pieces_hashes_len - 1 do
    if Bytes.get completed_pieces i <> '\000' then
      initial_completed := !initial_completed + 1
  done;

  if !initial_completed > 0 then
    Logs.info (fun m -> m "Resuming download: %d/%d pieces already completed"
      !initial_completed pieces_hashes_len);

  (* Setup TUI if enabled *)
  (if use_tui then Tui.create_tui file_name pieces_hashes_len
   else Lwt.return (None, ref { Tui.name = file_name; total_pieces = pieces_hashes_len;
                                 completed_pieces = 0; peers_count = 0; download_speed = 0.0 }))
  >>= fun (term_opt, stats_ref) ->

  (* Collect results *)
  let done_pieces = ref !initial_completed in
  let start_time = Unix.gettimeofday () in
  let bytes_downloaded = ref 0 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 !done_pieces < pieces_hashes_len && not !should_quit then (
        let elapsed = Unix.gettimeofday () -. start_time in
        let speed = if elapsed > 0.0 then float_of_int !bytes_downloaded /. elapsed else 0.0 in
        Tui.update_stats stats_ref ~completed:!done_pieces
          ~peers:(List.length torrent.peers) ~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
  );

  (* Periodic cancellation checker *)
  let rec cancellation_checker () =
    Lwt_unix.sleep 0.1 >>= fun () ->
    if !should_quit then
      Lwt.return_unit
    else
      cancellation_checker ()
  in

  (* Collect all pieces asynchronously *)
  let rec collect_pieces () =
    if !should_quit then
      Lwt.return_unit
    else if !done_pieces < pieces_hashes_len then
      (* Race between getting next piece and cancellation check *)
      Lwt.pick [
        (Lwt_stream.next result_stream >>= fun piece_result ->
         let length = calculate_piece_size torrent piece_result.index in
         let start, _ = calculate_bounds_for_piece torrent piece_result.index in
         done_pieces := !done_pieces + 1;
         bytes_downloaded := !bytes_downloaded + length;
         Bytes.set completed_pieces piece_result.index '\001';
         Bytes.blit piece_result.buf 0 final_buf start length;
         (* Save state after each piece *)
         save_state file_name torrent.info_hash completed_pieces >>= fun () ->
         Lwt.return `Continue);
        (cancellation_checker () >>= fun () -> Lwt.return `Cancelled)
      ] >>= function
      | `Continue -> collect_pieces ()
      | `Cancelled -> Lwt.return_unit
    else
      Lwt.return_unit
  in

  collect_pieces () >>= fun () ->
  Lwt_unix.disable_signal_handler sigint_handler;
  Tui.close_tui term_opt >>= fun () ->
  if !should_quit then (
    Logs.info (fun m -> m "Download cancelled by user");
    exit 130
  ) else
    Lwt.return final_buf
;;

let write_multifile base_dir files final_buf =
  (* Create base directory if needed *)
  (try Unix.mkdir base_dir 0o755 with Unix.Unix_error (Unix.EEXIST, _, _) -> ());

  let rec write_files offset = function
    | [] -> Lwt.return_unit
    | file :: rest ->
      let file_path = String.concat "/" (base_dir :: file.Torrent_file.path) in
      let file_dir = Filename.dirname file_path in
      (* Create directory structure *)
      let rec mkdir_p dir =
        if not (Sys.file_exists dir) then (
          mkdir_p (Filename.dirname dir);
          Unix.mkdir dir 0o755
        )
      in
      mkdir_p file_dir;

      let file_len = Int64.to_int file.Torrent_file.length in
      Lwt_io.with_file
        ~mode:Lwt_io.Output
        ~flags:[ Unix.O_WRONLY; Unix.O_CREAT; Unix.O_TRUNC ]
        ~perm:0o644
        file_path
        (fun oc ->
          Lwt_io.write_from_exactly oc final_buf offset file_len) >>= fun () ->
      write_files (offset + file_len) rest
  in
  write_files 0 files
;;

let download_file output_file torrent_file timeout connect_timeout use_tui =
  let open Torrent_file in
  let base_name =
    match output_file, torrent_file.name with
    | Some file, _ -> file
    | _, Some file -> file
    | _, _ -> "torrent_download"
  in
  (* Get Peers *)
  let random_peer = Bytes.create 20 in
  let client_port = 6881 in
  let uri = build_tracker_url torrent_file random_peer client_port in
  let total_len = total_length torrent_file |> Int64.to_int in
  let left = Stdint.Int64.of_int total_len in
  Peers.request_peers uri torrent_file.info_hash random_peer client_port left >>= fun peers ->
  Logs.debug (fun m -> m "Got %d peers\n" (List.length peers));
  (* Download *)
  let pieces_count = Array.length torrent_file.piece_hashes in

  (* Load resume state if exists *)
  load_state base_name torrent_file.info_hash pieces_count >>= fun state_opt ->
  let completed_pieces = match state_opt with
    | Some state -> state
    | None -> Bytes.make pieces_count '\000'
  in

  let torrent =
    create_torrent
      peers
      random_peer
      torrent_file.info_hash
      torrent_file.piece_hashes
      (torrent_file.piece_length |> Int64.to_int)
      total_len
  in
  download_torrent timeout connect_timeout use_tui torrent base_name completed_pieces >>= fun final_buf ->
  (* Write File(s) *)
  (match torrent_file.file_mode with
  | SingleFile _ ->
    Lwt_io.with_file
      ~mode:Lwt_io.Output
      ~flags:[ Unix.O_WRONLY; Unix.O_CREAT; Unix.O_TRUNC ]
      ~perm:0o644
      base_name
      (fun oc -> Lwt_io.write_from_exactly oc final_buf 0 (Bytes.length final_buf))
  | MultiFile files ->
    write_multifile base_name files final_buf) >>= fun () ->
  (* Clean up state file after successful download *)
  remove_state_file base_name;
  (* Announce completion to tracker for seeding *)
  Logs.info (fun m -> m "Download complete. Announcing to tracker as seeder...");
  let completed_uri = build_completed_url torrent_file random_peer client_port in
  Peers.announce_completed completed_uri torrent_file.info_hash random_peer client_port >>= fun () ->
  Logs.info (fun m -> m "Seeding announcement sent to tracker.");
  Lwt.return_unit
;;
