stack_picker/lib/srv.ml

open Tyxml

type endpoint_key = {
  path : string;
  callback : Http.Request.t -> Http.Response.t;
}

type opt_route_response = {
  headers : Http.Header.t option;
  status : Http.Status.t option;
  body : Cohttp_eio.Body.t option;
}

type route_response = {
  headers : Http.Header.t;
  status : Http.Status.t;
  body : Cohttp_eio.Body.t;
}

type route_callback = Http.Request.t -> route_response

module Endpath : sig
  type t

  val compare : t -> t -> int
  val of_tuple : Http.Method.t * string -> t
end = struct
  type t = Http.Method.t * string

  let compare (meth_a, path_a) (meth_b, path_b) =
    let meth_compare = Http.Method.compare meth_a meth_b in
    if meth_compare <> 0 then meth_compare
    else String.(compare (lowercase_ascii path_a) (lowercase_ascii path_b))

  let of_tuple (method_, path) : t = (method_, path)
end

module EndpointMap = Map.Make (Endpath)

(* let endpoints: route_callback EndpointMap.t ref = ref EndpointMap.empty *)

let default_headers =
  let headers_strings = [ ("Content-Type", "text/html; charset=UTF-8") ] in
  ref (Http.Header.of_list headers_strings)

let default_body = ref ("" |> Eio.Flow.string_source)
let log_warning ex = Logs.warn (fun f -> f "%a" Eio.Exn.pp ex)

let get (path : string) (callback : route_callback) =
  let key = Endpath.of_tuple (`GET, path) in
  (key, callback)

let post (path : string) (callback : route_callback) =
  let key = Endpath.of_tuple (`POST, path) in
  (key, callback)
(* endpoints :=  !endpoints |> EndpointMap.add key callback *)

let rec find_html_files (static_path : string) =
  let files = Sys.readdir static_path in
  Array.fold_left
    (fun acc file ->
      let path = Filename.concat static_path file in
      if Sys.is_directory path then acc @ find_html_files path
      else if Filename.check_suffix path ".html" then path :: acc
      else acc)
    [] files

let respond ?body ?headers ?status () =
  {
    body = Option.value body ~default:!default_body;
    headers = Option.value headers ~default:!default_headers;
    status = Option.value status ~default:`OK;
  }

let respondStr ?body ?headers ?status () =
  let eio_body =
    match body with Some b -> Eio.Flow.string_source b | None -> !default_body
  in
  {
    body = eio_body;
    headers = Option.value headers ~default:!default_headers;
    status = Option.value status ~default:`OK;
  }

let respondHtmlDoc ?body ?headers ?status () =
  let eio_body =
    match body with
    | Some b ->
        let content = Format.asprintf "%a" (Html.pp ()) b in
        Eio.Flow.string_source content
    | None -> !default_body
  in
  {
    body = eio_body;
    headers = Option.value headers ~default:!default_headers;
    status = Option.value status ~default:`OK;
  }

let respondHtml ?body ?headers ?status () =
  let eio_body =
    match body with
    | Some b ->
        let content = Format.asprintf "%a" (Html.pp_elt ()) b in
        Eio.Flow.string_source content
    | None -> !default_body
  in
  {
    body = eio_body;
    headers = Option.value headers ~default:!default_headers;
    status = Option.value status ~default:`OK;
  }

(** TODO not implemented, needs yojson ? *)
let respondJson body ?headers ?status () =
  {
    body;
    headers = Option.value headers ~default:!default_headers;
    status = Option.value status ~default:`OK;
  }

let run ~env ~sw ~port endpoint_list =
  let endpoints : route_callback EndpointMap.t =
    EndpointMap.of_list endpoint_list
  in
  let socket =
    Eio.Net.listen env#net ~sw ~backlog:128 ~reuse_addr:true
      (`Tcp (Eio.Net.Ipaddr.V4.loopback, port))
  in
  let handler _socket request _body =
    (* let req = Http.Request.resource request in *)
    let path = Http.Request.resource request in
    let end_key = Endpath.of_tuple (request.meth, path) in
    let opt_path = EndpointMap.find_opt end_key endpoints in
    match opt_path with
    | Some route ->
        (* NOTE needs an Exn.protect here? try / finally *)
        let response = route request in
        Cohttp_eio.Server.respond () ~headers:response.headers
          ~status:response.status ~body:response.body
    | None ->
        let not_found_body = "Route not defined\n" |> Eio.Flow.string_source in
        Cohttp_eio.Server.respond () ~status:`Not_found ~body:not_found_body
  in
  let server = Cohttp_eio.Server.make ~callback:handler () in
  Cohttp_eio.Server.run socket server ~on_error:log_warning