open Stdint
open Lwt.Infix

type t =
  { flow : Lwt_io.input_channel * Lwt_io.output_channel
  ; choked : bool ref
  ; bitfield : Bitfield.t
  ; peer : Peers.t
  ; info_hash : bytes
  ; peer_id : bytes
  }

let read flow =
  Tcp.Client.read_bytes flow 4 >>= fun length_buf ->
  let length = Uint32.to_int (Uint32.of_bytes_big_endian length_buf 0) in
  Tcp.Client.read_bytes flow length >>= fun msg_bytes ->
  Lwt.return (Message.read length_buf msg_bytes)
;;

let send_request t index start length =
  let msg = Message.format_request index start length in
  let msg_bytes = Message.serialize (Some msg) in
  Tcp.Client.write_bytes t.flow msg_bytes
;;

let send_interested t =
  let msg : Message.t =
    { id = Message.id_of_message_type Msg_interested
    ; payload = Bytes.create 0
    }
  in
  let msg_bytes = Message.serialize (Some msg) in
  Tcp.Client.write_bytes t.flow msg_bytes
;;

let send_not_interested flow =
  let msg : Message.t =
    { id = Message.id_of_message_type Msg_not_interested
    ; payload = Bytes.create 0
    }
  in
  let msg_bytes = Message.serialize (Some msg) in
  Tcp.Client.write_bytes flow msg_bytes
;;

let send_unchoke t =
  let msg : Message.t =
    { id = Message.id_of_message_type Msg_unchoke; payload = Bytes.create 0 }
  in
  let msg_bytes = Message.serialize (Some msg) in
  Tcp.Client.write_bytes t.flow msg_bytes
;;

let send_have client index =
  let msg = Message.format_have (Uint32.of_int index) in
  let msg_bytes = Message.serialize (Some msg) in
  Tcp.Client.write_bytes client.flow msg_bytes
;;

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

let complete_handshake flow info_hash peerID =
  with_timeout 3.0 (fun () ->
    let handshake = Handshake.create info_hash peerID in
    let handshake_bytes = Handshake.serialize_to_bytes handshake in
    Tcp.Client.write_bytes flow handshake_bytes >>= fun () ->
    Tcp.Client.read_bytes flow 1 >>= fun pstrlen_buf ->
    let pstrlen = Bytes.get_uint8 pstrlen_buf 0 in
    Tcp.Client.read_bytes flow (pstrlen + 48) >>= fun handshake_bytes ->
    let handshake_result = Handshake.read pstrlen handshake_bytes in
    match handshake_result with
    | Ok h when h.info_hash <> info_hash ->
      Lwt.return (Result.error `Info_hash_is_not_equal)
    | Ok h -> Lwt.return (Result.ok h)
    | Error e -> Lwt.return (Result.error e))
;;

let recv_bitfield flow =
  with_timeout 6.0 (fun () ->
    read flow >>= fun msg_result ->
    match msg_result with
    | None ->
      Lwt.return
        (Result.error
          (`Error
            (Printf.sprintf
               "Expected bitfield but got %s"
               (Message.to_string msg_result))))
    | Some m when m.id <> Message.id_of_message_type Msg_bitfield ->
      Lwt.return
        (Result.error
          (`Error
            (Printf.sprintf
               "Expected bitfield but got ID %d"
               (Uint8.to_int m.id))))
    | Some m -> Lwt.return (Result.ok m.payload))
;;

let connect (peer : Peers.t) info_hash peer_id =
  (* Log.debug "Connecting to %s" (Peers.show peer); *)
  Tcp.Client.open_connection ~host:peer.ip ~port:peer.port >>= fun flow ->
  (* Logs.debug (fun m -> m "Connected to %s" (Peers.show peer)); *)
  complete_handshake flow info_hash peer_id >>= fun complete_handshake_result ->
  match complete_handshake_result with
  | Error e -> Lwt.return (Result.error e)
  | Ok _ ->
    let () = Logs.debug (fun m -> m "Completed handshake") in
    recv_bitfield flow >>= fun recv_bitfield_result ->
    (match recv_bitfield_result with
     | Error e -> Lwt.return (Result.error e)
     | Ok bitfield ->
       Lwt.return
         (Result.ok
           { flow
           ; choked = ref true
           ; bitfield = Bitfield.of_bytes bitfield
           ; peer
           ; info_hash
           ; peer_id
           }))
;;
