open Stdint
open Ipaddr

type t =
  { ip : V4.t
  ; port : int
  }

let create peers_bin =
  let peer_size = 6 in
  let peers_bin_length = Bytes.length peers_bin in
  let num_peers = peers_bin_length / peer_size in
  if peers_bin_length mod peer_size <> 0
  then failwith "Received malformed peers";
  List.init num_peers (fun i ->
    let ip_size = 4 in
    let offset = i * peer_size in
    let ip_bin = Bytes.sub peers_bin offset ip_size in
    let port =
      Uint16.of_bytes_big_endian peers_bin (offset + ip_size) |> Uint16.to_int
    in
    let ip_array =
      Array.init ip_size (fun j ->
        Uint8.of_bytes_big_endian ip_bin j |> Uint8.to_int)
    in
    let ip =
      Ipaddr.V4.make ip_array.(0) ip_array.(1) ip_array.(2) ip_array.(3)
    in
    { ip; port })
;;

(* UDP Tracker Protocol (BEP 15) *)

let udp_protocol_id = Int64.of_string "0x41727101980"

let write_int64_be buf offset value =
  let bytes = Bytes.create 8 in
  Stdint.Int64.to_bytes_big_endian value bytes 0;
  Bytes.blit bytes 0 buf offset 8
;;

let write_int32_be buf offset value =
  let bytes = Bytes.create 4 in
  Stdint.Int32.to_bytes_big_endian value bytes 0;
  Bytes.blit bytes 0 buf offset 4
;;

let write_int16_be buf offset value =
  let bytes = Bytes.create 2 in
  Stdint.Uint16.to_bytes_big_endian value bytes 0;
  Bytes.blit bytes 0 buf offset 2
;;

let read_int64_be buf offset =
  Stdint.Int64.of_bytes_big_endian buf offset
;;

let read_int32_be buf offset =
  Stdint.Int32.of_bytes_big_endian buf offset
;;

let request_peers_udp host port info_hash peer_id client_port left =
  let open Lwt.Infix in
  let sock = Lwt_unix.socket Unix.PF_INET Unix.SOCK_DGRAM 0 in

  (* Resolve host *)
  Lwt_unix.gethostbyname host >>= fun host_entry ->
  let addr = host_entry.Unix.h_addr_list.(0) in
  let sockaddr = Unix.ADDR_INET (addr, port) in

  (* Connect request *)
  let transaction_id = Random.int32 Int32.max_int in
  let connect_req = Bytes.create 16 in
  write_int64_be connect_req 0 udp_protocol_id;
  write_int32_be connect_req 8 (Stdint.Int32.zero);
  write_int32_be connect_req 12 transaction_id;

  Lwt_unix.sendto sock connect_req 0 16 [] sockaddr >>= fun _ ->

  (* Connect response *)
  let connect_resp = Bytes.create 16 in
  Lwt_unix.recvfrom sock connect_resp 0 16 [] >>= fun (len, _) ->

  if len < 16 then
    Lwt.fail (Failure "UDP tracker: invalid connect response")
  else
    let resp_action = read_int32_be connect_resp 0 in
    let resp_transaction = read_int32_be connect_resp 4 in
    if resp_action <> Stdint.Int32.zero || resp_transaction <> transaction_id then
      Lwt.fail (Failure "UDP tracker: connect response mismatch")
    else
      let connection_id = read_int64_be connect_resp 8 in

      (* Announce request *)
      let announce_transaction = Random.int32 Int32.max_int in
      let announce_req = Bytes.create 98 in
      write_int64_be announce_req 0 connection_id;
      write_int32_be announce_req 8 (Stdint.Int32.one); (* action=1 for announce *)
      write_int32_be announce_req 12 announce_transaction;
      Bytes.blit info_hash 0 announce_req 16 20;
      Bytes.blit peer_id 0 announce_req 36 20;
      write_int64_be announce_req 56 (Stdint.Int64.zero); (* downloaded *)
      write_int64_be announce_req 64 left; (* left *)
      write_int64_be announce_req 72 (Stdint.Int64.zero); (* uploaded *)
      write_int32_be announce_req 80 (Stdint.Int32.zero); (* event=0 (none) *)
      write_int32_be announce_req 84 (Stdint.Int32.zero); (* ip=0 (default) *)
      write_int32_be announce_req 88 (Random.int32 Int32.max_int); (* key *)
      write_int32_be announce_req 92 (Stdint.Int32.of_int (-1)); (* num_want=-1 *)
      write_int16_be announce_req 96 (Stdint.Uint16.of_int client_port);

      Lwt_unix.sendto sock announce_req 0 98 [] sockaddr >>= fun _ ->

      (* Announce response *)
      let announce_resp = Bytes.create 65536 in
      Lwt_unix.recvfrom sock announce_resp 0 65536 [] >>= fun (len, _) ->

      Lwt_unix.close sock >>= fun () ->

      if len < 20 then
        Lwt.fail (Failure "UDP tracker: invalid announce response")
      else
        let resp_action = read_int32_be announce_resp 0 in
        let resp_transaction = read_int32_be announce_resp 4 in
        if resp_action <> Stdint.Int32.one || resp_transaction <> announce_transaction then
          Lwt.fail (Failure "UDP tracker: announce response mismatch")
        else
          let peers_data = Bytes.sub announce_resp 20 (len - 20) in
          let peers = create peers_data in
          Lwt.return peers
