#!/bin/python3
import subprocess
import re
import tempfile
import shlex
import sys
import os

dietc_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))

def check_result(spr, debug):
  if not spr.returncode: return spr.stdout
  print("Failed to run gcc wrapper:", ' '.join(sys.argv), file=sys.stderr)
  if spr.stderr:
    print(spr.stderr.decode("utf-8"), file=sys.stderr)
  open("/tmp/dietc.error.c", "wb").write(debug)
  assert False

def main(argv):
  argv = argv[1:]
  def argument(flag): return argv[argv.index(flag) + 1]
  CC = "gcc"
  tmpdir = tempfile.TemporaryDirectory()

  # get the arguments WITHOUT -c, -o, source files, etc.
  stripped_args = []
  i = 0
  while i < len(argv):
    if (argv[i] == "-c" or argv[i].endswith(".c") or argv[i].endswith(".a") or argv[i].endswith(".o")):
      i += 1
    elif argv[i] == "-o":
      i += 2
    else:
      stripped_args.append(argv[i])
      i += 1

  # ====== HANDLE -c ======
  if "-c" in argv:
    c_file = next(a for a in argv if a.endswith(".c"))
    o_file = c_file[:-2] + ".o"
    if "-o" in argv: o_file = argument("-o")

    # (1) RUN PREPROCESSOR
    c = open(c_file, "r").read()
    t_file = open(f"{tmpdir.name}/file.c", "wb")
    t_file.write(f"#line 1 \"{c_file}\"\n".encode("utf-8"))
    t_file.write(f"#include \"{dietc_dir}/scripts/stdincludes/dietc_defines.h\"\n".encode("utf-8"))
    t_file.write(c.encode("utf-8"))
    t_file.flush()
    dirname = os.path.dirname(os.path.realpath(c_file))
    argstr = list(map(shlex.quote, stripped_args))
    argstr.insert(0, f"-I{dietc_dir}/scripts/stdincludes")
    argstr = ' '.join(argstr)
    preproc_c = check_result(
      subprocess.run(f"{CC} -I{dirname} -E {argstr} {t_file.name}",
                     shell=True, capture_output=True), c)

    # (2) STRIP GCC-SPECIFIC STUFF AFTER PREPROCESSING
    preproc_c = re.sub(rb"\[\[[^\]]*\]\]", b"", preproc_c)
    preproc_c = re.sub(rb"__asm__", b"", preproc_c)

    # (3) RUN DIETC
    preproc_c = preproc_c.replace(b"# ", b"#line ")
    t_file = open(f"{tmpdir.name}/file.c", "wb")
    t_file.write(preproc_c)
    t_file.flush()
    dirname = os.path.dirname(os.path.realpath(c_file))
    # TODO: args = ' '.join(meta["dietc_args"])
    args = "--type-builtins"
    diet_c = check_result(
      subprocess.run(f"timeout 5s {dietc_dir}/dietc {t_file.name} {args}",
                     shell=True, stdout=subprocess.PIPE, stderr=sys.stderr), preproc_c)

    # (4) SAVE TO A .o
    with open(o_file, "wb") as out:
      argstr = list(map(shlex.quote, stripped_args))
      argstr = ' '.join(argstr)
      out.write(b"DIETC")
      out.write(len(argstr).to_bytes(4, "little"))
      out.write(argstr.encode("utf-8"))
      out.write(diet_c)

    return

  # ======= If any .cs, first convert them to .os =======
  for i in range(len(argv)):
    if not argv[i].endswith(".c"): continue
    o_file = f"{tmpdir.name}/arg_{i}.o"
    main(["dietcc", *stripped_args, "-c", argv[i], "-o", o_file])
    argv[i] = o_file

  # ======= If any .as, first ar -x them into .os =======
  all_os = []
  for i in range(len(argv)):
    if argv[i].endswith(".o"):
      all_os.append(argv[i])
    elif argv[i].endswith(".a"):
      subprocess.run(["ar", "-x", argv[i], "--output", f"{tmpdir.name}/{i}.unar"])
      for path in os.listdir(f"{tmpdir.name}/{i}.unar"):
        all_os.append(f"{tmpdir.name}/{i}.unar/{path}")

  # ======= Put all the .os together into an output folder with Makefile =======
  out_path = "a.out"
  if "-o" in argv: out_path = argument("-o")
  assert out_path not in ("", "/")

  subprocess.run(["rm", "-rf", out_path])
  os.mkdir(out_path)
  makefile = open(f"{out_path}/Makefile", "w")
  makefile.write("a.out: " + ' '.join([f"{i}.o" for i in range(len(all_os))]) + " $(OTHER_OBJS)\n")
  argstr = list(map(shlex.quote, stripped_args))
  argstr = ' '.join(argstr)
  makefile.write(f"\t$(CC) $(CFLAGS) $^ -o $@ {argstr}\n\n")
  for i, o_file in enumerate(all_os):
    f = open(o_file, "rb")
    if f.read(5) == b"DIETC":
      # dietc file!
      argstr = f.read(int.from_bytes(f.read(4), "little")).decode("utf-8")
      with open(f"{out_path}/{i}.c", "wb") as i_file:
        i_file.write(f.read())
      makefile.write(f"{i}.o: {i}.c\n")
      makefile.write(f"\t$(CC) $(CFLAGS) -c $^ -o $@ {argstr} -Wno-builtin-declaration-mismatch\n\n")
      # TODO: include a postprocessing pass ...
    else:
      # regular .o
      shutil.copy(o_file, f"{out_path}/{i}.o")
  makefile.close()

main(sys.argv.copy())