open Bencode_utils

(* File information for multi-file torrents *)
type file_info =
  { path : string list  (* path components *)
  ; length : int64
  }

(* Torrent can be single-file or multi-file *)
type file_mode =
  | SingleFile of int64  (* total length *)
  | MultiFile of file_info list

(* TORRENT *)
type t =
  { announce : string option
  ; info_hash : bytes
  ; piece_hashes : bytes array
  ; piece_length : int64
  ; file_mode : file_mode
  ; name : string option
  }

(* Get total length regardless of file mode *)
let total_length t =
  match t.file_mode with
  | SingleFile len -> len
  | MultiFile files -> List.fold_left (fun acc f -> Int64.add acc f.length) 0L files
;;

let parse_file_list files_bencode =
  match Bencode.as_list files_bencode with
  | None -> failwith "Files field is not a list"
  | Some file_list ->
    List.map (fun file_dict ->
      let length =
        match bencode_to_int file_dict "length" with
        | None -> failwith "File entry missing 'length' field"
        | Some l -> l
      in
      let path =
        match Bencode.dict_get file_dict "path" with
        | None -> failwith "File entry missing 'path' field"
        | Some path_bencode ->
          (match Bencode.as_list path_bencode with
           | None -> failwith "File path is not a list"
           | Some path_list ->
             List.map (fun p ->
               match Bencode.as_string p with
               | None -> failwith "Path component is not a string"
               | Some s -> s
             ) path_list)
      in
      { path; length }
    ) file_list
;;

let create_with_beencode bencode_root =
  let announce = bencode_to_string bencode_root "announce" in
  let info_beencode =
    match Bencode.dict_get bencode_root "info" with
    | None -> failwith "Torrent file is missing required 'info' dictionary"
    | Some x -> x
  in
  let piece_length =
    match bencode_to_int info_beencode "piece length" with
    | None -> failwith "Torrent file is missing required 'piece length' field in info dictionary"
    | Some x -> x
  in
  let pieces =
    match Bencode.dict_get info_beencode "pieces" with
    | None -> failwith "Torrent file is missing required 'pieces' field in info dictionary"
    | Some x -> x
  in
  let name = bencode_to_string info_beencode "name" in
  let info_hash = sha1_of_bencode info_beencode in
  let piece_hashes = split_piece_hashes pieces in

  (* Determine if single-file or multi-file *)
  let file_mode =
    match Bencode.dict_get info_beencode "files" with
    | Some files_bencode ->
      (* Multi-file torrent *)
      MultiFile (parse_file_list files_bencode)
    | None ->
      (* Single-file torrent *)
      (match bencode_to_int info_beencode "length" with
       | None -> failwith "Single-file torrent missing 'length' field"
       | Some len -> SingleFile len)
  in

  { announce; info_hash; piece_hashes; piece_length; file_mode; name }
;;

let open_file input_file =
  let bencode_file = Bencode.decode (`File_path input_file) in
  create_with_beencode bencode_file
;;

let build_tracker_url file peer_id port =
  let announce_url = match file.announce with
    | Some url -> url
    | None -> failwith "Torrent file does not contain an 'announce' URL. This torrent may use 'announce-list' which is not yet supported."
  in
  let query =
    [ "info_hash", [ Bytes.to_string file.info_hash ]
    ; "peer_id", [ Bytes.to_string peer_id ]; "port", [ Int.to_string port ]
    ; "uploaded", [ "0" ]; "downloaded", [ "0" ]; "compact", [ "1" ]
    ; "left", [ Int64.to_string (total_length file) ] ]
  in
  let uri = Uri.of_string announce_url in
  Uri.add_query_params uri query
;;

let build_completed_url file peer_id port =
  let announce_url = match file.announce with
    | Some url -> url
    | None -> failwith "Torrent file does not contain an 'announce' URL."
  in
  let query =
    [ "info_hash", [ Bytes.to_string file.info_hash ]
    ; "peer_id", [ Bytes.to_string peer_id ]
    ; "port", [ Int.to_string port ]
    ; "uploaded", [ "0" ]
    ; "downloaded", [ Int64.to_string (total_length file) ]
    ; "left", [ "0" ]
    ; "event", [ "completed" ]
    ]
  in
  let uri = Uri.of_string announce_url in
  Uri.add_query_params uri query
;;

