open Inshellah.Parser
open Inshellah.Manpage
open Inshellah.Nushell

let usage () =
  Printf.eprintf
    {|inshellah - generate nushell completions

Usage:
  inshellah generate BINDIR MANDIR -o OUTDIR
      Full generation: native completions, manpages, and --help fallback.
      One .nu file per command.
  inshellah manpage FILE            Parse a manpage and emit nushell extern
  inshellah manpage-dir DIR         Batch-process manpages under DIR
  inshellah help CMD [ARGS...]      Run CMD ARGS --help, parse and emit extern
  inshellah parse-help CMD          Read --help text from stdin, emit extern
  inshellah demo                    Run built-in demo

|};
  exit 1

let command_sections = [1; 8]

let contains_str s sub =
  try ignore (Str.search_forward (Str.regexp_string sub) s 0); true
  with Not_found -> false

let is_nushell_source text =
  String.length text > 20
  && (contains_str text "export extern"
      || contains_str text "export def"
      || (contains_str text "module " && contains_str text "export"))

let filename_of_cmd cmd =
  String.map (function
    | ('a'..'z' | 'A'..'Z' | '0'..'9' | '-' | '_' | '.') as c -> c | _ -> '-') cmd

let write_file path contents =
  let oc = open_out path in output_string oc contents; close_out oc

let cmd_name_of_manpage path =
  let base = Filename.basename path in
  let base =
    if Filename.check_suffix base ".gz" then Filename.chop_suffix base ".gz"
    else base in
  try Filename.chop_extension base with Invalid_argument _ -> base

let safe_env = lazy (
  Array.of_list (
    List.filter (fun s ->
      not (String.starts_with ~prefix:"DISPLAY=" s
           || String.starts_with ~prefix:"WAYLAND_DISPLAY=" s
           || String.starts_with ~prefix:"DBUS_SESSION_BUS_ADDRESS=" s
           || String.starts_with ~prefix:"XAUTHORITY=" s))
      (Array.to_list (Unix.environment ()))))

let run_cmd args timeout_ms =
  let (rd, wr) = Unix.pipe () in
  let devnull = Unix.openfile "/dev/null" [Unix.O_RDONLY] 0 in
  let argv = Array.of_list args in
  let pid =
    try Unix.create_process_env (List.hd args) argv
          (Lazy.force safe_env) devnull wr wr
    with Unix.Unix_error _ ->
      Unix.close rd; Unix.close wr; Unix.close devnull; -1 in
  Unix.close wr; Unix.close devnull;
  if pid < 0 then (Unix.close rd; None)
  else begin
    let buf = Buffer.create 4096 in
    let deadline = Unix.gettimeofday () +. (float_of_int timeout_ms /. 1000.0) in
    let chunk = Bytes.create 8192 in
    let alive = ref true in
    (try while !alive do
       let remaining = deadline -. Unix.gettimeofday () in
       if remaining <= 0.0 then alive := false
       else match Unix.select [rd] [] [] (min remaining 0.05) with
         | (_ :: _, _, _) ->
           let n = Unix.read rd chunk 0 8192 in
           if n = 0 then raise Exit
           else Buffer.add_subbytes buf chunk 0 n
         | _ -> ()
     done with Exit -> ());
    Unix.close rd;
    if not !alive then begin
      (try Unix.kill pid Sys.sigkill with Unix.Unix_error _ -> ());
      ignore (Unix.waitpid [] pid)
    end else
      ignore (Unix.waitpid [] pid);
    if Buffer.length buf > 0 then Some (Buffer.contents buf) else None
  end

let is_executable path =
  try let st = Unix.stat path in
    st.st_kind = Unix.S_REG && st.st_perm land 0o111 <> 0
  with Unix.Unix_error _ -> false

let is_shell_script path =
  try
    let real = Unix.realpath path in
    let ic = open_in_bin real in
    let line = (try let b = Bytes.create 256 in
      let n = input ic b 0 256 in
      let s = Bytes.sub_string b 0 n in
      (match String.index_opt s '\n' with Some i -> String.sub s 0 i | None -> s)
    with End_of_file -> "") in
    close_in ic;
    String.length line >= 2 && line.[0] = '#' && line.[1] = '!'
    && let shebang = String.lowercase_ascii line in
       List.exists (fun s -> contains_str shebang s)
         ["bash"; "/sh "; "/sh\n"; "zsh"; "fish"; "nushell"; "/nu "; "/nu\n";
          "dash"; "ksh"; "csh"]
  with _ -> false

