let (>>=) = Result.bind;;

type stage =
  | Lex
  | Parse
  | Validate
  | Tacky
  | Codegen
  | Emit
;;

(* preprocess a .c to .i *)
let preprocess src =
  let base = Filename.chop_extension src in
  let preprocessed = base ^ ".i" in
  let cmd = Filename.quote_command "gcc" ["-E"; "-P"; src; "-o"; preprocessed] in
  let ret = Sys.command cmd in
  if ret == 0 then Ok preprocessed else Error ret
;;

(* compile a .i to a .s *)
let compile stage src =
  let base = Filename.chop_extension src in
  let asm = base ^ ".s" in
  let args = [src; "-o"; asm] in
  let args = match stage with
    | Lex -> "--lex" :: args
    | Parse -> "--parse" :: args
    | Validate -> "--validate" :: args
    | Tacky -> "--tacky" :: args
    | Codegen -> "--codegen" :: args
    | Emit -> args
  in
  (* ~/sources/compiler/lc/_build/install/default/bin/lc *)
  let cmd = Filename.quote_command "/home/t/sources/compiler/lc/_build/install/default/bin/lc" args in
  let ret = Sys.command cmd in
  if ret == 0 && stage == Emit then Ok asm else Error ret
;;

(* assemble and link the assembly file using "gcc ASSEMBLY_FILE -o OUTPUT_FILE" *)
let assemble asm do_object =
  let out = Filename.chop_extension asm in
  let cmd = match do_object with
    | true -> Filename.quote_command "gcc" ["-c"; asm; "-o"; out ^ ".o"]
    | false -> Filename.quote_command "gcc" [asm; "-o"; out]
  in
  let ret = Sys.command cmd in
  if ret == 0 then Ok out else Error ret
;;


let () =
  let usage = "compile [--lex|--parse|--codegen] [-c] <file>" in
  let stage = ref Emit in
  let input_file = ref None in
  let do_object = ref false in
  let spec = [
    ("--lex", Arg.Unit (fun () -> stage := Lex), "run the lexer, but stop before parsing");
    ("--parse", Arg.Unit (fun () -> stage := Parse), "run the lexer and parser, but stop before assembly generation");
    ("--validate", Arg.Unit (fun () -> stage := Validate), "run through the semantic analysis stage, but stop before TACKY generation");
    ("--tacky", Arg.Unit (fun () -> stage := Tacky), "run the TACKY generation stage, but stop before assembly generation");
    ("--codegen", Arg.Unit (fun () -> stage := Codegen), "perform lexing, parsing, and assembly generation, but stop before code emission");
    ("-c", Arg.Set do_object, "generate an object file instead of executable")
    ] in

  let anon_fun filename =
    if Option.is_some !input_file
    then Arg.Bad "invalid extra argument" |> raise
    else input_file := Some filename in
  Arg.parse spec anon_fun usage;
  match !input_file with
  | None -> Arg.usage spec usage
  | Some input_file ->
     let r = preprocess input_file >>= fun preprocessed ->
             compile !stage preprocessed >>= fun asm ->
             assemble asm !do_object
     in
     exit @@ match r with
             | Result.Ok _out -> 0
             | Result.Error err -> err