open Rresult
open Cmdliner

let port =
  let doc = Arg.info ~doc:"Port of HTTP service." [ "p"; "port" ] in
  Arg.(value & opt int 3434 doc)

let use_tls =
  let doc =
    Arg.info ~doc:"Start an HTTP server with a TLS certificate." [ "tls" ]
  in
  Mirage_runtime.register_arg Arg.(value & flag doc)

let tls_port =
  let doc = Arg.info ~doc:"Port of HTTPS service." [ "tls-port" ] in
  Mirage_runtime.register_arg Arg.(value & opt int 4343 doc)

let alpn =
  let alpns = [ "h2"; "http/1.1" ] in
  let doc =
    Printf.sprintf "Protocols handled by the HTTP server. Must be %s."
      (Arg.doc_alts alpns)
  in
  let doc = Arg.info ~doc [ "alpn" ] in
  Mirage_runtime.register_arg
    Arg.(value & opt_all (enum (List.map (fun v -> (v, v)) alpns)) alpns doc)

let ( <.> ) f g x = f (g x)
let always x _ = x

let map_err_to_string pp_err res =
  Lwt.map (R.reword_error (R.msgf "%a" pp_err)) res

let list_get_ok l =
  let err = ref None in
  try
    l
    |> List.map (function
      | Error _e as e ->
          err := Some e;
          raise Exit
      | Ok v -> v)
    |> Result.ok
  with Exit -> ( match !err with None -> assert false | Some v -> v)

let lwt_list_get_ok l =
  let open Lwt.Syntax in
  let+ l = l in
  list_get_ok l

module Make
    (Assets_ro : Mirage_kv.RO)
    (Certificates_ro : Mirage_kv.RO)
    (Keys_ro : Mirage_kv.RO)
    (Tcp : Tcpip.Tcp.S with type ipaddr = Ipaddr.t)
    (Connect : Connect.S)
    (HTTP_server : Paf_mirage.S) =