let elf_scan path needles =
  let found = Hashtbl.create 4 in
  let remaining () = List.filter (fun n -> not (Hashtbl.mem found n)) needles in
  (try
    let real = Unix.realpath path in
    let ic = open_in_bin real in
    let magic = Bytes.create 4 in
    really_input ic magic 0 4;
    if Bytes.get magic 0 = '\x7f' && Bytes.get magic 1 = 'E'
       && Bytes.get magic 2 = 'L' && Bytes.get magic 3 = 'F' then begin
      let max_needle = List.fold_left (fun m n -> max m (String.length n)) 0 needles in
      let chunk_size = 65536 in
      let buf = Bytes.create (chunk_size + max_needle) in
      let carry = ref 0 in
      let eof = ref false in
      while not !eof && remaining () <> [] do
        let n = (try input ic buf !carry chunk_size with End_of_file -> 0) in
        if n = 0 then eof := true
        else begin
          let total = !carry + n in
          List.iter (fun needle ->
            if not (Hashtbl.mem found needle) then begin
              let nlen = String.length needle in
              let i = ref 0 in
              while !i <= total - nlen do
                if Bytes.get buf !i = needle.[0] then begin
                  let ok = ref true in
                  for j = 1 to nlen - 1 do
                    if Bytes.get buf (!i + j) <> needle.[j] then ok := false
                  done;
                  if !ok then (Hashtbl.replace found needle true; i := total)
                  else incr i
                end else incr i
              done
            end
          ) (remaining ());
          let new_carry = min max_needle total in
          Bytes.blit buf (total - new_carry) buf 0 new_carry;
          carry := new_carry
        end
      done
    end;
    close_in ic
  with _ ->
    List.iter (fun n -> Hashtbl.replace found n true) needles);
  found

let skip_name name =
  String.length name = 0 || name.[0] = '.'
  || String.starts_with ~prefix:"lib" name
  || String.ends_with ~suffix:"-daemon" name
  || String.ends_with ~suffix:"-wrapped" name
  || String.ends_with ~suffix:".so" name

type bin_class = Skip | Try_help | Try_native_and_help

let classify_binary bindir name =
  if is_nushell_builtin name || skip_name name then Skip
  else
    let path = Filename.concat bindir name in
    if not (is_executable path) then Skip
    else if is_shell_script path then Try_help
    else
      let scan = elf_scan path ["-h"; "completion"] in
      if not (Hashtbl.mem scan "-h") then Skip
      else if Hashtbl.mem scan "completion" then Try_native_and_help
      else Try_help

let num_cores () =
  try
    let ic = open_in "/proc/cpuinfo" in
    let n = ref 0 in
    (try while true do
       if String.starts_with ~prefix:"processor" (input_line ic) then incr n
     done with End_of_file -> ());
    close_in ic; max 1 !n
  with _ -> 4

let try_native_completion bin_path =
  let patterns = [
    [bin_path; "completions"; "nushell"];
    [bin_path; "completion"; "nushell"];
    [bin_path; "--completions"; "nushell"];
    [bin_path; "--completion"; "nushell"];
    [bin_path; "generate-completion"; "nushell"];
    [bin_path; "--generate-completion"; "nushell"];
    [bin_path; "shell-completions"; "nushell"];
  ] in
  let rec go = function
    | [] -> None
    | args :: rest ->
      match run_cmd args 500 with
      | Some text when is_nushell_source text -> Some text
      | _ -> go rest
  in
  go patterns

let cmd_manpage file =
  let contents = read_manpage_file file in
  let fallback = cmd_name_of_manpage file in
  let cmd = match extract_synopsis_command contents with
    | Some name -> name | None -> fallback in
  if not (is_nushell_builtin cmd) then
    let result = parse_manpage_string contents in
    if result.entries <> [] then
      print_string (generate_extern cmd result)

let cmd_manpage_dir dir =
  List.iter (fun section ->
    let subdir = Filename.concat dir (Printf.sprintf "man%d" section) in
    if Sys.file_exists subdir && Sys.is_directory subdir then
      Array.iter (fun file ->
        (try cmd_manpage (Filename.concat subdir file) with _ -> ())
      ) (Sys.readdir subdir)
  ) command_sections

let cmd_help args =
  match args with
  | [] -> Printf.eprintf "error: help requires a command name\n"; exit 1
  | cmd :: rest ->
    let name = Filename.basename cmd in
    (match run_cmd (cmd :: rest @ ["--help"]) 10_000 with
     | None -> Printf.eprintf "no output from %s --help\n" name; exit 1
     | Some text ->
       (match parse_help text with
        | Ok r -> print_string (generate_extern name r)
        | Error msg -> Printf.eprintf "parse error for %s: %s\n" name msg; exit 1))

let cmd_parse_help cmd =
  let buf = Buffer.create 4096 in
  (try while true do
     Buffer.add_string buf (input_line stdin); Buffer.add_char buf '\n'
   done with End_of_file -> ());
  (match parse_help (Buffer.contents buf) with
   | Ok r -> print_string (generate_extern cmd r)
   | Error msg -> Printf.eprintf "parse error for %s: %s\n" cmd msg; exit 1)

let process_manpage file =
  try
    let contents = read_manpage_file file in
    let fallback = cmd_name_of_manpage file in
    let cmd = match extract_synopsis_command contents with
      | Some name -> name | None -> fallback in
    if is_nushell_builtin cmd then None
    else
      let result = parse_manpage_string contents in
      if result.entries <> [] then Some (cmd, result) else None
  with _ -> None

