open Lwt.Infix

let safe_close s =
  Lwt.catch (fun () -> Lwt_unix.close s)
    (fun _ -> Lwt.return_unit)

let read_wire fd =
  let buf = Bytes.create 4 in
  let rec r b i l =
    Lwt.catch (fun () ->
        Lwt_unix.read fd b i l >>= function
        | 0 ->
          Lwt.return (Error `Eof)
        | n when n == l -> Lwt.return (Ok ())
        | n when n < l -> r b (i + n) (l - n)
        | _ ->
          Lwt.return (Error `Toomuch))
      (fun _ ->
         safe_close fd >|= fun () ->
         Error `Exception)
  in
  r buf 0 4 >>= function
  | Error e -> Lwt.return (Error e)
  | Ok () ->
    let len = Cstruct.BE.get_uint32 (Cstruct.of_bytes buf) 0 in
    if len > 0l then begin
      let b = Bytes.create (Int32.to_int len) in
      r b 0 (Int32.to_int len) >|= function
      | Error e -> Error e
      | Ok () ->
        match Wire_asn.wire_of_cstruct (Cstruct.of_bytes b) with
        | Error (`Msg _msg) ->
          Error `Exception
        | (Ok (hdr, _)) as w ->
          if not (hdr.version = `V1) then
            assert false; (* FIXME *)
          w
    end else begin
      Lwt.return (Error `Eof)
    end

let write_raw fd buf =
  let rec w off l =
    Lwt.catch (fun () ->
        Lwt_unix.send fd buf off l [] >>= fun n ->
        if n = l then
          Lwt.return (Ok ())
        else
          w (off + n) (l - n))
      (fun _ ->
         safe_close fd >|= fun () ->
         Error `Exception)
  in
  w 0 (Bytes.length buf)

let write_wire fd wire =
  let data = Wire_asn.wire_to_cstruct wire in
  let dlen = Cstruct.create 4 in
  Cstruct.BE.set_uint32 dlen 0 (Int32.of_int (Cstruct.length data)) ;
  let buf = Cstruct.(to_bytes (append dlen data)) in
  write_raw fd buf