open Lwt.Syntax

exception DerivationError of string

type store_path = { path : string }
type build_result = { destdir : string list; paths : string list }

let backend_conn = ref None

class virtual derivation =
  object
    method virtual name : string
    val stdout_ : Buffer.t option ref = ref None
    val stderr_ : Buffer.t option ref = ref None
    method stdout = stdout_
    method stderr = stderr_

    method set_stdout =
      let buf = Buffer.create 1024 in
      stdout_ := Some buf;
      buf

    method set_stderr =
      let buf = Buffer.create 1024 in
      stderr_ := Some buf;
      buf

    (* Returns the script to setup the build process, also used in shells. *)
    method setup : string Lwt.t = Lwt.return ""

    (* List of paths resulting from building the package. *)
    method build : build_result Lwt.t = failwith "Not implemented"
  end

open Grpc_lwt

let ubuntu_releases = Hashtbl.create 8
let ubuntu_release_lock = Lwt_mutex.create ()

let ubuntu_release connection ~release ~arch ~repository =
  Lwt_mutex.with_lock ubuntu_release_lock (fun () ->
      match Hashtbl.find_opt ubuntu_releases (release, arch, repository) with
      | None ->
          let req =
            Elpegrpc.Elpe.UbuntuReleaseRequest.make ~release ~arch ~repository
              ()
          in
          let open Ocaml_protoc_plugin in
          let encode, decode =
            Service.make_client_functions Elpegrpc.Elpe.Elpe.ubuntuRelease
          in
          let enc = encode req |> Writer.contents in
          let* r =
            Client.call ~service:"elpe.Elpe" ~rpc:"UbuntuRelease"
              ~do_request:
                (H2_lwt_unix.Client.request connection ~error_handler:(fun _ ->
                     failwith "Error"))
              ~handler:
                (Client.Rpc.unary enc ~f:(fun decoder ->
                     let+ decoder = decoder 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 -> Elpegrpc.Elpe.Elpe.Derivation.Response.make ()))
              ()
          in
          Hashtbl.add ubuntu_releases (release, arch, repository) r;
          Lwt.return r
      | Some x -> Lwt.return x)

let link_extra =
  [ ("libstd-rust-.*-dev", "libstd-rust.*"); ("clang", "clang-.*") ]

let ubuntu_packages = Hashtbl.create 64
let ubuntu_packages_locks = Hashtbl.create 64

let ubuntu_package connection ~index ~name =
  let lock =
    try Hashtbl.find ubuntu_packages_locks (index, name)
    with Not_found ->
      let lock = Lwt_mutex.create () in
      Hashtbl.add ubuntu_packages_locks (index, name) lock;
      lock
  in
  Lwt_mutex.with_lock lock (fun () ->
      match Hashtbl.find_opt ubuntu_packages (index, name) with
      | Some x -> Lwt.return x
      | None ->
          let link_extra =
            List.map
              (fun (pkg, dep) -> Elpegrpc.Elpe.LinkExtra.make ~pkg ~dep ())
              link_extra
          in
          let req =
            Elpegrpc.Elpe.UbuntuPackageRequest.make ~index ~name ~link_extra ()
          in
          let open Ocaml_protoc_plugin in
          let encode, decode =
            Service.make_client_functions Elpegrpc.Elpe.Elpe.ubuntuPackage
          in
          let enc = encode req |> Writer.contents in
          let result = ref None in
          let* x =
            Client.call ~service:"elpe.Elpe" ~rpc:"UbuntuPackage"
              ~do_request:
                (H2_lwt_unix.Client.request connection ~error_handler:(fun _ ->
                     failwith "Error"))
              ~handler:
                (Client.Rpc.server_streaming enc ~f:(fun responses ->
                     Lwt_stream.iter_s
                       (fun str ->
                         Reader.create str |> decode |> function
                         | Ok (`Ok v) ->
                             result := Some v;
                             Lwt.return ()
                         | Ok (`Loading v) ->
                             Printf.eprintf "Downloading %s\n%!" v;
                             Lwt.return ()
                         | Error e ->
                             failwith
                               (Printf.sprintf "Could not decode request: %s"
                                  (Result.show_error e))
                         | _ -> Lwt.return ())
                       responses))
              ()
          in
          let result =
            match x with Ok _ -> Ok (Option.get !result) | Error e -> Error e
          in
          Hashtbl.add ubuntu_packages (index, name) result;
          Lwt.return result)