let manpaged_commands mandir =
  let cmds = Hashtbl.create 128 in
  List.iter (fun section ->
    let subdir = Filename.concat mandir (Printf.sprintf "man%d" section) in
    if Sys.file_exists subdir && Sys.is_directory subdir then
      Array.iter (fun f -> Hashtbl.replace cmds (cmd_name_of_manpage f) true)
        (Sys.readdir subdir)
  ) command_sections;
  cmds

let cmd_generate bindir mandir outdir =
  let done_cmds = Hashtbl.create 256 in
  let bins = Sys.readdir bindir in
  Array.sort String.compare bins;
  let manpaged = manpaged_commands mandir in
  let max_jobs = num_cores () in
  let classified = Array.map (fun name ->
    if Hashtbl.mem manpaged name then (name, Skip)
    else (name, classify_binary bindir name)
  ) bins in
  let pending = ref [] in
  let reap () =
    pending := List.filter (fun pid ->
      match Unix.waitpid [Unix.WNOHANG] pid with
      | (0, _) -> true | _ -> false
      | exception Unix.Unix_error (Unix.ECHILD, _, _) -> false
    ) !pending in
  let wait_slot () =
    while List.length !pending >= max_jobs do
      reap ();
      if List.length !pending >= max_jobs then
        (try ignore (Unix.wait ()) with Unix.Unix_error _ -> ());
      reap ()
    done in
  Array.iter (fun (name, cls) ->
    match cls with
    | Skip -> ()
    | Try_help | Try_native_and_help ->
      wait_slot ();
      let pid = Unix.fork () in
      if pid = 0 then begin
        (try
          let path = Filename.concat bindir name in
          let native_ok = match cls with
            | Try_native_and_help ->
              (match try_native_completion path with
               | Some src ->
                 write_file (Filename.concat outdir (filename_of_cmd name ^ ".nu")) src;
                 true
               | None -> false)
            | _ -> false in
          if not native_ok then begin
            let text = match run_cmd [path; "--help"] 200 with
              | Some _ as r -> r
              | None -> run_cmd [path; "-h"] 200 in
            match text with
            | None -> ()
            | Some t ->
              (match parse_help t with
               | Ok r when r.entries <> [] ->
                 write_file (Filename.concat outdir (filename_of_cmd name ^ ".nu"))
                   (generate_module name r)
               | _ -> ())
          end;
          exit 0
        with _ -> exit 1)
      end else begin
        pending := pid :: !pending;
        Hashtbl.replace done_cmds name true
      end
  ) classified;
  while !pending <> [] do
    (try ignore (Unix.wait ()) with Unix.Unix_error _ -> ());
    reap ()
  done;
  List.iter (fun section ->
    let subdir = Filename.concat mandir (Printf.sprintf "man%d" section) in
    if Sys.file_exists subdir && Sys.is_directory subdir then begin
      let files = Sys.readdir subdir in
      Array.sort String.compare files;
      Array.iter (fun file ->
        match process_manpage (Filename.concat subdir file) with
        | None -> ()
        | Some (cmd, result) ->
          let base = List.hd (String.split_on_char ' ' cmd) in
          if Hashtbl.mem done_cmds base then ()
          else begin
            let outpath = Filename.concat outdir (filename_of_cmd base ^ ".nu") in
            if Sys.file_exists outpath then begin
              let existing =
                let ic = open_in outpath in
                let n = in_channel_length ic in
                let s = Bytes.create n in
                really_input ic s 0 n; close_in ic;
                Bytes.to_string s in
              let mod_name = module_name_of base in
              let use_line = Printf.sprintf "\nuse %s *\n" mod_name in
              let base_content =
                if contains_str existing use_line then
                  String.sub existing 0
                    (Str.search_forward (Str.regexp_string use_line) existing 0)
                else existing in
              write_file outpath
                (String.concat "" [base_content; generate_extern cmd result; use_line])
            end else
              write_file outpath (generate_module base result)
          end
      ) files
    end
  ) command_sections

let cmd_demo () =
  Printf.printf "# Generated by: inshellah demo\n\n";
  match parse_help
    {|Usage: ls [OPTION]... [FILE]...
  -a, --all                  do not ignore entries starting with .
  -A, --almost-all           do not list implied . and ..
      --block-size=SIZE      with -l, scale sizes by SIZE when printing
      --color[=WHEN]         color the output WHEN
  -h, --human-readable       with -l and -s, print sizes like 1K 234M 2G etc.
      --help                 display this help and exit
      --version              output version information and exit
|} with
  | Ok r -> print_string (generate_extern "ls" r)
  | Error msg -> Printf.eprintf "parse error: %s\n" msg

let () =
  match Array.to_list Sys.argv |> List.tl with
  | ["generate"; bindir; mandir; "-o"; outdir] -> cmd_generate bindir mandir outdir
  | ["manpage"; file] -> cmd_manpage file
  | ["manpage-dir"; dir] -> cmd_manpage_dir dir
  | "help" :: rest -> cmd_help rest
  | ["parse-help"; cmd] -> cmd_parse_help cmd
  | ["demo"] -> cmd_demo ()
  | _ -> usage ()