open Lwt.Infix

module Transport (R : Mirage_random.S) (C : Mirage_clock.MCLOCK) (Time : Mirage_time.S) = struct
  type +'a io = 'a Lwt.t
  type io_addr = Ipaddr.V4.t * int
  module Dispatcher = Dispatcher.Make(R)(C)(Time)
  type stack = Dispatcher.t * (src_port:int -> dst:Ipaddr.V4.t -> dst_port:int -> Cstruct.t -> (unit, [ `Msg of string ]) result Lwt.t) * (Udp_packet.t * Cstruct.t) Lwt_mvar.t

  module IM = Map.Make(Int)

  type t = {
    protocol : Dns.proto ;
    nameserver : io_addr ;
    stack : stack ;
    timeout_ns : int64 ;
    mutable requests : Cstruct.t Lwt_condition.t IM.t ;
  }
  type context = t

  let nameservers { protocol ; nameserver ; _ } = protocol, [ nameserver ]
  let rng = R.generate ?g:None
  let clock = C.elapsed_ns

  let rec read t =
    let _, _, answer = t.stack in
    Lwt_mvar.take answer >>= fun (_, data) ->
    if Cstruct.length data > 2 then begin
      match IM.find_opt (Cstruct.BE.get_uint16 data 0) t.requests with
      | Some cond -> Lwt_condition.broadcast cond data
      | None -> ()
    end;
    read t

  let create ?nameservers ~timeout stack =
    let protocol, nameserver = match nameservers with
      | None | Some (_, []) -> invalid_arg "no nameserver found"
      | Some (proto, ns :: _) -> proto, ns
    in
    let t =
      { protocol ; nameserver ; stack ; timeout_ns = timeout ; requests = IM.empty }
    in
    Lwt.async (fun () -> read t);
    t

  let with_timeout timeout_ns f =
    let timeout = Time.sleep_ns timeout_ns >|= fun () -> Error (`Msg "DNS request timeout") in
    Lwt.pick [ f ; timeout ]

  let connect (t : t) = Lwt.return (Ok (t.protocol, t))

  let send_recv (ctx : context) buf : (Cstruct.t, [> `Msg of string ]) result Lwt.t =
    let dst, dst_port = ctx.nameserver in
    let router, send_udp, _ = ctx.stack in
    let src_port, evict =
      My_nat.free_udp_port router.nat ~src:router.config.our_ip ~dst ~dst_port:53
    in
    let id = Cstruct.BE.get_uint16 buf 0 in
    with_timeout ctx.timeout_ns
      (let cond = Lwt_condition.create () in
       ctx.requests <- IM.add id cond ctx.requests;
       (send_udp ~src_port ~dst ~dst_port buf >|= Rresult.R.open_error_msg) >>= function
       | Ok () -> Lwt_condition.wait cond >|= fun dns_response -> Ok dns_response
       | Error _ as e -> Lwt.return e) >|= fun result ->
    ctx.requests <- IM.remove id ctx.requests;
    evict ();
    result

  let close _ = Lwt.return_unit

  let bind = Lwt.bind

  let lift = Lwt.return
end