(* Copyright (C) 2022--2024  Petter A. Urkedal <paurkedal@gmail.com>
 *
 * This library is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or (at your
 * option) any later version, with the LGPL-3.0 Linking Exception.
 *
 * This library is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
 * License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * and the LGPL-3.0 Linking Exception along with this library.  If not, see
 * <http://www.gnu.org/licenses/> and <https://spdx.org>, respectively.
 *)

open Lwt.Infix
open Caqti_platform

module type SOCKET_OPS =
  Caqti_platform.System_sig.SOCKET_OPS with type 'a fiber := 'a Lwt.t

module Make
  (STACK : Tcpip.Stack.V4V6)
  (DNS : Dns_client_mirage.S) =
struct
  module TCP = STACK.TCP
  module TLS = Tls_mirage.Make (TCP)
  module TCP_channel = Mirage_channel.Make (TCP)
  module TLS_channel = Mirage_channel.Make (TLS)

  module System_core = struct
    include Caqti_lwt.System_core

    type stdenv = {
      stack: STACK.t;
      dns: DNS.t;
    }
  end

  module Alarm = struct
    type t = {cancel: unit -> unit}

    let schedule ~sw ~stdenv:_ t f =
      let t_now = Mtime_clock.now () in
      let dt_ns =
        if Mtime.is_later t ~than:t_now then 0L else
        Mtime.Span.to_uint64_ns (Mtime.span t t_now)
      in
      let task = Mirage_sleep.ns dt_ns >|= f in
      let hook =
        Caqti_lwt.Switch.on_release_cancellable sw
          (fun () -> Lwt.cancel task; Lwt.return_unit)
      in
      {cancel = (fun () -> Caqti_lwt.Switch.remove_hook hook; Lwt.cancel task)}

    let unschedule alarm = alarm.cancel ()
  end

  module Pool = Caqti_platform.Pool.Make (System_core) (Alarm)

  module System = struct
    include System_core
    module Pool = Pool

    module Net = struct

      module Sockaddr = struct
        type t = [`Tcp of Ipaddr.t * int | `Unix of string]
        let unix s = `Unix s
        let tcp (host, port) = `Tcp (host, port)
      end

      let getaddrinfo_ipv4 dns host port =
        let extract (_, ips) =
          Ipaddr.V4.Set.elements ips
            |> List.map (fun ip -> `Tcp (Ipaddr.V4 ip, port))
        in
        DNS.getaddrinfo dns Dns.Rr_map.A host >|= Result.map extract

      let getaddrinfo_ipv6 dns host port =
        let extract (_, ips) =
          Ipaddr.V6.Set.elements ips
            |> List.map (fun ip -> `Tcp (Ipaddr.V6 ip, port))
        in
        DNS.getaddrinfo dns Dns.Rr_map.Aaaa host >|= Result.map extract

      let getaddrinfo ~stdenv:{stack; dns} host port =
        let laddrs = STACK.IP.configured_ips (STACK.ip stack) in
        (match
          List.exists Ipaddr.(function V4 _ -> true | V6 _ -> false) laddrs,
          List.exists Ipaddr.(function V4 _ -> false | V6 _ -> true) laddrs
         with
         | true, true ->
            getaddrinfo_ipv4 dns host port >>= fun r4 ->
            getaddrinfo_ipv6 dns host port >|= fun r6 ->
            (match r4, r6 with
             | Ok addrs4, Ok addrs6 -> Ok (addrs4 @ addrs6)
             | Ok addrs, Error _ | Error _, Ok addrs -> Ok addrs
             | Error (`Msg msg4), Error (`Msg msg6) ->
                if String.equal msg4 msg6 then Error (`Msg msg4) else
                Error (`Msg ("IPv4: " ^ msg4 ^ " IPv6: " ^ msg6)))
         | true, false -> getaddrinfo_ipv4 dns host port
         | false, true -> getaddrinfo_ipv6 dns host port
         | false, false ->
            Lwt.return (Error (`Msg "No IP address assigned to host.")))

      let convert_io_exception = function
       | Failure msg -> Some (Caqti_error.Msg msg) (* Channel.S.error *)
       | _ -> None

      module Make_stream_ops (Channel : Mirage_channel.S) = struct
        type t = Channel.t

        let output_char channel c =
          Channel.write_char channel c;
          Lwt.return_unit

        let output_string channel s =
          Channel.write_string channel s 0 (String.length s);
          Lwt.return_unit

        let flush channel =
          Channel.flush channel >>= function
           | Ok () -> Lwt.return_unit
           | Error err ->
              Lwt.fail_with (Format.asprintf "%a" Channel.pp_write_error err)

        let input_char channel =
          Channel.read_char channel >>= function
           | Ok (`Data c) -> Lwt.return c
           | Ok `Eof -> Lwt.fail End_of_file
           | Error err ->
              Lwt.fail_with (Format.asprintf "%a" Channel.pp_error err)

        let really_input channel buf off len =
          Channel.read_exactly ~len channel >>= function
           | Ok (`Data bufs) ->
              let content = Cstruct.copyv bufs in
              Bytes.blit_string content 0 buf off len;
              Lwt.return_unit
           | Ok `Eof -> Lwt.fail End_of_file
           | Error err ->
              Lwt.fail_with (Format.asprintf "%a" Channel.pp_error err)

        let close channel =
          Channel.close channel >>= function
           | Ok () -> Lwt.return_unit
           | Error err ->
              Lwt.fail_with (Format.asprintf "%a" Channel.pp_write_error err)
      end

      module TCP_stream_ops = Make_stream_ops (TCP_channel)
      module TLS_stream_ops = Make_stream_ops (TLS_channel)

      module Socket = struct
        type t = V : {
          tcp_flow: TCP_channel.flow option;
          ops: (module SOCKET_OPS with type t = 'a);
          channel: 'a;
        } -> t

        let output_char (V {ops = (module Ops); channel; _}) =
          Ops.output_char channel
        let output_string (V {ops = (module Ops); channel; _}) =
          Ops.output_string channel
        let flush (V {ops = (module Ops); channel; _}) =
          Ops.flush channel
        let input_char (V {ops = (module Ops); channel; _}) =
          Ops.input_char channel
        let really_input (V {ops = (module Ops); channel; _}) =
          Ops.really_input channel
        let close (V {ops = (module Ops); channel; _}) =
          Ops.close channel
      end

      type tcp_flow = TCP_channel.flow
      type tls_flow = Tls_flow : {
        ops: (module SOCKET_OPS with type t = 'a);
        channel: 'a;
      } -> tls_flow

      let connect_tcp ~sw:_ ~stdenv:{stack; _} sockaddr =
        (match sockaddr with
         | `Unix _ ->
            Lwt.return_error
              (Caqti_error.Msg "Unix sockets are not available under MirageOS.")
         | `Tcp (ipaddr, port) ->
            TCP.create_connection (STACK.tcp stack) (ipaddr, port) >|=
            (function
             | Ok flow ->
                let channel = TCP_channel.create flow in
                Ok (Socket.V {
                  tcp_flow = Some flow;
                  ops = (module TCP_stream_ops);
                  channel;
                })
             | Error err ->
                let msg = Format.asprintf "%a" TCP.pp_error err in
                Error (Caqti_error.Msg msg)))

      let tcp_flow_of_socket (Socket.V {tcp_flow; _}) = tcp_flow

      let socket_of_tls_flow ~sw:_ (Tls_flow {ops; channel}) =
        Socket.V {tcp_flow = None; ops; channel}

      module type TLS_PROVIDER = Caqti_platform.System_sig.TLS_PROVIDER
        with type 'a fiber := 'a Lwt.t
         and type tcp_flow := tcp_flow
         and type tls_flow := tls_flow

      module Tls_provider = struct
        type tls_config = Tls.Config.client

        let tls_config_key = Caqti_tls.Config.client

        let start_tls ~config ?host flow =
          TLS.client_of_flow config ?host flow >|=
          (function
           | Ok tls_flow ->
              Ok (Tls_flow {
                ops = (module TLS_stream_ops);
                channel = TLS_channel.create tls_flow;
              })
           | Error err ->
              let msg = Format.asprintf "%a" TLS.pp_write_error err in
              Error (Caqti_error.Msg msg))
      end

      let tls_providers_r : (module TLS_PROVIDER) list ref =
        ref [(module Tls_provider : TLS_PROVIDER)]

      let tls_providers _ = !tls_providers_r

      let register_tls_provider p = tls_providers_r := p :: !tls_providers_r

    end
  end

  module Loader = Caqti_platform.Driver_loader.Make (System)

  include Connector.Make (System) (Pool) (Loader)

  let connect
        ?subst ?env ?config ?tweaks_version ?(sw = Caqti_lwt.Switch.eternal)
        stack dns uri =
    connect ?subst ?env ?config ?tweaks_version ~sw ~stdenv:{stack; dns} uri

  let with_connection ?subst ?env ?config ?tweaks_version stack dns uri f =
    with_connection ?subst ?env ?config ?tweaks_version ~stdenv:{stack; dns} uri f

  let connect_pool
        ?pool_config ?post_connect ?subst ?env ?config ?tweaks_version
        ?(sw = Caqti_lwt.Switch.eternal) stack dns uri =
    connect_pool
      ?pool_config ?post_connect ?subst ?env ?config ?tweaks_version
      ~sw ~stdenv:{stack; dns} uri

end