open Syntax

let src = Logs.Src.create "server"

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

let is_digit = function '0' .. '9' -> true | _ -> false

let root = {|Hello!
  This is MTE, the MirageOS Taler Exchange unikernel.
|}

module type S = sig
  type response
  type request

  val version : [ `HTTP_1_1 | `HTTP_2_0 ]
  val create : (string * string) list -> H2.Status.t -> response
  val with_etag : string -> response -> response
  val with_status : H2.Status.t -> response -> response
  val get : request -> string -> string option
end

let transmit_over_http : to_close:_ -> Mimic.flow -> Mimic.flow -> unit Lwt.t =
 fun ~to_close src dst ->
  let closed = Lwt_mvar.create_empty () in
  let rec loop ~src ~dst () =
    let* res =
      Lwt.pick
        [
          (let+? v = Mimic.read src in
           (v :> [ `Closed | _ Mirage_flow.or_eof ]));
          (let+ v = Lwt_mvar.take closed in
           Ok v);
        ]
    in
    match res with
    | Error err ->
        Log.err (fun m ->
            m "Got an error while we reading the source (CONNECT): %a."
              Mimic.pp_error err);
        if Lwt_mvar.is_empty closed then Lwt_mvar.put closed `Closed
        else Lwt.return_unit
    | Ok `Closed -> Lwt.return_unit
    | Ok `Eof ->
        if Lwt_mvar.is_empty closed then Lwt_mvar.put closed `Closed
        else Lwt.return_unit
    | Ok (`Data cs) -> (
        Log.debug (fun m -> m "Transfer over HTTP:");
        Log.debug (fun m ->
            m "@[<hov>%a@]." (Hxd_string.pp Hxd.default) (Cstruct.to_string cs));
        let* res = Mimic.write dst cs in
        match res with
        | Ok () ->
            let* () = Lwt.pause () in
            loop ~src ~dst ()
        | Error err ->
            Log.err (fun m ->
                m
                  "Got an error while we writing into the destination \
                   (CONNECT): %a."
                  Mimic.pp_write_error err);
            if Lwt_mvar.is_empty closed then Lwt_mvar.put closed `Closed
            else Lwt.return_unit)
  in
  let* () = Lwt.join [ loop ~src ~dst (); loop ~src:dst ~dst:src () ] in
  to_close src;
  let* () = Lwt.join [ Mimic.close src; Mimic.close dst ] in
  Log.debug (fun m -> m "Connection closed properly on both side.");
  Lwt.return_unit

(***** HTTP/1.1 *****)

module S_HTTP_1_1 = struct
  type request = H1.Request.t
  type response = H1.Response.t

  let version = `HTTP_1_1
  let get request name = H1.Headers.get request.H1.Request.headers name

  let create headers = function
    | #H1.Status.t as status ->
        H1.Response.create ~headers:(H1.Headers.of_list headers) status
    | _ -> assert false

  let with_etag etag response =
    let headers = response.H1.Response.headers in
    { response with H1.Response.headers= H1.Headers.add headers "etag" etag }

  let with_status (status : H2.Status.t) response =
    match status with
    | #H1.Status.t as status -> { response with H1.Response.status }
    | _ -> assert false
end

let connect_http_1_1 ~ctx ~authenticator ~to_close flow reqd =
  let request = H1.Reqd.request reqd in
  match H1.Headers.get request.H1.Request.headers "host" with
  | Some uri ->
      let uri = "http://" ^ uri in
      Lwt.async (fun () ->
          let* res = Connect.create_connection ~ctx ~authenticator uri in
          match res with
          | Ok dst ->
              let headers = H1.Headers.of_list [ ("connection", "close") ] in
              let response =
                H1.Response.create ~reason:"CONNECT" ~headers `OK
              in
              H1.Reqd.respond_with_string reqd response "";
              H1.Body.Reader.close (H1.Reqd.request_body reqd);
              transmit_over_http ~to_close flow dst
          | Error err ->
              Log.err (fun m ->
                  m "Got an error while connection to %S: %a." uri
                    Mimic.pp_error err);
              let contents = Fmt.str "Invalid URI: %S" uri in
              let headers =
                H1.Headers.of_list
                  [
                    ("content-length", string_of_int (String.length contents));
                    ("connection", "close"); ("content-type", "text/plain");
                  ]
              in
              let response =
                H1.Response.create ~reason:"CONNECT" ~headers `Bad_request
              in
              H1.Reqd.respond_with_string reqd response contents;
              Lwt.return_unit)
  | None ->
      let contents = "Missing Host field." in
      let headers =
        H1.Headers.of_list
          [
            ("content-length", string_of_int (String.length contents));
            ("connection", "close"); ("content-type", "text/plain");
          ]
      in
      let response =
        H1.Response.create ~reason:"CONNECT" ~headers `Bad_request
      in
      H1.Reqd.respond_with_string reqd response contents

let http_1_1_request_handler ~ctx ~authenticator ~to_close =
 fun flow reqd ->
  let request = H1.Reqd.request reqd in
  Log.debug (fun m ->
      m "(HTTP/1.1) request-handler: %S" request.H1.Request.target);
  match request.H1.Request.meth with
  | `CONNECT ->
      Log.debug (fun m -> m "Start to transmit data over HTTP/1.1.");
      connect_http_1_1 ~ctx ~authenticator ~to_close flow reqd
  | _meth -> (
      match String.split_on_char '/' request.H1.Request.target with
      | [ ""; "" ] ->
          let headers =
            H1.Headers.of_list
              [
                ("content-length", string_of_int (String.length root));
                ("connection", "close"); ("content-type", "text/plain");
              ]
          in
          let response = H1.Response.create ~reason:"root" ~headers `OK in
          H1.Reqd.respond_with_string reqd response root
      | _ ->
          let contents = "Not found." in
          let headers =
            H1.Headers.of_list
              [
                ("content-type", "text/plain"); ("connection", "close");
                ("content-length", string_of_int (String.length contents));
              ]
          in
          let response =
            H1.Response.create ~reason:"not-found" ~headers `Not_found
          in
          H1.Reqd.respond_with_string reqd response contents)

(***** H2 *****)

module S_HTTP_2_0 = struct
  type request = H2.Request.t
  type response = H2.Response.t

  let version = `HTTP_2_0
  let get request name = H2.Headers.get request.H2.Request.headers name

  let create headers status =
    H2.Response.create ~headers:(H2.Headers.of_list headers) status

  let with_etag etag response =
    let headers = response.H2.Response.headers in
    { response with H2.Response.headers= H2.Headers.add headers "etag" etag }

  let with_status status response = { response with H2.Response.status }
end

let transmit src dst =
  let rec on_eof () = H2.Body.Reader.close src; H2.Body.Writer.close dst
  and on_read buf ~off ~len =
    H2.Body.Writer.write_bigstring dst ~off ~len buf;
    H2.Body.Reader.schedule_read src ~on_eof ~on_read
  in
  H2.Body.Reader.schedule_read src ~on_eof ~on_read

let connect_http_2_0 ~ctx ~authenticator ~to_close flow reqd =
  let request = H2.Reqd.request reqd in
  match H2.Headers.get request.H2.Request.headers "host" with
  | Some uri ->
      let uri = "http://" ^ uri in
      Lwt.async (fun () ->
          let* res = Connect.create_connection ~ctx ~authenticator uri in
          match res with
          | Ok dst ->
              let response = H2.Response.create `OK in
              H2.Reqd.respond_with_string reqd response "";
              H2.Body.Reader.close (H2.Reqd.request_body reqd);
              transmit_over_http ~to_close flow dst
          | Error err ->
              Log.err (fun m ->
                  m "Got an error while connection to %S: %a." uri
                    Mimic.pp_error err);
              let contents = Fmt.str "Invalid URI: %S" uri in
              let headers =
                H2.Headers.of_list
                  [
                    ("content-length", string_of_int (String.length contents));
                    ("content-type", "text/plain");
                  ]
              in
              let response = H2.Response.create ~headers `Bad_request in
              H2.Reqd.respond_with_string reqd response contents;
              Lwt.return_unit)
  | None ->
      let contents = "Missing Host field." in
      let headers =
        H2.Headers.of_list
          [
            ("content-length", string_of_int (String.length contents));
            ("content-type", "text/plain");
          ]
      in
      let response = H2.Response.create ~headers `Bad_request in
      H2.Reqd.respond_with_string reqd response contents

let http_2_0_request_handler ~ctx ~authenticator ~to_close =
 fun flow reqd ->
  let request = H2.Reqd.request reqd in
  Log.debug (fun m -> m "(H2) request-handler: %S" request.H2.Request.target);
  match request.H2.Request.meth with
  | `CONNECT ->
      Log.debug (fun m -> m "Start to transmit data over H2.");
      connect_http_2_0 ~ctx ~authenticator ~to_close flow reqd
  | _meth -> (
      match String.split_on_char '/' request.H2.Request.target with
      | [ ""; "" ] ->
          let headers =
            H2.Headers.of_list
              [
                ("content-length", string_of_int (String.length root));
                ("content-type", "text/plain");
              ]
          in
          let response = H2.Response.create ~headers `OK in
          H2.Reqd.respond_with_string reqd response root
      | _ ->
          let contents = "Not found." in
          let headers =
            H2.Headers.of_list
              [
                ("content-type", "text/plain");
                ("content-length", string_of_int (String.length contents));
              ]
          in
          let response = H2.Response.create ~headers `Not_found in
          H2.Reqd.respond_with_string reqd response contents)

let alpn_request_handler : type reqd headers request response ro wo.
    ctx:_ ->
    authenticator:_ ->
    to_close:_ ->
    ?shutdown:_ ->
    _ ->
    _ ->
    reqd ->
    (reqd, headers, request, response, ro, wo) Alpn.protocol ->
    unit =
 fun ~ctx ~authenticator ~to_close ?shutdown:_ flow _edn reqd -> function
  | Alpn.HTTP_1_1 _ ->
      http_1_1_request_handler ~ctx ~authenticator ~to_close flow reqd
  | Alpn.H2 _ ->
      http_2_0_request_handler ~ctx ~authenticator ~to_close flow reqd

let headers_of_list : type reqd headers request response ro wo.
    (reqd, headers, request, response, ro, wo) Alpn.protocol ->
    (string * string) list ->
    headers =
 fun protocol lst ->
  match protocol with
  | Alpn.HTTP_1_1 _ -> H1.Headers.of_list lst
  | Alpn.H2 _ -> H2.Headers.of_list lst

let respond_with_string : type reqd headers request response ro wo.
    (reqd, headers, request, response, ro, wo) Alpn.protocol ->
    headers:headers ->
    respond:(headers -> wo) ->
    string ->
    unit =
 fun protocol ~headers ~respond str ->
  let body = respond headers in
  match protocol with
  | Alpn.HTTP_1_1 _ ->
      H1.Body.Writer.write_string body str;
      H1.Body.Writer.close body
  | Alpn.H2 _ ->
      H2.Body.Writer.write_string body str;
      H2.Body.Writer.close body

let alpn_error_handler : type reqd headers request response ro wo.
    _ ->
    (reqd, headers, request, response, ro, wo) Alpn.protocol ->
    ?request:_ ->
    _ ->
    (headers -> wo) ->
    unit =
 fun _edn protocol ?request:_ error respond ->
  let contents =
    match error with
    | `Bad_gateway -> {|Bad gateway.|}
    | `Bad_request -> {|Bad request.|}
    | `Exn (Paf.Flow err) | `Exn (Paf.Flow_write err) ->
        Fmt.str {|I/O error: %s.|} err
    | `Exn exn -> Fmt.str {|Unknown error: %S.|} (Printexc.to_string exn)
    | `Internal_server_error -> {|Internal server error.|}
  in
  let headers =
    headers_of_list protocol
      [
        ("content-type", "text/plain");
        ("content-length", string_of_int (String.length contents));
      ]
  in
  respond_with_string protocol ~respond ~headers contents

let http_1_1_error_handler edn ?request error respond =
  alpn_error_handler edn ?request Alpn.http_1_1
    (error :> Alpn.server_error)
    respond

let http_2_0_error_handler edn ?request error respond =
  alpn_error_handler edn ?request Alpn.h2 (error :> Alpn.server_error) respond