;;

let request_peers_http uri =
  let open Lwt.Infix in
  Cohttp_lwt_unix.Client.get uri >>= fun (resp, body) ->
  Cohttp_lwt.Body.to_string body >>= fun body_string ->
  let tracker_bencode = Bencode.decode (`String body_string) in
  let peers_string =
    Bencode_utils.bencode_to_string tracker_bencode "peers" |> Option.get
  in
  let peers_bytes = peers_string |> Bytes.of_string in
  let peers = create peers_bytes in
  Lwt.return peers
;;

let request_peers uri info_hash peer_id client_port left =
  let scheme = Uri.scheme uri in
  match scheme with
  | Some "udp" ->
    let host = Uri.host uri |> Option.get in
    let port = Uri.port uri |> Option.value ~default:80 in
    request_peers_udp host port info_hash peer_id client_port left
  | _ -> request_peers_http uri
;;

let to_string peer = Printf.sprintf "%s:%d" (V4.to_string peer.ip) peer.port

(* Announce completion to tracker for seeding *)

let announce_completed_udp host port info_hash peer_id client_port =
  let open Lwt.Infix in
  let sock = Lwt_unix.socket Unix.PF_INET Unix.SOCK_DGRAM 0 in

  Lwt_unix.gethostbyname host >>= fun host_entry ->
  let addr = host_entry.Unix.h_addr_list.(0) in
  let sockaddr = Unix.ADDR_INET (addr, port) in

  (* Connect request *)
  let transaction_id = Random.int32 Int32.max_int in
  let connect_req = Bytes.create 16 in
  write_int64_be connect_req 0 udp_protocol_id;
  write_int32_be connect_req 8 (Stdint.Int32.zero);
  write_int32_be connect_req 12 transaction_id;

  Lwt_unix.sendto sock connect_req 0 16 [] sockaddr >>= fun _ ->
  let connect_resp = Bytes.create 16 in
  Lwt_unix.recvfrom sock connect_resp 0 16 [] >>= fun (len, _) ->

  if len < 16 then
    Lwt_unix.close sock >>= fun () ->
    Lwt.fail (Failure "UDP tracker: invalid connect response")
  else
    let resp_action = read_int32_be connect_resp 0 in
    let resp_transaction = read_int32_be connect_resp 4 in
    if resp_action <> Stdint.Int32.zero || resp_transaction <> transaction_id then
      Lwt_unix.close sock >>= fun () ->
      Lwt.fail (Failure "UDP tracker: connect response mismatch")
    else
      let connection_id = read_int64_be connect_resp 8 in

      (* Announce with event=completed (1) *)
      let announce_transaction = Random.int32 Int32.max_int in
      let announce_req = Bytes.create 98 in
      write_int64_be announce_req 0 connection_id;
      write_int32_be announce_req 8 (Stdint.Int32.one);
      write_int32_be announce_req 12 announce_transaction;
      Bytes.blit info_hash 0 announce_req 16 20;
      Bytes.blit peer_id 0 announce_req 36 20;
      write_int64_be announce_req 56 (Stdint.Int64.zero); (* downloaded=0 *)
      write_int64_be announce_req 64 (Stdint.Int64.zero); (* left=0 (complete) *)
      write_int64_be announce_req 72 (Stdint.Int64.zero); (* uploaded=0 *)
      write_int32_be announce_req 80 (Stdint.Int32.one); (* event=1 (completed) *)
      write_int32_be announce_req 84 (Stdint.Int32.zero);
      write_int32_be announce_req 88 (Random.int32 Int32.max_int);
      write_int32_be announce_req 92 (Stdint.Int32.of_int 0); (* num_want=0 (no more peers needed) *)
      write_int16_be announce_req 96 (Stdint.Uint16.of_int client_port);

      Lwt_unix.sendto sock announce_req 0 98 [] sockaddr >>= fun _ ->

      let announce_resp = Bytes.create 65536 in
      Lwt_unix.recvfrom sock announce_resp 0 65536 [] >>= fun (_len, _) ->

      Lwt_unix.close sock >>= fun () ->
      Lwt.return_unit
;;

let announce_completed_http uri =
  let open Lwt.Infix in
  Lwt.catch
    (fun () ->
      Cohttp_lwt_unix.Client.get uri >>= fun (_resp, body) ->
      Cohttp_lwt.Body.drain_body body >>= fun () ->
      Lwt.return_unit)
    (fun _ex -> Lwt.return_unit)
;;

let announce_completed uri info_hash peer_id client_port =
  let scheme = Uri.scheme uri in
  match scheme with
  | Some "udp" ->
    let host = Uri.host uri |> Option.get in
    let port = Uri.port uri |> Option.value ~default:80 in
    announce_completed_udp host port info_hash peer_id client_port
  | _ -> announce_completed_http uri
;;
