(**************************************************************************)
(*  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 Ben
open Core

let ( let@ ) x f = x f
let input_source = ref 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 -> Error.raise (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 = 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 Modules
open Marshallable
open 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 := 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 ignore_indep = Dependencies.get_ignore_indep config in
  let { src_map = sources; bin_map = binaries } =
    filter_affected data is_affected config
  in
  let src_of_bin : (Package.binary, Package.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 ~ignore_indep sources src_of_bin in
  match StringMap.find_opt "filter" config with
  | None -> Lwt.return raw_graph
  | Some (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
              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)
          | _ -> 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 : (Package.source, status) M.t;
  pool : unit Lwt_pool.t;
  mutable building : Package.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
    (fun xs ->
      let packages = List.sort (fun x y -> compare !!!x !!!y) xs in
      List.iter
        (fun src ->
          let deps = M.find src dep_graph in
          print_dep_line oc src deps)
        packages)
    rounds

let lwt_main () =
  let config =
    match !input_source with
    | Types.NoSource -> Error.raise Error.Missing_configuration_file
    | _ as source -> Frontend.read_config ~multi:true source
  in
  let archs_list =
    Frontend.to_string_l "architectures"
      (Clflags.get_config config "architectures")
  in
  let data = 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 String.ends_with file ~suffix:".ben" then input_source := Types.File file

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