open Grpc_lwt
open Lwt.Syntax
include Derivation

module Rust = struct
  include Rust
end

exception Error of H2.Client_connection.error

let rec walk_dir_rec encode f buf path path_name =
  let open Ocaml_protoc_plugin in
  let* dir = Lwt_unix.opendir path in
  Lwt.finalize
    (fun () ->
      let open Elpegrpc.Elpe in
      let rec walk () =
        Lwt.catch
          (fun () ->
            let* entry = Lwt_unix.readdir dir in
            let entry : string = entry in
            if entry = ".." || entry = "." then walk ()
            else
              let path = Filename.concat path entry in
              let path_name = Filename.concat path_name entry in
              let* stat = Lwt_unix.lstat path in
              let* () =
                match stat.st_kind with
                | Unix.S_DIR ->
                    let req =
                      AddPathRequest.make
                        ~request:
                          (`Directory
                             (AddPathRequest.Directory.make ~name:path_name
                                ~permissions:0o644 ()))
                        ()
                    in
                    let enc = encode req |> Writer.contents in
                    f (Some enc);
                    walk_dir_rec encode f buf path path_name
                | Unix.S_REG ->
                    let* file = Lwt_unix.openfile path [ O_RDONLY ] 0 in
                    let ff =
                      AddPathRequest.File.make ~name:path_name
                        ~length:stat.st_size ~permissions:0o644 ()
                    in
                    let req = AddPathRequest.make ~request:(`File ff) () in
                    let enc = encode req |> Writer.contents in
                    let () = f (Some enc) in
                    let rec read_all n =
                      let* r = Lwt_unix.read file buf 0 4096 in
                      if r != 0 then
                        let req =
                          AddPathRequest.make
                            ~request:
                              (`Contents
                                 (AddPathRequest.FileContents.make ~start:n
                                    ~content:(Bytes.sub buf 0 r) ()))
                            ()
                        in
                        let enc = encode req |> Writer.contents in
                        let () = f (Some enc) in
                        read_all (n + r)
                      else Lwt.return ()
                    in
                    read_all 0
                | _ -> Lwt.return ()
              in
              walk ())
          (function End_of_file -> Lwt.return () | e -> Lwt.fail e)
      in
      walk ())
    (fun () -> Lwt_unix.closedir dir)

let add_path connection path0 =
  let open Ocaml_protoc_plugin in
  let open Elpegrpc.Elpe in
  let encode, decode = Service.make_client_functions Elpe.addPath in

  Client.call ~service:"elpe.Elpe" ~rpc:"AddPath"
    ~do_request:
      (H2_lwt_unix.Client.request connection ~error_handler:(fun _ ->
           failwith "Error"))
    ~handler:
      (Client.Rpc.client_streaming ~f:(fun f response ->
           let buf = Bytes.create 4096 in
           let* _ = walk_dir_rec encode f buf path0 "" in
           f None;

           let+ decoder = response in
           match decoder with
           | Some decoder -> (
               Reader.create decoder |> decode |> function
               | Ok v -> v
               | Error e ->
                   failwith
                     (Printf.sprintf "Could not decode request: %s"
                        (Result.show_error e)))
           | None -> Elpe.Derivation.Response.make ()))
    ()

let local_src p =
  object
    inherit derivation
    method name = Filename.basename p

    method! build =
      let c =
        match !backend_conn with None -> failwith "no conn" | Some c -> c
      in
      let* res = add_path c p in
      let res, _ = Result.get_ok res in
      match res with
      | `Ok r -> Lwt.return { destdir = r.destdir; paths = r.paths }
      | `Error e -> failwith e
      | _ -> assert false
  end

let last_built_module : std_derivation option ref = ref None
let build (spec : std_derivation) = last_built_module := Some spec