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

CC = "gcc"

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"{CC} {' '.join(map(shlex.quote, 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, 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 preprocess_pass(c, meta):
  with tempfile.TemporaryDirectory() as tmpdir:
    t_file = open(f"{tmpdir}/file.c", "wb")
    t_file.write(f"#line 1 \"{meta['c_path']}\"\n".encode("utf-8"))
    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"{CC} -I{dirname} -E {meta['argstr']} {t_file.name}",
                     shell=True, capture_output=True), c)

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):
  c = c.replace(b"# ", b"#line ")
  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"]))
    args = ' '.join(meta["dietc_args"])
    return check_result(
      subprocess.run(f"timeout 5s {dietc_dir}/dietc {t_file.name} {args}",
                     shell=True, stdout=subprocess.PIPE, stderr=sys.stderr), c)

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

def make_external_pass(command):
  def pass_(dietc, meta):
    with tempfile.TemporaryDirectory() as tmpdir:
      t_file = open(f"{tmpdir}/file.c", "wb")
      t_file.write(dietc)
      t_file.flush()
      return check_result(subprocess.run(f"{command} {t_file.name}",
                          shell=True, stdout=subprocess.PIPE, stderr=sys.stderr), dietc)
  return pass_

def strip_line_info(dietc, meta):
  if "--line-numbers" not in meta["dietc_args"]: return dietc
  lines = dietc.split(b"\n")
  insert_after = dict()
  insert_before = dict()
  for i, line in enumerate(lines):
    if not line.startswith(b"#line"): continue
    before_line, after_line = lines[i-1], lines[i+1]
    insert_after[before_line] = line
    insert_before[after_line] = line
  meta["lineinfo_insert_after"] = insert_after
  meta["lineinfo_insert_before"] = insert_before
  return b"\n".join(l for l in lines if not l.startswith(b"#line"))

def reinsert_line_info(dietc, meta):
  if "--line-numbers" not in meta["dietc_args"]: return dietc
  lines = dietc.split(b"\n")
  insert_after = meta["lineinfo_insert_after"]
  insert_before = meta["lineinfo_insert_before"]
  final_lines = []
  for line in lines:
    if final_lines and final_lines[-1].startswith(b"#line"):
      final_lines.append(line)
    elif line in insert_before:
      final_lines.append(insert_before[line])
      final_lines.append(line)
    elif line in insert_after:
      final_lines.append(line)
      final_lines.append(insert_after[line])
    else:
      final_lines.append(line)
  return b"\n".join(final_lines)

PASSES = [preprocess_pass,
          strip_after_preprocess_pass,
          dietc_pass,
          strip_line_info,
          final_cleanup_pass,
          reinsert_line_info]

def main():
  args = list(map(shlex.quote, sys.argv[1:]))
  args.insert(0, f"-I{dietc_dir}/scripts/stdincludes")
  # first, parse out any DietC passes
  while "--dietc-pass" in args:
    i = args.index("--dietc-pass")
    args.pop(i)
    PASSES.insert(-1, make_external_pass(args.pop(i)))
  dietc_args = []
  while "--dietc-arg" in args:
    i = args.index("--dietc-arg")
    args.pop(i)
    dietc_args.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)

    pass_history = []
    last = open(c_file, "rb").read()
    meta = {"c_path": c_file, "argstr": argstr, "dietc_args": dietc_args}
    for dietpass in PASSES:
      last = dietpass(last, meta)
      pass_history.append(last)

    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"{CC} {' '.join(args)} -Wno-int-to-pointer-cast -Wno-builtin-declaration-mismatch", shell=True)
  check_result(gcc, pass_history[-1] if c_files else b"")

main()