let header = "\n  ----- MEAT ----------------------------------------\n"
let footer = "\n  ---------------------------------------------------\n"

let help =
  "\n\
   \tYUM    -  CONSUME DELICIOUS MEATS\n\
   \tCOOK   -  ONLY PREPARE MEATS\n\
   \tPOKE   -  TASTE SUSPICIOUS MEATS\n\
   \tGUT    -  CLEAN MEAT STORES\n\
   \tLOOK   -  INSPECT MEAT\n\
   \tFRESH  -  HUNT FRESH MEATS\n\
   \t..N    -  ..MEATS OF N\n\
   \t..-S N -  ..SUBMEAT N\n\
   \t..-A   -  ..ALL MEATS\n"

open Sys

let pass_args () =
  let len = Array.length argv and sconcat acc el = acc ^ " " ^ el in
  match len with
  | 3 -> argv.(2)
  | n when n > 3 ->
      print_int (n - 1);
      Array.fold_left sconcat "" (Array.sub argv 2 (n - 2))
  | _ -> ""

let do_cmd ?(args = true) cmd =
  match command (if args then cmd ^ pass_args () else cmd) with _ -> ()

let meat_print text = print_endline ("\n   \t" ^ text ^ "\n")

let yum () =
  print_string header;
  meat_print "CONSUMING DELICIOUS MEATS..";
  do_cmd "nh os switch";
  print_string footer

let cook () =
  print_string header;
  meat_print "PREPARING DELICIOUS MEATS..";
  do_cmd "nh os build";
  print_string footer

let poke () =
  print_string header;
  meat_print "PREPARING SUSPICIOUS MEATS..";
  do_cmd "nh os build --show-trace";
  print_string footer

let gut () =
  print_string header;
  meat_print "CLEANING MEAT STORES..";
  do_cmd "nh clean all";
  print_string footer

let look () =
  print_string header;
  meat_print "INSPECTING MEAT..";
  meat_print "CONSUMING MEATS:";
  do_cmd "nix flake metadata $NH_FLAKE";
  meat_print "PRODUCING MEATS:";
  do_cmd "nix flake show $NH_FLAKE";
  print_string footer

(* let all_flag () = Array.mem "-a" argv || Array.mem "--all" argv
let sub_flag () = Array.mem "-s" argv || Array.mem "--subflake" argv *)
let all_flag () = if Array.length argv >= 3 then match Array.get argv 2 with
  | "-a" | "--all" -> true
  | _ -> false
else
  false
let sub_flag () = if Array.length argv >= 3 then match Array.get argv 2 with
  | "-s" | "--sub" -> true
  | _ -> false
else
  false

let derelativise base = List.map (fun a -> base ^ "/" ^ a)

let filter_dirs fullpath dirs =
  dirs |> derelativise fullpath
  |> List.filter (fun d -> (Unix.stat d).st_kind = S_DIR)

let readdir d = try Sys.readdir d with Sys_error _ -> [||]

let walk entry =
  let rec loop dir : string list =
    let contents = readdir dir |> Array.to_list in
    let is_flake = List.mem "flake.nix" contents in
    if dir = entry then
      let subdirs = contents |> filter_dirs dir in
      List.flatten (List.map loop subdirs)
    else if is_flake then
      let subdirs = contents |> filter_dirs dir in
      let children = List.flatten (List.map loop subdirs) in
      [ dir ] @ children
    else []
  in
  loop entry

let countdepth s =
  s |> String.fold_left (fun acc el -> acc + if el = '/' then 1 else 0) 0

let compdepth a b =
  let ad = countdepth a and bd = countdepth b in
  let dif = ad - bd in
  match dif with 0 -> 0 | _ -> dif / abs dif

let all_caps s = s |> String.map (fun c -> Char.uppercase_ascii c)
let all_low s = s |> String.map (fun c -> Char.lowercase_ascii c)
let fmt_dir d = String.split_on_char '/' d |> List.rev |> List.hd |> all_caps

let submeats s =
  s
  |> List.iter (fun d ->
         meat_print ("PROCESSING SUBMEAT " ^ fmt_dir d ^ "..");
         do_cmd ~args:false ("nix flake update --flake " ^ d))

let fresh () =
  print_string header;
  meat_print "HUNTING FRESH MEATS..";
  let argv_len = Array.length argv in
  let root = Sys.getenv "NH_FLAKE" in
  let base_dir = root ^ "/flakes" in
  (match (sub_flag (), all_flag ()) with
  | false, false ->
      (* no flags, plain flake inputs *)
      (if argv_len > 3 then
        let flakes = Array.sub argv 2 (argv_len - 2) in
        flakes
        |> Array.iter (fun f ->
               meat_print ("PROCESSING FRESH MEAT " ^ all_caps f ^ "..");
               do_cmd ~args:false
                 ("nix flake update " ^ all_low f ^ " --flake $NH_FLAKE"));
      else
        do_cmd ("nix flake update --flake $NH_FLAKE"))
  | _, true ->
      (* all flag, update all subflakes and main flake inputs *)
      walk base_dir |> List.sort compdepth |> List.rev |> submeats;
      print_string footer;
      meat_print "PROCESSING FRESH MEATS..";
      do_cmd ~args:false ("nix flake update --flake " ^ root)
  | true, _ when argv_len >= 4 ->
      (* sub flag, update inputs as subflakes *)
      let subflakes =
        (match argv_len with
        | 4 -> [ argv.(3) ]
        | _ -> Array.sub argv 4 (argv_len - 4) |> Array.to_list)
      in
      let sf_full = subflakes |> derelativise base_dir in
      sf_full |> List.map walk |> List.flatten |> List.sort compdepth
      |> List.rev |> List.append sf_full |> submeats;
      print_string footer;
      meat_print "PROCESSING FRESH MEATS..";
      subflakes
      |> List.iter (fun f ->
             do_cmd ~args:false
               ("nix flake update " ^ all_low f ^ " --flake $NH_FLAKE"))
  | _ -> print_string help);
  print_string footer;
  print_newline ()

let help () = print_string (header ^ help ^ footer ^ "\n")