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

if sys.argv[1:] == ["--version"]:
  print("DIETC")
  sys.exit(0)

# Have to do this because ./configure seems to do some nasty stuff
# https://lists.gnu.org/archive/html/bug-make/2022-11/msg00041.html
if "-E" in sys.argv[1:]:
  os.system(f"gcc {' '.join(sys.argv[1:])}")
  sys.exit(0)

# https://stackoverflow.com/questions/595305
dietc_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))

def check_result(spr):
  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)
  assert False

def preprocess_pass(c, meta):
  with tempfile.TemporaryDirectory() as tmpdir:
    t_file = open(f"{tmpdir}/file.c", "wb")
    t_file.write(f"#include \"{dietc_dir}/scripts/stdincludes/dietc_defines.h\"\n".encode("utf-8"))
    t_file.write(c)
    t_file.flush()
    dirname = os.path.dirname(os.path.realpath(meta["c_path"]))
    return check_result(
      subprocess.run(f"gcc -I{dirname} -E {meta['argstr']} {t_file.name}",
                     shell=True, capture_output=True))

def strip_after_preprocess_pass(c, meta):
  c = re.sub(rb"\[\[[^\]]*\]\]", b"", c)
  c = re.sub(rb"__asm__", b"", c)
  return c

def dietc_pass(c, meta):
  with tempfile.TemporaryDirectory() as tmpdir:
    t_file = open(f"{tmpdir}/file.c", "wb")
    t_file.write(c)
    t_file.flush()
    dirname = os.path.dirname(os.path.realpath(meta["c_path"]))
    return check_result(
      subprocess.run(f"timeout 5s {dietc_dir}/dietc {t_file.name}",
                     shell=True, capture_output=True))

def final_cleanup_pass(dietc, meta):
  # (1) treat builtins as builtins
  patterns = {
      rb"extern Type_[0-9]* __builtin_[^ ]* ;\n": b"",
      rb"extern Type_[0-9]* __(div|mul)[ds]c3 ;\n": b"",
      rb"extern Type_[0-9]* __btowc_alias ;\n": b"",
  }
  for pattern, replace in patterns.items():
    dietc = re.sub(pattern, b"", dietc)

  # (2) builtins must be called directly
  # NOTE: if passes mess this up, e.g., by doing optimizations, we could be in
  # trouble here
  pattern = rb"\t(t[0-9]*) = (__builtin_[^ ]*|alloca) ;\n"
  match_ = re.search(pattern, dietc)
  while match_ is not None:
    temp = match_.groups()[0]
    builtin = match_.groups()[1]
    # (2a) delete the match
    dietc = dietc.replace(b"\t" + temp + b" = " + builtin + b" ;\n", b"")
    # (2b) delete the declaration of the builtin
    dietc = re.sub(b"\tType_[0-9]* " + temp + b" ;\n", b"", dietc)
    # (2c) replace uses of the temp with the builtin
    dietc = dietc.replace(b" " + temp + b" ", b" " + builtin + b" ")
    dietc = dietc.replace(b"\t" + temp + b" ", b"\t" + builtin + b" ")

    match_ = re.search(pattern, dietc)

  # (3) __builtin_va_arg_pack* must be called *super* directly, and must be in
  # an always-inline method
  # NOTE: if passes mess this up, e.g., by doing optimizations, we could be in
  # trouble here
  pattern = rb"\t(t[0-9]*) = (__builtin_va_arg_pack[^ ]*) \( \) ;\n"
  match_ = re.search(pattern, dietc)
  any_arg_pack = (match_ is not None)
  while match_ is not None:
    temp = match_.groups()[0]
    builtin = match_.groups()[1]
    # (2a) delete the match
    dietc = dietc.replace(b"\t" + temp + b" = " + builtin + b" ( ) ;\n", b"")
    # (2b) delete the declaration of the result
    dietc = re.sub(b"\tType_[0-9]* " + temp + b" ;\n", b"", dietc)
    # (2c) replace uses of the temp with the builtin
    dietc = dietc.replace(b" " + temp + b" ", b" " + builtin + b" ( ) ")

    match_ = re.search(pattern, dietc)

  # (3b) any functions using __builtin_va_arg_pack* must be inlined
  if any_arg_pack or b"\t__builtin_va_start " in dietc:
    lines = dietc.split(b"\n")
    inline_is = set()
    last_decl_i = 0
    for i, line in enumerate(lines):
      if not line.startswith(b"\t"):
        last_decl_i = i
      elif b"__builtin_va_arg_pack" in line:
        inline_is.add(last_decl_i)
      elif b"\t__builtin_va_start " in line:
        # second argument of __builtin_va_start must be last named argument to
        # function
        fn, _po, arg1, _c, arg2, _pc, _sc = line.split()
        # {, ), ..., <,>, [arg]
        real_arg2 = lines[last_decl_i].split()[-5]
        lines[i] = b"\t" + b" ".join([fn, _po, arg1, _c, real_arg2, _pc, _sc])
    for i in sorted(inline_is):
      if "static" in lines[i]:
        lines[i] = b"inline __attribute__ ((__gnu_inline__)) " + lines[i]
      else:
        lines[i] = b"extern inline __attribute__ ((__gnu_inline__)) " + lines[i]
    dietc = b"\n".join(lines)

  return dietc

PASSES = [preprocess_pass,
          strip_after_preprocess_pass,
          dietc_pass,
          final_cleanup_pass]

def main():
  args = sys.argv[1:]
  args.insert(0, f"-I{dietc_dir}/scripts/stdincludes")
  # first, parse out any DietC passes
  dietc_passes = []
  while "--dietc-pass" in args:
    i = args.index("--dietc-pass")
    args.pop(i)
    dietc_passes.append(args.pop(i))

  # then process all of the C files
  out_dir = tempfile.TemporaryDirectory()
  c_files = [arg for arg in args if arg.endswith(".c")]
  diet_files = []
  for i, c_file in enumerate(c_files):
    subargs = [arg for arg in args if not arg.endswith(".c")]
    if "-o" in subargs:
      i = subargs.index("-o")
      subargs.pop(i)
      subargs.pop(i)
    argstr = " ".join(subargs)

    last = open(c_file, "rb").read()
    for dietpass in PASSES:
      last = dietpass(last, {"c_path": c_file, "argstr": argstr})

    diet_files.append(f"{out_dir.name}/file{i}.c")
    with open(diet_files[-1], "wb") as f:
      f.write(last)

  # then reconstruct the arguments, using the processed C files
  renaming = dict({c: diet for c, diet in zip(c_files, diet_files)})
  args = [renaming.get(arg, arg) for arg in args]
  if c_files and "-o" not in args:
    if "-c" in args:
      args += ["-o", c_files[0].replace(".c", ".o")]
  gcc = subprocess.run(f"gcc {' '.join(args)} -Wno-int-to-pointer-cast -Wno-builtin-declaration-mismatch", shell=True)
  check_result(gcc)

main()