struct
  module Assets = struct
    let ( let*? ) = Lwt_result.bind
    let ( let+? ) x f = Lwt_result.map f x
    let map_err_to_string = map_err_to_string Assets_ro.pp_error

    let get_subdirs ro k =
      let+? keys = Assets_ro.list ro k |> map_err_to_string in
      List.filter (fun (_, t) -> t = `Dictionary) keys |> List.map fst

    let get_values ro k =
      let+? keys = Assets_ro.list ro k |> map_err_to_string in
      List.filter (fun (_, t) -> t = `Value) keys |> List.map fst

    let find keys name =
      keys
      |> List.find_opt Mirage_kv.Key.(equal (v name))
      |> Option.to_result ~none:(`Msg (Fmt.str "missing `%s` directory" name))
      |> Lwt_result.lift

    let get ro s =
      let*? base_dir = get_subdirs ro Mirage_kv.Key.empty in
      let*? dir = find base_dir s in
      let*? subdirs = get_subdirs ro dir in
      let*? l =
        lwt_list_get_ok
        @@ Lwt_list.map_s
             (fun lang_dir ->
               let+? files = get_values ro lang_dir in
               (lang_dir, files))
             subdirs
      in
      Lwt.return
      @@
      let open Result.Syntax in
      let* l =
        list_get_ok
        @@ List.map
             (fun (lang, l) ->
               let lang = Mirage_kv.Key.basename lang in
               let* l =
                 List.map Mirage_kv.Key.basename l
                 |> List.map (fun s ->
                     match String.split_on_char '.' s with
                     | [] -> assert false
                     | [ etag; ext ] -> Ok (etag, ext)
                     | _ -> Error (`Msg (Fmt.str "invalid file name: `%s`" s)))
                 |> list_get_ok
               in
               let etag_l, ext_l = List.split l in
               Ok (lang, etag_l, ext_l))
             l
      in
      let lang_l =
        l |> List.map (fun (v, _, _) -> v) |> List.sort String.compare
      in
      let* etag =
        let etag_l = l |> List.map (fun (_, v, _) -> v) |> List.concat in
        match List.sort_uniq String.compare etag_l with
        | [ etag ] -> Ok etag
        | [] ->
            Fmt.error_msg "directory `%s` is empty"
              (Mirage_kv.Key.to_string dir)
        | l ->
            Fmt.error_msg "directory `%s` contains more than one etag: %a"
              (Mirage_kv.Key.to_string dir)
              (Fmt.list Fmt.string) l
      in
      let* ext_l =
        let ll =
          l
          |> List.map (fun (_, _, v) -> v)
          |> List.map (List.sort String.compare)
        in
        match List.sort_uniq Stdlib.compare ll with
        | [ l ] -> Ok l
        | [] -> assert false
        | _ll ->
            Fmt.error_msg
              "directory `%s` does not has the same set of file extensions for \
               each language"
              (Mirage_kv.Key.to_string dir)
      in
      Ok (etag, lang_l, ext_l)

    let assets ro =
      let*? terms_assoc = get ro "terms" in
      let*? privacy_assoc = get ro "privacy" in
      Lwt_result.return (terms_assoc, privacy_assoc)
  end

  let tls certificate_ro key_ro =
    let ( let*? ) = Lwt_result.bind in
    let ( let+? ) x f = Lwt_result.map f x in
    let*? keys =
      Keys_ro.list key_ro Mirage_kv.Key.empty
      |> map_err_to_string Keys_ro.pp_error
    in
    let keys = List.filter (fun (_, t) -> t = `Value) keys in
    let*? certificates =
      Certificates_ro.list certificate_ro Mirage_kv.Key.empty
      |> map_err_to_string Certificates_ro.pp_error
    in
    let certificates = List.filter (fun (_, t) -> t = `Value) certificates in
    let fold acc (name, _) =
      match Mirage_kv.Key.basename name with
      | ".gitkeep" -> Lwt.return acc
      | _ ->
          let*? data =
            Certificates_ro.get certificate_ro name
            |> map_err_to_string Certificates_ro.pp_error
          in
          let*? certificates =
            Lwt.return (X509.Certificate.decode_pem_multiple data)
          in
          let+? acc = Lwt.return acc in
          (name, certificates) :: acc
    in
    let*? certificates = Lwt_list.fold_left_s fold (Ok []) certificates in
    let fold acc (name, _) =
      match Mirage_kv.Key.basename name with
      | ".gitkeep" -> Lwt.return acc
      | _ ->
          let*? data =
            Keys_ro.get key_ro name |> map_err_to_string Keys_ro.pp_error
          in
          let*? key = Lwt.return (X509.Private_key.decode_pem data) in
          let+? acc = Lwt.return acc in
          (name, key) :: acc
    in
    let+? keys = Lwt_list.fold_left_s fold (Ok []) keys in
    let tbl = Hashtbl.create 0x10 in
    List.iter
      (fun (name, certificates) ->
        match List.assoc_opt name keys with
        | Some key -> Hashtbl.add tbl name (certificates, key)
        | None -> ())
      certificates;
    match Hashtbl.fold (fun _ certchain acc -> certchain :: acc) tbl [] with
    | [] -> `None
    | [ certchain ] -> `Single certchain
    | certchains -> `Multiple certchains

  let http_1_1_request_handler ~ctx ~authenticator flow _edn =
    let module R = (val Mimic.repr HTTP_server.tcp_protocol) in
    fun reqd ->
      match (H1.Reqd.request reqd).H1.Request.meth with
      | `CONNECT ->
          HTTP_server.TCP.no_close flow;
          let to_close = function
            | R.T flow -> HTTP_server.TCP.to_close flow
            | _ -> ()
          in
          Server.http_1_1_request_handler ~ctx ~authenticator ~to_close
            (R.T flow) reqd
      | _ ->
          Server.http_1_1_request_handler ~ctx ~authenticator
            ~to_close:(always ()) (R.T flow) reqd

  let alpn_handler ~ctx ~authenticator =
    let module R = (val Mimic.repr HTTP_server.tls_protocol) in
    let to_close = function
      | R.T flow -> HTTP_server.TLS.to_close flow
      | _ -> ()
    in
    {
      Alpn.error= Server.alpn_error_handler;
      Alpn.request=
        (fun flow edn reqd protocol ->
          Server.alpn_request_handler ~ctx ~authenticator ~to_close (R.T flow)
            edn reqd protocol);
    }

  let run_with_tls ~ctx ~authenticator ~tls http_server tls_port tcpv4v6 =
    let open Lwt.Infix in
    let alpn_service =
      HTTP_server.alpn_service ~tls (alpn_handler ~ctx ~authenticator)
    in
    let http_1_1_service =
      HTTP_server.http_service ~error_handler:Server.http_1_1_error_handler
        (http_1_1_request_handler ~ctx ~authenticator)
    in
    HTTP_server.init ~port:tls_port tcpv4v6 >|= Paf.serve alpn_service
    >>= fun (`Initialized th0) ->
    Paf.serve http_1_1_service http_server |> fun (`Initialized th1) ->
    Lwt.both th0 th1 >>= fun ((), ()) -> Lwt.return_unit

  let run ~ctx ~authenticator http_server =
    let http_1_1_service =
      HTTP_server.http_service ~error_handler:Server.http_1_1_error_handler
        (http_1_1_request_handler ~ctx ~authenticator)
    in
    Paf.serve http_1_1_service http_server |> fun (`Initialized th) -> th

  let start assets_ro certificate_ro key_ro tcpv4v6 ctx http_server =
    let open Lwt.Infix in
    let authenticator = Connect.authenticator in
    Assets.assets assets_ro >>= fun res ->
    match res with
    | Error (`Msg m) -> Fmt.failwith "Assets configuration error: %s." m
    | Ok (terms, _privacy) -> (
        let etag, lang_l, ext_l = terms in
        Fmt.pr "ETAG: %s@\nlanguages: %a@\nextensions: %a@." etag
          (Fmt.list ~sep:(Fmt.any ", ") Fmt.string)
          lang_l
          (Fmt.list ~sep:(Fmt.any ", ") Fmt.string)
          ext_l;
        tls certificate_ro key_ro >>= fun tls ->
        match use_tls () with
        | false -> run ~ctx ~authenticator http_server
        | true -> (
            match tls with
            | Error (`Msg m) ->
                Fmt.failwith
                  "A TLS server requires, at least, one certificate and one \
                   private key. Received error %s."
                  m
            | Ok certificates -> (
                let alpn_protocols = alpn () in
                match Tls.Config.server ~certificates ~alpn_protocols () with
                | Error (`Msg m) ->
                    Fmt.failwith "TLS configuration error: %s." m
                | Ok tls ->
                    run_with_tls ~ctx ~authenticator ~tls http_server
                      (tls_port ()) tcpv4v6)))
end