let ubuntu ?(release = "plucky") ?(platform = !Platform.target) name =
  object
    inherit derivation
    method name = name
    val cached_build = ref None
    val build_lock = Lwt_mutex.create ()

    method! build =
      match !cached_build with
      | Some x -> Lwt.return x
      | None ->
          Lwt_mutex.with_lock build_lock (fun () ->
              match !cached_build with
              | Some x -> Lwt.return x
              | None ->
                  let c =
                    match !backend_conn with
                    | None -> failwith "no conn"
                    | Some c -> c
                  in
                  let* index =
                    Lwt_list.map_p
                      (fun repository ->
                        let arch =
                          match platform.arch with
                          | Amd64 -> Elpegrpc.Elpe.Arch.Amd64
                          | Aarch64 -> Elpegrpc.Elpe.Arch.Aarch64
                        in
                        let* res =
                          ubuntu_release c ~release ~arch ~repository
                        in
                        let res, _ = Result.get_ok res in
                        match res with
                        | `Ok r -> Lwt.return r.destdir
                        | `Error e -> raise (DerivationError e)
                        | _ -> assert false)
                      [ "main"; "universe" ]
                  in
                  let index = List.concat index in
                  let* res = ubuntu_package c ~index ~name in
                  let r = Result.get_ok res in
                  let x = { destdir = r.destdir; paths = r.paths } in
                  cached_build := Some x;
                  Lwt.return x)
  end

let derivation ?(stdout = None) ?(stderr = None) connection ~name ~builder
    ~paths ~target ~output_hash =
  let req =
    match output_hash with
    | None ->
        Elpegrpc.Elpe.DerivationRequest.make ~name ~builder ~paths ~target ()
    | Some output_hash ->
        Elpegrpc.Elpe.DerivationRequest.make ~name ~builder ~paths ~target
          ~output_hash ()
  in
  let open Ocaml_protoc_plugin in
  let encode, decode =
    Service.make_client_functions Elpegrpc.Elpe.Elpe.derivation
  in
  let enc = encode req |> Writer.contents in
  let result = ref None in
  let* _ =
    Client.call ~service:"elpe.Elpe" ~rpc:"Derivation"
      ~do_request:
        (H2_lwt_unix.Client.request connection ~error_handler:(fun _ ->
             failwith "Error"))
      ~handler:
        (Client.Rpc.server_streaming enc ~f:(fun responses ->
             Lwt_stream.iter_s
               (fun str ->
                 Reader.create str |> decode |> function
                 | Ok (`Ok path) ->
                     result := Some (`Ok path);
                     Lwt.return ()
                 | Ok (`Error path) ->
                     result := Some (`Error path);
                     Lwt.return ()
                 | Ok (`Stdout buf) ->
                     let _ =
                       match stdout with
                       | Some b -> Buffer.add_bytes b buf
                       | None ->
                           let buf = Bytes.to_string buf in
                           Printf.printf "%s%!" buf
                     in
                     Lwt.return ()
                 | Ok (`Stderr buf) ->
                     let _ =
                       match stderr with
                       | Some b -> Buffer.add_bytes b buf
                       | None ->
                           let buf = Bytes.to_string buf in
                           Printf.eprintf "%s%!" buf
                     in
                     Lwt.return ()
                 | Error e ->
                     failwith
                       (Printf.sprintf "Could not decode request: %s"
                          (Result.show_error e))
                 | _ -> Lwt.return ())
               responses))
      ()
  in
  match !result with Some v -> Lwt.return v | None -> assert false

