(**************************************************************************)
(*  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 parse_filter_line line =
  match String.split_on_char ':' line with
  | [ pkg; filtered_deps ] -> (
      match String.split_on_char '[' filtered_deps with
      | [] -> assert false
      | filtered_deps :: all ->
          ( Package.Name.of_string pkg,
            filtered_deps |> String.trim |> String.split_on_char ' '
            |> List.map Package.Name.of_string,
            all = [ "all]" ] ))
  | _ -> Printf.ksprintf failwith "bad filter line: %s" line

let compute_graph data config =
  let ignore_indep = Dependencies.get_ignore_indep config in
  let { src_map; bin_map } = filter_affected data is_affected config in
  let raw_graph = Dependencies.get_dep_graph ~ignore_indep src_map bin_map in
  match StringMap.find_opt "filter" config with
  | None -> Lwt.return (raw_graph, S.empty)
  | Some (Types.EString filter) ->
      let* lines = Lwt_io.lines_of_file filter |> Lwt_stream.to_list in
      List.fold_left
        (fun (graph, all) line ->
          let pkg, filtered_deps, all' = parse_filter_line line in
          let graph =
            match M.find_opt pkg graph with
            | None -> graph
            | Some deps ->
                let new_deps =
                  List.fold_left
                    (fun accu x -> S.remove x accu)
                    deps filtered_deps
                in
                M.add pkg new_deps graph
          in
          let all = if all' then S.add pkg all else all in
          (graph, all))
        (raw_graph, S.empty) lines
      |> Lwt.return
  | _ -> Lwt.fail_with "invalid filter in configuration"

type state = {
  mutable state : (Package.source, bool Lwt.t) M.t;
  pool : unit Lwt_pool.t;
  mutable building : Package.source Package.Name.t list;
  mutex : Lwt_mutex.t;
}

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 = [];
    mutex = Lwt_mutex.create ();
  }

let string_of_state_state state =
  state |> M.bindings
  |> List.filter_map (fun (pkg, t) ->
      match Lwt.state t with Sleep -> Some !!!pkg | Return _ | Fail _ -> None)
  |> String.concat " "
  |> Printf.sprintf "pending = {%s}"

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
  Printf.sprintf "%s\nis_locked = %b\nbuilding = {%s}" state is_locked building

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 rebuild dep_graph all config =
  let rebuild_command =
    let cmd = run_command config "rebuild_command" in
    fun pkg ->
      let no_arch_all = if S.mem pkg all then [ "--no-arch-all" ] else [] in
      cmd (!!!pkg :: no_arch_all)
  in
  let check_command = run_command config "check_command" in
  let promote_command = run_command config "promote_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 (Lwt.return (i = 0)) accu)
        | Some 1 -> Lwt.return 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@ () = 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* result =
      if !global_stop then Lwt.return_none
      else
        let* r1 = rebuild_command pkg in
        let* r2 = promote_command [ !!!pkg ] in
        Lwt.return_some (r1 = Some 0 && r2 = Some 0)
    in
    let* () =
      let@ () = Lwt_mutex.with_lock state.mutex in
      let building = List.filter (fun x -> x <> pkg) state.building in
      let* () =
        let result =
          match result with None -> "skipped" | Some b -> string_of_bool b
        in
        debugf 0 "[END] %s (%s) [%s]" !!!pkg result
          (building |> List.map ( !!! ) |> String.concat " ")
      in
      state.building <- building;
      Lwt.return_unit
    in
    Lwt.return (result = Some true)
  in
  let rec build_pkg closure pkg =
    match S.mem pkg closure with
    | true ->
        (* build cycle detected *)
        let* () = debugf 0 "[CYCLE] %s" !!!pkg in
        Lwt.return_false
    | false -> (
        match M.find_opt pkg state.state with
        | Some t -> t
        | None ->
            let closure = S.add pkg closure in
            let deps = S.elements (M.find pkg dep_graph) in
            let t =
              let* b = Lwt_list.for_all_p (build_pkg closure) deps in
              if b then rebuild pkg else Lwt.return_false
            in
            state.state <- M.add pkg t state.state;
            t)
  in
  let* _ = Lwt_list.map_p (build_pkg S.empty) pkgs in
  let* () = debugf 0 "[FINISH]" in
  let unsuccessful pkg =
    match Lwt.state (M.find pkg state.state) with
    | Return 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 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, all = compute_graph data config in
  let* dep_graph = rebuild dep_graph all 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 () =
  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);
  Lwt_main.run (lwt_main ())

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