open Eio.Std

module Flow = Eio.Flow

exception Tls_alert   of Tls.Packet.alert_type
exception Tls_failure of Tls.Engine.failure

type Eio.Exn.Backend.t += Tls_socket_closed
let () = Eio.Exn.Backend.register_pp (fun f -> function
    | Tls_socket_closed -> Fmt.pf f "TLS_socket_closed"; true
    | _ -> false
  )

type ty = [ `Tls | Eio.Flow.two_way_ty | Eio.Resource.close_ty ]
type t = ty r

module Raw = struct

  (* We could replace [`Eof] with [`Error End_of_file] and then use
     a regular [result] type here. *)
  type t = {
    flow           : [Flow.two_way_ty | Eio.Resource.close_ty] r;
    mutable state  : [ `Active of Tls.Engine.state
                     | `Read_closed of Tls.Engine.state
                     | `Write_closed of Tls.Engine.state
                     | `Closed
                     | `Error of exn ] ;
    mutable linger : Cstruct.t option ;
    recv_buf       : Cstruct.t ;
  }

  let half_close state mode =
    match state, mode with
    | `Active tls, `read -> `Read_closed tls
    | `Active tls, `write -> `Write_closed tls
    | `Active _, `read_write -> `Closed
    | `Read_closed tls, `read -> `Read_closed tls
    | `Read_closed _, (`write | `read_write) -> `Closed
    | `Write_closed tls, `write -> `Write_closed tls
    | `Write_closed _, (`read | `read_write) -> `Closed
    | (`Closed | `Error _) as e, (`read | `write | `read_write) -> e

  let inject_state tls = function
    | `Active _ -> `Active tls
    | `Read_closed _ -> `Read_closed tls
    | `Write_closed _ -> `Write_closed tls
    | (`Closed | `Error _) as e -> e

  let write_t t s =
    try Flow.copy_string s t.flow
    with exn ->
      (match t.state with
       | `Error _ -> ()
       | _ -> t.state <- `Error exn) ;
      raise exn

  let try_write_t t cs =
    try write_t t cs
    with _ -> Eio.Fiber.check ()      (* Error is in [t.state] *)

  let rec read_react t =

    let handle tls buf =
      match Tls.Engine.handle_tls tls buf with
      | Ok (state', eof, `Response resp, `Data data) ->
          let state' = inject_state state' t.state in
          let state' = Option.(value ~default:state' (map (fun `Eof -> half_close state' `read) eof)) in
          t.state <- state' ;
          Option.iter (try_write_t t) resp;
          Option.map Cstruct.of_string data

      | Error (fail, `Response resp) ->
          t.state <- `Error (match fail with `Alert a -> Tls_alert a | f -> Tls_failure f) ;
          write_t t resp; read_react t
    in

    match t.state with
    | `Error e  -> raise e
    | `Closed
    | `Read_closed _ -> raise End_of_file
    | _ ->
        match Flow.single_read t.flow t.recv_buf with
          | exception End_of_file ->
            t.state <- half_close t.state `read;
            raise End_of_file
          | exception exn ->
            (match t.state with
             | `Error _ -> ()
             | _ -> t.state <- `Error exn) ;
            raise exn
          | n ->
            match t.state with
            | `Error e -> raise e
            | `Active tls | `Read_closed tls | `Write_closed tls ->
              handle tls (Cstruct.to_string t.recv_buf ~off:0 ~len:n)
            | `Closed -> raise End_of_file

  let rec single_read t buf =

    let writeout res =
      let open Cstruct in
      let rlen = length res in
      let n    = min (length buf) rlen in
      blit res 0 buf 0 n ;
      t.linger <-
        (if n < rlen then Some (sub res n (rlen - n)) else None) ;
      n in

    match t.linger with
    | Some res -> writeout res
    | None     ->
        match read_react t with
          | None     -> single_read t buf
          | Some res -> writeout res

  let writev t css =
    match t.state with
    | `Error err  -> raise err
    | `Write_closed _ | `Closed -> raise (Eio.Net.err (Connection_reset Tls_socket_closed))
    | `Active tls | `Read_closed tls ->
        let css = List.map Cstruct.to_string css in
        match Tls.Engine.send_application_data tls css with
        | Some (tls, tlsdata) ->
            ( t.state <- inject_state tls t.state ; write_t t tlsdata )
        | None -> invalid_arg "tls: write: socket not ready"

  let single_write t bufs =
    writev t bufs;
    Cstruct.lenv bufs

  (*
   * XXX bad XXX
   * This is a point that should particularly be protected from concurrent r/w.
   * Doing this before a `t` is returned is safe; redoing it during rekeying is
   * not, as the API client already sees the `t` and can mistakenly interleave
   * writes while this is in progress.
   * *)
  let rec drain_handshake t =
    let push_linger t mcs =
      match (mcs, t.linger) with
      | (None, _)         -> ()
      | (scs, None)       -> t.linger <- scs
      | (Some cs, Some l) -> t.linger <- Some (Cstruct.append l cs)
    in
    match t.state with
    | `Active tls when not (Tls.Engine.handshake_in_progress tls) ->
        t
    | _ ->
        let cs = read_react t in
        push_linger t cs; drain_handshake t

  let reneg ?authenticator ?acceptable_cas ?cert ?(drop = true) t =
    match t.state with
    | `Error err  -> raise err
    | `Closed | `Read_closed _ | `Write_closed _ -> invalid_arg "tls: closed socket"
    | `Active tls ->
        match Tls.Engine.reneg ?authenticator ?acceptable_cas ?cert tls with
        | None -> invalid_arg "tls: can't renegotiate"
        | Some (tls', buf) ->
           if drop then t.linger <- None ;
           t.state <- inject_state tls' t.state ;
           write_t t buf;
           ignore (drain_handshake t : t)

  let key_update ?request t =
    match t.state with
    | `Error err  -> raise err
    | `Write_closed _ | `Closed -> invalid_arg "tls: closed socket"
    | `Active tls | `Read_closed tls ->
      match Tls.Engine.key_update ?request tls with
      | Error f -> Fmt.invalid_arg "tls: can't update key: %a" Tls.Engine.pp_failure f
      | Ok (tls', buf) ->
        t.state <- inject_state tls' t.state ;
        write_t t buf

  let shutdown t = function
    | `Receive -> ()
    | `Send | `All ->
      match t.state with
      | `Active tls | `Read_closed tls ->
        let tls', buf = Tls.Engine.send_close_notify tls in
        t.state <- inject_state tls' (half_close t.state `write) ;
        write_t t buf
      | _ -> ()

  let server_of_flow config flow =
    drain_handshake {
      state    = `Active (Tls.Engine.server config) ;
      flow     = (flow :> [Flow.two_way_ty | Eio.Resource.close_ty] r) ;
      linger   = None ;
      recv_buf = Cstruct.create 4096
    }

  let client_of_flow config ?host flow =
    let config' = match host with
      | None -> config
      | Some host -> Tls.Config.peer config host
    in
    let (tls, init) = Tls.Engine.client config' in
    let t = {
      state    = `Active tls ;
      flow     = (flow :> [Flow.two_way_ty | Eio.Resource.close_ty] r);
      linger   = None ;
      recv_buf = Cstruct.create 4096
    } in
    write_t t init;
    drain_handshake t


  let epoch t =
    match t.state with
    | `Active tls | `Read_closed tls | `Write_closed tls -> Tls.Engine.epoch tls
    | `Closed | `Error _ -> Error ()

  let copy t ~src = Eio.Flow.Pi.simple_copy ~single_write t ~src

  let read_methods = []

  let close t = Eio.Resource.close t.flow

  type (_, _, _) Eio.Resource.pi += T : ('t, 't -> t, ty) Eio.Resource.pi
end

let raw (Eio.Resource.T (t, ops)) = Eio.Resource.get ops Raw.T t

let handler =
  Eio.Resource.handler [
    H (Eio.Flow.Pi.Source, (module Raw));
    H (Eio.Flow.Pi.Sink, (module Raw));
    H (Eio.Flow.Pi.Shutdown, (module Raw));
    H (Eio.Resource.Close, Raw.close);
    H (Raw.T, Fun.id);
  ]

let of_t t = Eio.Resource.T (t, handler)

let server_of_flow config       flow = Raw.server_of_flow config       flow |> of_t
let client_of_flow config ?host flow = Raw.client_of_flow config ?host flow |> of_t

let reneg ?authenticator ?acceptable_cas ?cert ?drop (t:t) = Raw.reneg ?authenticator ?acceptable_cas ?cert ?drop (raw t)
let key_update ?request (t:t) = Raw.key_update ?request (raw t)
let epoch (t:t) = Raw.epoch (raw t)

let () =
  Printexc.register_printer (function
      | Tls_alert typ ->
        Some ("TLS alert from peer: " ^ Tls.Packet.alert_type_to_string typ)
      | Tls_failure f ->
        Some ("TLS failure: " ^ Tls.Engine.string_of_failure f)
      | _ -> None)