open Parser

(* Nushell built-in commands and keywords *)
let nushell_builtins = [
  "alias"; "all"; "ansi"; "any"; "append"; "ast"; "attr";
  "bits"; "break"; "bytes";
  "cal"; "cd"; "char"; "chunk-by"; "chunks"; "clear"; "collect";
  "columns"; "commandline"; "compact"; "complete"; "config"; "const";
  "continue"; "cp";
  "date"; "debug"; "decode"; "def"; "default"; "describe"; "detect";
  "do"; "drop"; "du";
  "each"; "echo"; "encode"; "enumerate"; "error"; "every"; "exec";
  "exit"; "explain"; "explore"; "export"; "export-env"; "extern";
  "fill"; "filter"; "find"; "first"; "flatten"; "for"; "format"; "from";
  "generate"; "get"; "glob"; "grid"; "group-by";
  "hash"; "headers"; "help"; "hide"; "hide-env"; "histogram";
  "history"; "http";
  "if"; "ignore"; "input"; "insert"; "inspect"; "interleave"; "into";
  "is-admin"; "is-empty"; "is-not-empty"; "is-terminal"; "items";
  "job"; "join";
  "keybindings"; "kill";
  "last"; "length"; "let"; "let-env"; "lines"; "load-env"; "loop"; "ls";
  "match"; "math"; "merge"; "metadata"; "mkdir"; "mktemp"; "module";
  "move"; "mut"; "mv";
  "nu-check"; "nu-highlight";
  "open"; "overlay";
  "panic"; "par-each"; "parse"; "path"; "plugin"; "port"; "prepend"; "print"; "ps";
  "query";
  "random"; "reduce"; "reject"; "rename"; "return"; "reverse"; "rm";
  "roll"; "rotate"; "run-external";
  "save"; "schema"; "scope"; "select"; "seq"; "shuffle"; "skip"; "sleep";
  "slice"; "sort"; "sort-by"; "source"; "source-env"; "split"; "start";
  "stor"; "str"; "sys";
  "table"; "take"; "tee"; "term"; "timeit"; "to"; "touch"; "transpose";
  "try"; "tutor";
  "ulimit"; "umask"; "uname"; "uniq"; "uniq-by"; "unlet"; "update";
  "upsert"; "url"; "use";
  "values"; "version"; "view";
  "watch"; "where"; "which"; "while"; "whoami"; "window"; "with-env"; "wrap";
  "zip";
]

let builtin_set = lazy (
  let tbl = Hashtbl.create (List.length nushell_builtins) in
  List.iter (fun s -> Hashtbl.replace tbl s true) nushell_builtins;
  tbl)

let is_nushell_builtin cmd =
  Hashtbl.mem (Lazy.force builtin_set) cmd

let dedup_entries entries =
  let key_of entry =
    match entry.switch with
    | Short c -> Printf.sprintf "-%c" c
    | Long l | Both (_, l) -> Printf.sprintf "--%s" l
  in
  let score entry =
    let sw = match entry.switch with Both _ -> 10 | _ -> 0 in
    let p = match entry.param with Some _ -> 5 | None -> 0 in
    let d = min 5 (String.length entry.desc / 10) in
    sw + p + d
  in
  let best = Hashtbl.create 64 in
  List.iter (fun e ->
    let k = key_of e in
    match Hashtbl.find_opt best k with
    | Some prev when score prev >= score e -> ()
    | _ -> Hashtbl.replace best k e
  ) entries;
  let covered_shorts = Hashtbl.create 16 in
  Hashtbl.iter (fun _ e ->
    match e.switch with
    | Both (c, _) -> Hashtbl.replace covered_shorts c true
    | _ -> ()
  ) best;
  let seen = Hashtbl.create 64 in
  List.filter_map (fun e ->
    let k = key_of e in
    if Hashtbl.mem seen k then None
    else
      match e.switch with
      | Short c when Hashtbl.mem covered_shorts c -> None
      | _ ->
        Hashtbl.add seen k true;
        Hashtbl.find_opt best k
  ) entries

let nushell_type_of_param = function
  | "FILE" | "file" | "PATH" | "path" | "DIR" | "dir" | "DIRECTORY"
  | "FILENAME" | "PATTERNFILE" -> "path"
  | "NUM" | "N" | "COUNT" | "NUMBER" | "int" | "INT" | "COLS" | "WIDTH"
  | "LINES" | "DEPTH" | "depth" -> "int"
  | _ -> "string"

let escape_nu s =
  if not (String.contains s '"') && not (String.contains s '\\') then s
  else begin
    let buf = Buffer.create (String.length s + 4) in
    String.iter (fun c -> match c with
      | '"' -> Buffer.add_string buf "\\\""
      | '\\' -> Buffer.add_string buf "\\\\"
      | _ -> Buffer.add_char buf c
    ) s;
    Buffer.contents buf
  end

let format_flag entry =
  let name = match entry.switch with
    | Both (s, l) -> Printf.sprintf "--%s(-%c)" l s
    | Long l -> Printf.sprintf "--%s" l
    | Short s -> Printf.sprintf "-%c" s
  in
  let typed = match entry.param with
    | Some (Mandatory p) | Some (Optional p) -> ": " ^ nushell_type_of_param p
    | None -> ""
  in
  let flag = "    " ^ name ^ typed in
  if String.length entry.desc = 0 then flag
  else
    let pad_len = max 1 (40 - String.length flag) in
    flag ^ String.make pad_len ' ' ^ "# " ^ entry.desc

let write_extern buf cmd_name result =
  let entries = dedup_entries result.entries in
  Printf.bprintf buf "export extern \"%s\" [\n" (escape_nu cmd_name);
  List.iter (fun e ->
    Buffer.add_string buf (format_flag e); Buffer.add_char buf '\n'
  ) entries;
  Buffer.add_string buf "]\n";
  List.iter (fun (sc : subcommand) ->
    Printf.bprintf buf "\nexport extern \"%s %s\" [  # %s\n]\n"
      (escape_nu cmd_name) (escape_nu sc.name) (escape_nu sc.desc)
  ) result.subcommands

let generate_extern cmd_name result =
  let buf = Buffer.create 1024 in
  write_extern buf cmd_name result;
  Buffer.contents buf

let module_name_of cmd_name =
  let s = String.map (function
    | ('a'..'z' | 'A'..'Z' | '0'..'9' | '-' | '_') as c -> c | _ -> '-') cmd_name in
  s ^ "-completions"

let generate_module cmd_name result =
  let m = module_name_of cmd_name in
  let buf = Buffer.create 1024 in
  Printf.bprintf buf "module %s {\n" m;
  write_extern buf cmd_name result;
  Printf.bprintf buf "}\n\nuse %s *\n" m;
  Buffer.contents buf

let generate_extern_from_entries cmd_name entries =
  generate_extern cmd_name { entries; subcommands = [] }