class virtual std_derivation =
  object (self)
    inherit derivation
    method environment : (string * string) List.t Lwt.t = Lwt.return []
    method build_inputs : derivation list Lwt.t = Lwt.return []
    method target : string = "x86_64-linux-gnu"
    val extra_paths : string list ref = ref []
    val extra_paths_h : (string, unit) Hashtbl.t = Hashtbl.create 16
    method output_hash : string option Lwt.t = Lwt.return None

    method add_path r =
      if Hashtbl.find_opt extra_paths_h r = None then
        let _ = Hashtbl.add extra_paths_h r () in
        extra_paths := r :: !extra_paths
      else ()

    method derivation (drv : derivation) =
      let* path = drv#build in
      List.iter self#add_path path.destdir;
      Lwt.return path

    method src : derivation Lwt.t = failwith ("No src defined for " ^ self#name)
    val post_setup = ""

    method! setup =
      let* bash = self#derivation (ubuntu "bash-static") in
      let bash = List.hd bash.destdir in
      let* build_inputs = self#build_inputs in
      let* paths =
        Lwt_list.map_p
          (fun d ->
            let* p = d#build in
            Lwt.return p.destdir)
          build_inputs
      in
      List.iter (List.iter self#add_path) paths;

      let path =
        List.fold_left
          (List.fold_left (fun acc path ->
               (if acc = "" then acc ^ path else acc ^ ":" ^ path) ^ "/usr/bin"))
          "" paths
      in
      Lwt.return
        ("#!" ^ bash ^ "/usr/bin/bash-static\nexport PATH=" ^ path
       ^ "\nexport LIBRARY_PATH=/lib/" ^ self#target ^ ":/lib64/" ^ self#target
       ^ ":/usr/lib:/usr/lib64:/usr/lib/x86_64-linux-gnu\n" ^ post_setup ^ "\n"
        )

    val pre_unpack = ""
    val post_unpack = ""

    method unpack_phase =
      let* src = self#src in
      let* src_built = src#build in
      List.iter self#add_path src_built.destdir;

      Lwt.return
        (pre_unpack ^ "\ncp -R " ^ List.hd src_built.destdir ^ "/. .\n"
       ^ post_unpack)

    val pre_configure = ""
    val post_configure = ""

    method configure_phase =
      Lwt.return
        (pre_configure ^ "\nif [[ -e configure ]]; then ./configure; fi\n"
       ^ post_configure)

    val pre_build = ""
    val post_build = ""

    method build_phase =
      Lwt.return
        (pre_build ^ "\nif [[ -e Makefile ]]; then make; fi\n" ^ post_build)

    val pre_install = ""
    val post_install = ""

    method install_phase =
      Lwt.return
        (pre_install ^ "\nif [[ -e Makefile ]]; then make install; fi\n"
       ^ post_install)

    val cached_build = ref None
    val build_lock = Lwt_mutex.create ()

    method! build =
      Lwt_mutex.with_lock build_lock (fun () ->
          match !cached_build with
          | Some cached -> Lwt.return cached
          | None -> (
              let* setup = self#setup in
              let* unpack_phase = self#unpack_phase in
              let* configure_phase = self#configure_phase in
              let* build_phase = self#build_phase in
              let* install_phase = self#install_phase in
              let phases =
                [
                  setup;
                  unpack_phase;
                  configure_phase;
                  build_phase;
                  install_phase;
                ]
              in
              let builder = String.concat "\n" phases in
              let c =
                match !backend_conn with
                | None -> failwith "no conn"
                | Some c -> c
              in
              let* output_hash = self#output_hash in
              let* r =
                derivation c ~name:self#name ~builder ~paths:!extra_paths
                  ~stdout:!(self#stdout) ~stderr:!(self#stderr)
                  ~target:self#target ~output_hash
              in
              match r with
              | `Ok r ->
                  let c = { destdir = r.destdir; paths = r.paths } in
                  cached_build := Some c;
                  Lwt.return c
              | `Error e -> failwith e
              | _ -> assert false))
  end

let http ~name ~url ~hash_algorithm ~hash =
  object
    inherit derivation
    method name = "http" ^ if name = "" then "" else "-" ^ name

    method! build =
      let connection =
        match !backend_conn with None -> failwith "no conn" | Some c -> c
      in
      let req =
        Elpegrpc.Elpe.AddUrlRequest.make ~url ~hash_algorithm ~hash ()
      in
      let* res =
        let open Ocaml_protoc_plugin in
        let encode, decode =
          Service.make_client_functions Elpegrpc.Elpe.Elpe.addUrl
        in
        let enc = encode req |> Writer.contents in
        Client.call ~service:"elpe.Elpe" ~rpc:"AddUrl"
          ~do_request:
            (H2_lwt_unix.Client.request connection ~error_handler:(fun _ ->
                 failwith "Error"))
          ~handler:
            (Client.Rpc.unary enc ~f:(fun decoder ->
                 let+ decoder = decoder 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 -> Elpegrpc.Elpe.Elpe.Derivation.Response.make ()))
          ()
      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 unzip ?(name = "") pkg =
  object (self)
    inherit std_derivation
    method name = "unzip" ^ if name = "" then "" else "-" ^ name
    method! unpack_phase = Lwt.return ""

    method! build_inputs =
      Lwt.return [ ubuntu "coreutils"; ubuntu "tar"; ubuntu "gzip" ]

    method! build_phase =
      let* zipped = self#derivation pkg in
      Lwt.return
        ("mkdir -p $DESTDIR; tar -xf " ^ List.hd zipped.destdir ^ " -C $DESTDIR")
  end

type cc = Gcc | Clang

let cc x =
  let x = match x with Gcc -> "gcc" | Clang -> "clang" in
  object (self)
    inherit std_derivation
    method name = "cc"
    method! build_inputs = Lwt.return [ ubuntu "coreutils"; ubuntu x ]
    method! unpack_phase = Lwt.return ""

    method! build_phase =
      let* cc = (ubuntu x)#build in
      List.iter self#add_path cc.destdir;
      Lwt.return
        ("mkdir -p $DESTDIR/usr/bin; ln -s " ^ List.hd cc.destdir ^ "/usr/bin/"
       ^ x ^ " $DESTDIR/usr/bin/cc")
  end