(**************************************************************************)
(*  Copyright © 2009-2022 Stéphane Glondu <steph@glondu.net>              *)
(*                                                                        *)
(*  This program is free software: you can redistribute it and/or modify  *)
(*  it under the terms of the GNU Affero 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 additional   *)
(*  exemption that compiling, linking, and/or using OpenSSL is allowed.   *)
(*                                                                        *)
(*  This program 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     *)
(*  Affero General Public License for more details.                       *)
(*                                                                        *)
(*  You should have received a copy of the GNU Affero General Public      *)
(*  License along with this program.  If not, see                         *)
(*  <http://www.gnu.org/licenses/>.                                       *)
(**************************************************************************)

open Printf
open Lwt.Syntax
open Benl_core

let ( let@ ) x f = x f

let input_source = ref Benl_types.NoSource
let output_file = ref None
let nproc = ref 1

let get_config config key =
  try StringMap.find key config
  with Not_found -> Benl_error.raise (Benl_error.Missing_configuration_item key)

let is_affected config =
  lazy (Query.of_expr (get_config config "is_affected"))

let run_command config name =
  let c = Benl_frontend.to_string name (get_config config name) in
  fun args ->
  let args = Array.of_list (c :: args) in
  let* r = Lwt_process.exec ~stdin:`Dev_null ~stdout:`Dev_null ~stderr:`Dev_null (c, args) in
  match r with
  | Unix.WEXITED i -> Lwt.return_some i
  | _ -> Lwt.return_none

open Benl_modules
open Marshallable
open Benl_data

let print_dep_line oc src deps =
  fprintf oc "%s:" !!!src;
  S.iter (fun dep -> fprintf oc " %s" !!!dep) deps;
  fprintf oc "\n%!"

let spec =
  Arg.align [
      "-stdin",
      Arg.Unit (fun () -> input_source := Benl_types.Stdin),
      " Use stdin to read the input file";
      "--output",
      Arg.String (fun x -> output_file := Some x),
      " Path to output file";
      "-o",
      Arg.String (fun x -> output_file := Some x),
      " Path to output file";
      "-j",
      Arg.Int (fun x -> nproc := x),
      " Maximum number of jobs";
    ]

let compute_graph data config =
  let {src_map = sources; bin_map = binaries} = filter_affected data is_affected config in
  let src_of_bin : ([`binary], [`source] Package.Name.t) M.t =
    PAMap.fold
      (fun (name, _) pkg accu ->
         let source = Package.get "source" pkg in
         M.add name (Package.Name.of_string source) accu)
      binaries
      M.empty
  in
  Dependencies.get_dep_graph sources src_of_bin

module type MUTEX = sig
  type 'a t
  val create : unit -> 'a t
  val lock : 'a t -> 'a option Lwt.t
  val unlock : 'a t -> ('a, exn) result -> 'a Lwt.t
end

module Mutex : MUTEX = struct
  type 'a t = { mutable locked : bool; waiters : 'a option Lwt.u Queue.t }
  let create () = { locked = false; waiters = Queue.create () }
  let lock m =
    if m.locked then (
      let t, u = Lwt.task () in
      Queue.push u m.waiters;
      t
    ) else (
      m.locked <- true;
      Lwt.return_none
    )
  let unlock m r =
    if m.locked then (
      match Queue.take_opt m.waiters with
      | None -> m.locked <- false; Lwt.of_result r
      | Some u ->
         match r with
         | Ok r -> Lwt.wakeup_later u (Some r); Lwt.return r
         | Error e -> Lwt.wakeup_later_exn u e; Lwt.fail e
    ) else Lwt.of_result r
end

let rebuild dep_graph config =
  let pool = Lwt_pool.create !nproc Lwt.return in
  let mutex = Mutex.create () in
  let pending = ref Lwt.return_unit in
  let push_pending x =
    pending := (let* () = !pending in let* _ = x in Lwt.return_unit);
    x
  in
  let rebuild_command = run_command config "rebuild_command" in
  let check_command = run_command config "check_command" in
  let stop_command = run_command config "stop_command" in
  let promote_command = run_command config "promote_command" in
  let update_command = run_command config "update_command" in
  let pkgs = M.bindings dep_graph |> List.map fst in
  let* initial_state =
    Lwt_list.fold_left_s (fun accu x ->
        let* b = check_command [!!!x] in
        match b with
        | Some ((0 | 2) as i) -> Lwt.return (M.add x (`Built (i = 0)) accu)
        | _ -> Lwt.return accu
      ) M.empty pkgs
  in
  let state = ref initial_state in
  let rec build closure pkg =
    if S.mem pkg closure then (
      (* dependency cycle detected *)
      Lwt.return (`Built, false)
    ) else (
      match (try Some (M.find pkg !state) with Not_found -> None) with
      | None ->
         let t, u = Lwt.task () in
         state := M.add pkg (`Building t) !state;
         let* b =
           let deps = M.find pkg dep_graph in
           let* rs =
             let closure = S.add pkg closure in
             Lwt_list.map_p (fun x ->
                 let* b = build closure x in
                 Lwt.return (x, b)
               ) (S.elements deps)
           in
           let successful = List.for_all (fun (_, (_, b)) -> b) rs in
           let justbuilt =
             List.filter_map
               (function
                | (pkg, (`JustBuilt, _)) -> Some pkg
                | _ -> None
               ) rs
           in
           let* justbuilt_processed =
             if justbuilt <> [] then (
               let* update_needed =
                 Lwt_list.fold_left_s (fun accu pkg ->
                     let* r = push_pending @@ promote_command [!!!pkg] in
                     Lwt.return (accu || r <> Some 0)
                   ) false justbuilt
               in
               if update_needed then (
                 let* previous_update = Mutex.lock mutex in
                 Lwt.catch (fun () ->
                     match previous_update with
                     | None ->
                        let* () = !pending in
                        let* r = update_command [] in
                        Mutex.unlock mutex (Ok (r = Some 0))
                     | Some b -> Mutex.unlock mutex (Ok b)
                   ) (fun e -> Mutex.unlock mutex (Error e))
               ) else Lwt.return_true
             ) else Lwt.return_true
           in
           if successful && justbuilt_processed then (
             let@ () = Lwt_pool.use pool in
             let@ () = fun cont ->
               let* r = stop_command [] in
               if r = Some 0 then Lwt.return (`Built, false)
               else cont ()
             in
             let* () =
               let open Lwt_io in
               write_line stderr (Printf.sprintf "[BGN] %s" !!!pkg)
             in
             let* r = push_pending @@ rebuild_command [!!!pkg] in
             let* () =
               let open Lwt_io in
               let result = if r = Some 0 then "success" else "failure" in
               write_line stderr (Printf.sprintf "[END] %s (%s)" !!!pkg result)
             in
             Lwt.return (`JustBuilt, r = Some 0)
           ) else (
             (* one of the dependencies failed *)
             Lwt.return (`Built, false)
           )
         in
         state := M.add pkg (`Built (snd b)) !state;
         Lwt.wakeup_later u b;
         Lwt.return b
      | Some (`Building t) -> t
      | Some (`Built b) -> Lwt.return (`Built, b)
    )
  in
  let* rs =
    Lwt_list.map_p (fun pkg ->
        let* r = build S.empty pkg in
        Lwt.return (pkg, r)
      ) pkgs
  in
  let* needs_update =
    Lwt_list.fold_left_s (fun accu (pkg, (r, _)) ->
        match r with
        | `JustBuilt ->
           let* r = promote_command [!!!pkg] in
           Lwt.return (accu || r <> Some 0)
        | `Built -> Lwt.return accu
      ) false rs
  in
  let* () =
    if needs_update then (
      let* _ = update_command [] in
      Lwt.return_unit
    ) else Lwt.return_unit
  in
  let unsuccessful pkg =
    match M.find pkg !state with
    | `Building _ -> assert false
    | `Built b -> not b
  in
  M.fold
    (fun pkg deps accu ->
      if unsuccessful pkg then (
        let deps = S.filter unsuccessful deps in
        M.add pkg deps accu
      ) else accu
    ) dep_graph M.empty
  |> Lwt.return

let print_dependency_levels oc dep_graph rounds =
  List.iter begin fun xs ->
    let packages = List.sort (fun x y -> compare !!!x !!!y) xs in
    List.iter begin fun src ->
      let deps = M.find src dep_graph in
      print_dep_line oc src deps
    end packages
  end rounds

let lwt_main () =
  let config = match !input_source with
    | Benl_types.NoSource -> Benl_error.raise Benl_error.Missing_configuration_file
    | _ as source -> Benl_frontend.read_config ~multi:true source
  in
  let archs_list = Benl_frontend.to_string_l
    "architectures"
    (Benl_clflags.get_config config "architectures")
  in
  let data = Benl_data.load_cache archs_list in
  let dep_graph = compute_graph data config in
  let* dep_graph = rebuild dep_graph config in
  let rounds = Dependencies.topo_split dep_graph in
  let oc, close = match !output_file with
    | None -> stdout, fun () -> ()
    | Some x -> let oc = open_out x in oc, fun () -> close_out oc
  in
  print_dependency_levels oc dep_graph rounds;
  close ();
  Lwt.return_unit

let main () =
  Lwt_main.run (lwt_main ())

let anon_fun file =
  if Benl_core.ends_with file ".ben" then
    input_source := Benl_types.File file

let frontend = {
  Benl_frontend.name = "rebuild";
  Benl_frontend.main = main;
  Benl_frontend.anon_fun = anon_fun;
  Benl_frontend.help = spec;
}
