module type RUNTIME = sig
  type t

  val next_read_operation : t -> [ `Read | `Yield | `Close | `Upgrade ]
  (** [next_read_connection t] returns a value describing the next operation
      that the caller should conduit on behalf of the connection. *)

  val read : t -> Bigstringaf.t -> off:int -> len:int -> int
  (** [read t bigstring ~off ~len] reads bytes of input from the provided range
      of [bigstring] an returns the number of bytes consumed by the connection.
      {!read} should be called after {!next_read_operation} returns a [`Read]
      value an additional input is available for the connection to consume. *)

  val read_eof : t -> Bigstringaf.t -> off:int -> len:int -> int
  (** [read_eof t bigstring ~off ~len] reads bytes of input from the provided
      range of [bigstring] and returns the number of bytes consumed by the
      connection. {!read_eof} should be called after {!next_read_operation}
      returns a [`Read] and an EOF has been received from the communication
      channel. The connection will attempt to consume any buffered input and
      then shutdown the HTTP parser for the connection. *)

  val yield_reader : t -> (unit -> unit) -> unit
  (** [yield_reader t continue] registers with the connection to call [continue]
      when reading should resume. {!yield_reader} should be called after
      {!next_read_operation} returns a [`Yield] value. *)

  val next_write_operation :
    t ->
    [ `Write of Bigstringaf.t Faraday.iovec list
    | `Yield
    | `Close of int
    | `Upgrade ]
  (** [next_write_operation t] returns a value describing the next operation
      that the caller should conduct on behalf the connection. *)

  val report_write_result : t -> [ `Ok of int | `Closed ] -> unit
  (** [report_write_result t result] reports the result of the latest write
      attempt to the connection. {!report_write_result} should be called after a
      call to {!next_write_operation} that returns a [`Write buffer] value.

      - [`Ok n] indicates that the caller successfully wrote [n] bytes of output
        from the buffer that the caller was provided by {!next_write_operation}
        that returns a [`Write buffer] value.
      - [`Closed] indicates that the output destination will no longer accept
        bytes from the write processor. *)

  val yield_writer : t -> (unit -> unit) -> unit
  (** [yield_writer t continue] registers with the connection to call [continue]
      when writing should resume. {!yield_writer} should be called after
      {!next_write_operation} returns a [`Yield] value. *)

  val report_exn : t -> exn -> unit
  (** [report_exn t exn] reports that an error [exn] has been caught and that it
      has been attributed to [t]. Calling this function will switch [t] into an
      error state. Depending on the tate [t] is transitioning from, it may call
      its error handler before terminating the connection. *)

  val is_closed : t -> bool
  (** [is_closed t] is [true] if both the read and write processors have been
      shutdown. When this is the case {!next_read_operation} will return
      [`Close _] and {!next_write_operation} will return a [`Write _] until all
      buffered output has been flushed, at which point it will return [`Close].
  *)

  val shutdown : t -> unit
  (** [shutdown t] asks to shutdown the connection. *)
end

type 'conn runtime = (module RUNTIME with type t = 'conn)

exception Flow of string
exception Flow_write of string

let src = Logs.Src.create "paf-flow"

module Log_flow = (val Logs.src_log src : Logs.LOG)

module Make (Flow : Mirage_flow.S) = struct
  type flow = {
    flow : Flow.flow;
    queue : (char, Bigarray.int8_unsigned_elt) Ke.Rke.t;
    mutable rd_closed : bool;
    mutable wr_closed : bool;
  }

  let create flow =
    let queue = Ke.Rke.create ~capacity:0x1000 Bigarray.char in
    Lwt.return { flow; queue; rd_closed = false; wr_closed = false }

  let safely_close flow =
    if flow.rd_closed && flow.wr_closed
    then (
      Log_flow.debug (fun m -> m "Close the connection.") ;
      Flow.close flow.flow)
    else Lwt.return ()

  let blit src src_off dst dst_off len =
    let dst = Cstruct.of_bigarray ~off:dst_off ~len dst in
    Cstruct.blit src src_off dst 0 len

  open Lwt.Infix

  type eof = [ `Eof ]

  let recv flow ~report_error ~report_closed ~read ~read_eof =
    Ke.Rke.compress flow.queue ;
    Flow.read flow.flow >>= function
    | (Error _ | Ok #eof) as v ->
        flow.rd_closed <- true ;
        safely_close flow >>= fun () ->
        let _shift =
          match
            Ke.Rke.compress flow.queue ;
            Ke.Rke.N.peek flow.queue
          with
          | [] -> read_eof Bigstringaf.empty ~off:0 ~len:0
          | [ slice ] -> read_eof slice ~off:0 ~len:(Bigstringaf.length slice)
          | _ -> assert false
          (* XXX(dinosaure): impossible due to [compress]. *) in
        (match v with
        | Ok `Eof -> report_closed ()
        | Error err -> report_error err) ;
        Lwt.return `Closed
    | Ok (`Data v) ->
        let len = Cstruct.length v in
        Ke.Rke.N.push flow.queue ~blit ~length:Cstruct.length ~off:0 ~len v ;
        let[@warning "-8"] (slice :: _) = Ke.Rke.N.peek flow.queue in
        let shift = read slice ~off:0 ~len:(Bigstringaf.length slice) in
        Ke.Rke.N.shift_exn flow.queue shift ;
        Lwt.return `Continue

  let writev ~report_error flow iovecs =
    let iovecs =
      List.map
        (fun { Faraday.buffer; off; len } ->
          Cstruct.to_string (Cstruct.of_bigarray buffer ~off ~len) ~off:0 ~len)
        iovecs in
    let iovecs = List.map Cstruct.of_string iovecs in
    (* XXX(dinosaure): the copy is needed:
       1) [Mirage_flow.S] explicitly says that [write] takes the ownership on
          the given [Cstruct.t]
       2) [ocaml-h2] wants to keep the ownership on given [Faraday.iovec]s

       To protect one from the other, copying is necessary. *)
    Log_flow.debug (fun m ->
        m "Start to write %d byte(s)."
          (List.fold_left (fun acc cs -> Cstruct.length cs + acc) 0 iovecs)) ;
    Flow.writev flow.flow iovecs >>= function
    | Ok () ->
        Lwt.return
          (`Ok (List.fold_left (fun acc cs -> acc + Cstruct.length cs) 0 iovecs))
    | Error err ->
        Log_flow.err (fun m ->
            m "Got an error when we wrote something: %a." Flow.pp_write_error
              err) ;
        report_error err ;
        flow.wr_closed <- true ;
        safely_close flow >>= fun () -> Lwt.return `Closed

  let send ~report_error flow iovecs =
    if flow.wr_closed
    then safely_close flow >>= fun () -> Lwt.return `Closed
    else writev ~report_error flow iovecs

  let close flow =
    match (flow.rd_closed, flow.wr_closed) with
    | true, true -> Lwt.return_unit
    | _ ->
        flow.rd_closed <- true ;
        flow.wr_closed <- true ;
        Flow.close flow.flow
end

let src = Logs.Src.create "paf-server"

module Log_server = (val Logs.src_log src : Logs.LOG)

module Server (Flow : Mirage_flow.S) (Runtime : RUNTIME) : sig
  val server : Runtime.t -> Flow.flow -> unit Lwt.t
end = struct
  module Easy_flow = Make (Flow)
  open Lwt.Infix

  let to_flow_exception err : exn = Flow (Fmt.str "%a" Flow.pp_error err)

  let to_flow_write_exception err : exn =
    Flow_write (Fmt.str "%a" Flow.pp_write_error err)

  let server connection flow =
    Easy_flow.create flow >>= fun flow ->
    let rd_exit, notify_rd_exit = Lwt.wait () in
    let wr_exit, notify_wr_exit = Lwt.wait () in

    let rec rd_fiber () =
      let report_error err =
        Runtime.report_exn connection (to_flow_exception err) in
      let rec go () =
        Log_server.debug (fun m -> m "Compute next read operation.") ;
        match Runtime.next_read_operation connection with
        | `Upgrade -> failwith "Unimplemented"
        | `Read ->
            Log_server.debug (fun m -> m "next read operation: `read") ;
            Easy_flow.recv flow ~report_error ~report_closed:ignore
              ~read:(Runtime.read connection)
              ~read_eof:(Runtime.read_eof connection)
            >>= fun _ -> Lwt.pause () >>= go
        | `Yield ->
            Log_server.debug (fun m -> m "next read operation: `yield") ;
            Runtime.yield_reader connection rd_fiber ;
            Lwt.pause ()
        | `Close ->
            Log_server.debug (fun m -> m "next read operation: `close") ;
            Lwt.wakeup_later notify_rd_exit () ;
            Flow.shutdown flow.flow `read in
      Lwt.async @@ fun () ->
      Lwt.catch go (fun exn ->
          Runtime.report_exn connection exn ;
          Lwt.return_unit) in
    let rec wr_fiber () =
      let report_error err =
        Runtime.report_exn connection (to_flow_write_exception err) in
      let rec go () =
        Log_server.debug (fun m -> m "Compute next write operation.") ;
        match Runtime.next_write_operation connection with
        | `Upgrade -> failwith "Unimplemented"
        | `Write iovecs ->
            Log_server.debug (fun m -> m "next write operation: `write") ;
            Easy_flow.send ~report_error flow iovecs >>= fun res ->
            Runtime.report_write_result connection res ;
            Lwt.pause () >>= go
        | `Yield ->
            Log_server.debug (fun m -> m "next write operation: `yield") ;
            Runtime.yield_writer connection wr_fiber ;
            Lwt.pause ()
        | `Close _ ->
            Log_server.debug (fun m -> m "next write operation: `close") ;
            Lwt.wakeup_later notify_wr_exit () ;
            Flow.shutdown flow.flow `write in
      Lwt.async @@ fun () ->
      Lwt.catch go (fun exn ->
          (* Runtime.report_write_result connection `Closed ; *)
          Runtime.report_exn connection exn ;
          Lwt.return_unit) in
    rd_fiber () ;
    wr_fiber () ;
    Lwt.join [ rd_exit; wr_exit ] >>= fun () ->
    Log_server.debug (fun m -> m "End of transmission.") ;
    Easy_flow.close flow
end

let src = Logs.Src.create "paf-client"

module Log_client = (val Logs.src_log src : Logs.LOG)

module Client (Flow : Mirage_flow.S) (Runtime : RUNTIME) : sig
  val run : Runtime.t -> Flow.flow -> unit Lwt.t
end = struct
  open Lwt.Infix
  module Easy_flow = Make (Flow)

  let to_flow_exception err : exn = Flow (Fmt.str "%a" Flow.pp_error err)

  let to_flow_write_exception err : exn =
    Flow_write (Fmt.str "%a" Flow.pp_write_error err)

  let run connection flow =
    Easy_flow.create flow >>= fun flow ->
    let rd_exit, notify_rd_exit = Lwt.wait () in
    let wr_exit, notify_wr_exit = Lwt.wait () in

    let rec rd_fiber () =
      let report_error err =
        Runtime.report_exn connection (to_flow_exception err) in
      let rec go () =
        match Runtime.next_read_operation connection with
        | `Upgrade -> failwith "Unimplemented"
        | `Read ->
            Log_client.debug (fun m -> m "next read operation: `read") ;
            Easy_flow.recv flow ~report_error ~report_closed:ignore
              ~read:(Runtime.read connection)
              ~read_eof:(Runtime.read_eof connection)
            >>= fun _ -> Lwt.pause () >>= go
        | `Yield ->
            Log_client.debug (fun m -> m "next read operation: `yield") ;
            Runtime.yield_reader connection rd_fiber ;
            Lwt.pause ()
        | `Close ->
            Log_client.debug (fun m -> m "next read operation: `close.") ;
            Lwt.wakeup_later notify_rd_exit () ;
            flow.Easy_flow.rd_closed <- true ;
            Easy_flow.safely_close flow in
      Lwt.async @@ fun () ->
      Lwt.catch go (fun exn ->
          Runtime.report_exn connection exn ;
          Lwt.return_unit) in
    let rec wr_fiber () =
      let report_error err =
        Runtime.report_exn connection (to_flow_write_exception err) in
      let rec go () =
        match Runtime.next_write_operation connection with
        | `Upgrade -> failwith "Unimplemented"
        | `Write iovecs ->
            Log_client.debug (fun m -> m "next write operation: `write.") ;
            Easy_flow.send ~report_error flow iovecs >>= fun res ->
            Runtime.report_write_result connection res ;
            Lwt.pause () >>= go
        | `Yield ->
            Log_client.debug (fun m -> m "next write operation: `yield.") ;
            Runtime.yield_writer connection wr_fiber ;
            Lwt.pause ()
        | `Close _ ->
            Log_client.debug (fun m -> m "next write operation: `close.") ;
            Lwt.wakeup_later notify_wr_exit () ;
            flow.Easy_flow.wr_closed <- true ;
            Easy_flow.safely_close flow in
      Lwt.async @@ fun () ->
      Lwt.catch go (fun exn ->
          Runtime.report_exn connection exn ;
          Lwt.return ()) in
    wr_fiber () ;
    rd_fiber () ;
    Lwt.join [ rd_exit; wr_exit ] >>= fun () ->
    Log_client.debug (fun m -> m "End of transmission.") ;
    Easy_flow.close flow
end

type impl = Runtime : 'conn runtime * 'conn -> impl

type 't service =
  | Service : {
      accept : 't -> ('socket, ([> `Closed ] as 'error)) result Lwt.t;
      handshake : 'socket -> ('flow, ([> `Closed ] as 'error)) result Lwt.t;
      connection : 'flow -> (Mimic.flow * impl, 'error) result Lwt.t;
      close : 't -> unit Lwt.t;
    }
      -> 't service

and ('t, 'socket, 'flow, 'error) posix = {
  accept : 't -> ('socket, 'error) result Lwt.t;
  handshake : 'socket -> ('flow, 'error) result Lwt.t;
  close : 't -> unit Lwt.t;
}
  constraint 'error = [> `Closed ]

let service connection handshake accept close =
  Service { accept; connection; handshake; close }

open Lwt.Infix

let serve_when_ready : type t socket flow.
    (t, socket, flow, _) posix ->
    ?stop:Lwt_switch.t ->
    handler:(flow -> unit Lwt.t) ->
    t ->
    [ `Initialized of unit Lwt.t ] =
 fun service ?stop ~handler t ->
  let { accept; handshake; close } = service in
  `Initialized
    (let switched_off =
       let t, u = Lwt.wait () in
       Lwt_switch.add_hook stop (fun () ->
           Lwt.wakeup_later u (Ok `Stopped) ;
           Lwt.return_unit) ;
       t in
     let rec loop () =
       accept t >>= function
       | Ok socket ->
           Lwt.async (fun () ->
               handshake socket >>= function
               | Ok flow -> handler flow
               | Error `Closed ->
                   Logs.info (fun m -> m "Connection closed by peer") ;
                   Lwt.return ()
               | Error _err ->
                   Logs.err (fun m ->
                       m "Got an error from a TCP/IP connection.") ;
                   Lwt.return ()) ;
           loop ()
       | Error `Closed -> Lwt.return_error `Closed
       | Error _ -> Lwt.pause () >>= loop in
     let stop_result =
       Lwt.pick [ switched_off; loop () ] >>= function
       | Ok `Stopped -> close t >>= fun () -> Lwt.return_ok ()
       | Error _ as err -> close t >>= fun () -> Lwt.return err in
     stop_result >>= function Ok () | Error `Closed -> Lwt.return_unit)

let server : type t. t runtime -> t -> Mimic.flow -> unit Lwt.t =
 fun (module Runtime) conn flow ->
  let module Server = Server (Mimic) (Runtime) in
  Server.server conn flow

let serve ?stop service t =
  let (Service { accept; handshake; connection; close }) = service in
  let handler flow =
    connection flow >>= function
    | Ok (flow, Runtime (runtime, conn)) -> server runtime conn flow
    | Error _ -> Lwt.return_unit in
  serve_when_ready ?stop ~handler { accept; handshake; close } t

let run : type t. t runtime -> t -> Mimic.flow -> unit Lwt.t =
 fun (module Runtime) conn flow ->
  let module Client = Client (Mimic) (Runtime) in
  Client.run conn flow