(**************************************************************************)
(*  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 debug_level = ref 0

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 ("setsid" :: "--wait" :: c :: args) in
  let* r = Lwt_process.exec ~stdin:`Dev_null ~stdout:`Dev_null ~stderr:`Dev_null ("setsid", 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";
      "-d",
      Arg.Int (fun x -> debug_level := x),
      " Debugging level";
    ]

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
  let raw_graph = Dependencies.get_dep_graph sources src_of_bin in
  match StringMap.find_opt "filter" config with
  | None -> Lwt.return raw_graph
  | Some (Benl_types.EString filter) ->
     let* lines = Lwt_io.lines_of_file filter |> Lwt_stream.to_list in
     List.fold_left
       (fun graph line ->
         match String.split_on_char ':' line with
         | [pkg; filtered_deps] ->
            let pkg = Package.Name.of_string pkg in
            begin match M.find_opt pkg graph with
            | None -> graph
            | Some deps ->
               let new_deps =
                 List.fold_left
                   (fun accu x -> S.remove (Package.Name.of_string x) accu)
                   deps (String.split_on_char ' ' filtered_deps)
               in
               M.add pkg new_deps graph
            end
         | _ -> failwith "syntax error in filter"
       ) raw_graph lines
     |> Lwt.return
  | _ -> Lwt.fail_with "invalid filter in configuration"

type status =
  | Unknown
  | Built of bool
  | Pending of bool
  | Building of unit Lwt.t

type state = {
    mutable state : ([`source], status) M.t;
    pool : unit Lwt_pool.t;
    mutable building : [`source] Package.Name.t list;
    pending_updates : unit Lwt.u Queue.t;
    pending_builds : unit Lwt.u Queue.t;
    mutex : Lwt_mutex.t;
    mutable pending_update : bool Lwt.t option;
    mutable pending_update_contexts : string list;
}

let global_stop = ref false
let global_state = ref None

let make_state state = {
    state;
    pool = Lwt_pool.create !nproc (fun () -> Lwt.return_unit);
    building = [];
    pending_updates = Queue.create ();
    pending_builds = Queue.create ();
    mutex = Lwt_mutex.create ();
    pending_update = None;
    pending_update_contexts = [];
}

let string_of_state_state state =
  let u, b, p =
    M.bindings state
    |> List.rev
    |> List.fold_left
         (fun (u, b, p) (pkg, status) ->
           match status with
           | Unknown -> !!!pkg :: u, b, p
           | Building _ -> u, !!!pkg :: b, p
           | Pending _ -> u, b, !!!pkg :: p
           | Built _ -> u, b, p
         ) ([], [], [])
  in
  Printf.sprintf "unknown = {%s}\nbuilding = {%s}\npending = {%s}"
    (u |> String.concat " ") (b |> String.concat " ") (p |> String.concat " ")

let string_of_state s =
  let state = string_of_state_state s.state in
  let is_locked = Lwt_mutex.is_locked s.mutex in
  let building = s.building |> List.map (!!!) |> String.concat " " in
  let pending_update =
    match s.pending_update with
    | None -> "N"
    | Some _ -> "P"
  in
  let pending_update_contexts = s.pending_update_contexts |> String.concat " " in
  Printf.sprintf
    "%s\nis_locked = %b\nbuilding = {%s}\npending_updates = %s{%s}"
    state is_locked building pending_update pending_update_contexts

let debugf level fmt =
  Printf.ksprintf
    (fun x ->
      if level <= !debug_level then (
        let x =
          let open Unix in
          let now = localtime (gettimeofday ()) in
          Printf.sprintf "%02d:%02d:%02d %s"
            now.tm_hour now.tm_min now.tm_sec x
        in
        let open Lwt_io in
        let* () = write_line stderr x in
        flush stderr
      ) else Lwt.return_unit
    ) fmt

let failf fmt = Printf.ksprintf Lwt.fail_with fmt

let rec exhaust q =
  match Queue.take_opt q with
  | None -> ()
  | Some u -> Lwt.wakeup_later u (); exhaust q

let rebuild dep_graph config =
  let rebuild_command = run_command config "rebuild_command" in
  let check_command = run_command config "check_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)
        | Some 1 -> Lwt.return (M.add x Unknown accu)
        | _ -> failf "could not determine initial state of %s" !!!x
      ) M.empty pkgs
  in
  let state = make_state initial_state in
  global_state := Some state;
  let rebuild pkg =
    let* () = debugf 0 "[SCHEDULE] %s" !!!pkg in
    let* () =
      if Queue.is_empty state.pending_updates then (
        Lwt.return_unit
      ) else (
        let t, u = Lwt.wait () in
        Queue.push u state.pending_builds;
        t
      )
    in
    let@ () = Lwt_pool.use state.pool in
    let* () =
      let@ () = Lwt_mutex.with_lock state.mutex in
      let building = state.building in
      let* () = debugf 0 "[BEGIN] %s [%s]" !!!pkg (building |> List.map (!!!) |> String.concat " ") in
      state.building <- pkg :: state.building;
      Lwt.return_unit
    in
    let* r = rebuild_command [!!!pkg] in
    let result = r = Some 0 in
    let* () =
      let@ () = Lwt_mutex.with_lock state.mutex in
      let building = List.filter (fun x -> x <> pkg) state.building in
      let* () = debugf 0 "[END] %s (%b) [%s]" !!!pkg result (building |> List.map (!!!) |> String.concat " ") in
      state.building <- building;
      if building = [] then exhaust state.pending_updates;
      Lwt.return_unit
    in
    Lwt.return result
  in
  let update () =
    match state.pending_update with
    | Some t -> t
    | None ->
       let t, u = Lwt.wait () in
       state.pending_update <- Some t;
       let@ () = Lwt_mutex.with_lock state.mutex in
       let context = state.pending_update_contexts |> String.concat " " in
       let* () = debugf 0 "[BEGIN UPDATE] {%s}" context in
       let* r = update_command [] in
       let success = r = Some 0 in
       let* () = debugf 0 "[END UPDATE] {%s} (%b)" context success in
       state.pending_update_contexts <- [];
       state.pending_update <- None;
       Lwt.wakeup_later u success;
       t
  in
  let promote context pkgs =
    let* () = debugf 1 "[BEGIN PROMOTIONS] %s" context in
    let* promotions =
      Lwt_list.map_p
        (fun pkg ->
          let@ () = Lwt_mutex.with_lock state.mutex in
          match M.find pkg state.state with
          | Pending b ->
             let* () = debugf 0 "[BEGIN PROMOTE] %s (%b)" !!!pkg b in
             let* r = promote_command [!!!pkg] in
             let b = b && r = Some 0 in
             let* () = debugf 0 "[END PROMOTE] %s (%b)" !!!pkg b in
             state.state <- M.add pkg (Built b) state.state;
             Lwt.return (`Promoted, b)
          | Built b ->
             let* () = debugf 1 "[ALREADY PROMOTED] %s (%b)" !!!pkg b in
             Lwt.return (`Built, b)
          | _ -> Lwt.fail_with "unexpected state in promote"
        ) pkgs
    in
    let promoted = List.filter (fun (x, _) -> x = `Promoted) promotions in
    let success = List.for_all (fun (_, x) -> x) promotions in
    let* success =
      if promoted = [] then (
        Lwt.return success
      ) else (
        state.pending_update_contexts <- context :: state.pending_update_contexts;
        let* r = update () in
        let success = r && success in
        let* () = debugf 0 "[UPDATED] %s (%b)" context success in
        Lwt.return success
      )
    in
    let* () = debugf 1 "[END PROMOTIONS] %s (%b)" context success in
    Lwt.return success
  in
  let rec build_pkg closure pkg =
    let@ () = fun cont ->
      if M.mem pkg closure then (
        (* build cycle detected *)
        let@ () = Lwt_mutex.with_lock state.mutex in
        Lwt_list.iter_p
          (fun (pkg, u) ->
            match M.find pkg state.state with
            | Building _ ->
               let* () = debugf 0 "[CANCEL] %s" !!!pkg in
               state.state <- M.add pkg (Built false) state.state;
               let () =
                 match !u with
                 | None -> ()
                 | Some u' -> u := None; Lwt.wakeup_later u' ()
               in
               Lwt.return_unit
            | Unknown -> assert false
            | Built _ | Pending _ -> Lwt.return_unit
          ) (M.bindings closure)
      ) else cont ()
    in
    match M.find pkg state.state with
    | (Pending _ | Built _) -> Lwt.return_unit
    | Building t -> t
    | Unknown ->
       let t, u = Lwt.wait () in
       let u = ref (Some u) in
       let closure = M.add pkg u closure in
       state.state <- M.add pkg (Building t) state.state;
       let deps = S.elements (M.find pkg dep_graph) in
       let* () = Lwt_list.iter_p (build_pkg closure) deps in
       match !u with
       | None ->
          (* build was interrupted while processing build deps *)
          Lwt.return_unit
       | Some u ->
          let* () =
            if state.building = [] then (
              Lwt.return_unit
            ) else (
              let t, u = Lwt.wait () in
              Queue.push u state.pending_updates;
              t
            )
          in
          let* b = promote !!!pkg deps in
          exhaust state.pending_builds;
          let* () =
            if b && not !global_stop then (
              let* result = rebuild pkg in
              state.state <- M.add pkg (Pending result) state.state;
              Lwt.return_unit
            ) else (
              state.state <- M.add pkg (Built false) state.state;
              Lwt.return_unit
            )
          in
          Lwt.wakeup_later u ();
          Lwt.return_unit
  in
  let* () = Lwt_list.iter_p (build_pkg M.empty) pkgs in
  let* r = promote "<MAIN>" pkgs in
  let unsuccessful pkg =
    match M.find pkg state.state with
    | Built b -> not b
    | _ -> assert false
  in
  let depgraph =
    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
  in
  Lwt.return (r, depgraph)

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* r, 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 r

let main () =
  let handler s =
    if s = Sys.sigint then (
      if not !global_stop then (
        Printf.eprintf "Stop requested, finishing pending jobs...\n%!";
        global_stop := true;
      )
    ) else if s = Sys.sigusr1 then (
      match !global_state with
      | None -> ()
      | Some state ->
         Printf.eprintf
           "-----[BEGIN STATE]-----\n%s\n-----[END STATE]-----\n%!"
           (string_of_state state)
    )
  in
  let handler = Sys.Signal_handle handler in
  Sys.(set_signal sigint handler);
  Sys.(set_signal sigusr1 handler);
  let r = Lwt_main.run (lwt_main ()) in
  if not r then exit 10

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;
}
