#!/bin/python3
import argparse;
import subprocess as proc;
import os;
import hashlib;
import pickle;
import tempfile;
import multiprocessing as mp;
import time;
from pathlib import Path;

mp.set_start_method("fork")

# This is used to build and test the compiler only.
# Large part of the scripts are generated by ChatGPT and Deepseek.

parser = argparse.ArgumentParser()
parser.add_argument("-g", "--gdb", action="store_true")
parser.add_argument("-V", "--valgrind", action="store_true")
parser.add_argument("-v", "--verbose", action="store_true")
parser.add_argument("--verify", action="store_true")
parser.add_argument("-s", "--stats", action="store_true")
parser.add_argument("--arm", action="store_true")
parser.add_argument("-n", "--no-execute", action="store_true")
parser.add_argument("-S", "--no-link", action="store_true")
parser.add_argument("-O1", action="store_true")
parser.add_argument("-r", "--run", action="store_true")
parser.add_argument("--sat", action="store_true")
parser.add_argument("--bv", action="store_true")
parser.add_argument("-d", "--directory", type=str)
parser.add_argument("--timeout", type=float, default=2)
parser.add_argument("--asm", type=str)
parser.add_argument("-t", "--test", type=str)
parser.add_argument("-i", "--input", type=str) # Input to the executable
parser.add_argument("-p", "--print-after", type=str)
parser.add_argument("-q", "--print-before", type=str)

args = parser.parse_args()

SRC_DIR = Path("src")
BUILD_DIR = Path("build")
FINAL_BINARY = BUILD_DIR / "sysc"
COMPILER = "clang++"
AR = "ar"
CFLAGS = [
  "-c", "-std=c++17", "-O2", "-g",
  "-Wall", "-Wextra", "-Wuninitialized", "-Wstrict-aliasing",
  "-Wno-sign-compare", "-Wno-parentheses"
]
LDFLAGS = []
CACHE_FILE = BUILD_DIR / ".build_cache.pkl"
INCLUDE_CACHE_FILE = BUILD_DIR / ".include_cache.pkl"
TIMES_FILE = "rank/time.txt"

def hash_file(path):
  h = hashlib.sha256()
  with open(path, 'rb') as f:
    h.update(f.read())
  return h.hexdigest()

def find_files() -> tuple[list[Path], list[Path]]:
  cpp_files = []
  h_files = []
  for path in SRC_DIR.rglob("*"):
    if path.suffix == ".cpp":
      cpp_files.append(path)
    elif path.suffix == ".h":
      h_files.append(path)
  return cpp_files, h_files

def load_cache():
  if CACHE_FILE.exists():
    with open(CACHE_FILE, "rb") as f:
      return pickle.load(f)
  return {}

def save_cache(cache):
  BUILD_DIR.mkdir(parents=True, exist_ok=True)
  with open(CACHE_FILE, "wb") as f:
    pickle.dump(cache, f)

def needs_recompile(src_path, obj_path, cache, dependencies):
  src_hash = hash_file(src_path)
  dep_hashes = { str(dep): hash_file(dep) for dep in dependencies }

  prev = cache.get(str(src_path))
  if not prev:
      return True
  if prev['src_hash'] != src_hash:
      return True
  if prev['dep_hashes'] != dep_hashes:
      return True
  if not obj_path.exists():
      return True
  return False

include_cache = {}
include_hashes = {}

def load_include_cache():
  if INCLUDE_CACHE_FILE.exists():
    with open(INCLUDE_CACHE_FILE, "rb") as f:
      return pickle.load(f)
  return {}

def save_include_cache(cache):
  BUILD_DIR.mkdir(parents=True, exist_ok=True)
  with open(INCLUDE_CACHE_FILE, "wb") as f:
    pickle.dump(cache, f)

def get_all_includes(src_path: Path, visited=None) -> set[Path]:
  if visited is None:
    visited = set()

  resolved_path = src_path.resolve()
  if resolved_path in visited:
    return set()

  visited.add(resolved_path)
  file_hash = hash_file(resolved_path)

  cached_entry = include_cache.get(resolved_path)
  if cached_entry and include_hashes.get(resolved_path) == file_hash:
    return cached_entry

  includes = set()

  with open(resolved_path, "r") as f:
    for line in f:
      line = line.strip()
      if line.startswith("#include \""):
        header = line.split("\"")[1]
        include_path = resolved_path.parent / header
        if include_path.exists():
          includes.add(include_path.resolve())
          includes.update(get_all_includes(include_path, visited))

  include_cache[resolved_path] = includes
  include_hashes[resolved_path] = file_hash
  return includes

# Note that the type of counter is not easily representable.
def compile_cpp(src_path: Path, obj_path: Path):
  obj_path.parent.mkdir(parents=True, exist_ok=True)
  proc.check_call([COMPILER] + CFLAGS + ["-o", str(obj_path), str(src_path)])

def archive_objects(obj_files, lib_path: Path):
  if lib_path.exists():
    lib_path.unlink()
  print(f"Creating archive {lib_path}")
  proc.check_call([AR, "rcs", str(lib_path)] + [str(obj) for obj in obj_files])

def link_libraries(lib_files, output_binary):
  print(f"Linking {output_binary}")
  proc.check_call([COMPILER] + LDFLAGS + ["-o", str(output_binary)] + ["-Wl,--start-group"] + [str(lib) for lib in lib_files] + ["-Wl,--end-group"])

def build():
  global include_cache, include_hashes
  include_cache_data = load_include_cache()
  include_cache = include_cache_data.get("cache", {})
  include_hashes = include_cache_data.get("hashes", {})

  cpp_files, _ = find_files()
  cache = load_cache()

  # Step 1: Compile .cpp to .o
  obj_files: list[tuple[Path, Path]] = []
  folder_changed: dict[Path, bool] = {}
  tasks: list[tuple[Path, Path]] = []
  for cpp in cpp_files:
    rel_dir = cpp.relative_to(SRC_DIR).parent
    obj_dir = BUILD_DIR / rel_dir
    obj_path = obj_dir / (cpp.stem + ".o")

    dependencies = get_all_includes(cpp)
    if needs_recompile(cpp, obj_path, cache, dependencies):
      tasks.append((cpp, obj_path))
      cache[str(cpp)] = {
        'src_hash': hash_file(cpp),
        'dep_hashes': {str(dep): hash_file(dep) for dep in dependencies},
      }
      folder_changed[rel_dir] = True
    obj_files.append((rel_dir, obj_path))

  total = len(tasks)
  file_prompt = "file" if total == 1 else "files"
  if total > 0:
    print(f"Compiling {total} {file_prompt}")
  with mp.Pool() as pool:
    pool.starmap(compile_cpp, tasks)
  
  save_cache(cache)

  # Step 2: Archive .o's in same folder into .a
  folder_objs = {}
  for rel_dir, obj in obj_files:
    folder_objs.setdefault(rel_dir, []).append(obj)

  lib_files = []
  for folder, objs in folder_objs.items():
    lib_path = BUILD_DIR / folder / (folder.name + ".a")
    need_archive = folder_changed.get(folder, False) or not lib_path.exists()
    if need_archive:
      archive_objects(objs, lib_path)
    lib_files.append(lib_path)

  # Step 3: Link all .a's into final binary
  link_libraries(lib_files, FINAL_BINARY)

  save_include_cache({
    "cache": include_cache,
    "hashes": include_hashes
  })


def run_asm(file: str):
  basename = os.path.splitext(os.path.basename(file))[0]

  # Invoke gcc.
  gcc = "aarch64-linux-gnu-gcc" if args.arm else "riscv64-linux-gnu-gcc"
  proc.run([gcc, file, "test/official/sylib.c", "-static", "-o", f"temp/{basename}"], check=True)

  # Run the file.
  qemu = "qemu-aarch64-static" if args.arm else "qemu-riscv64-static"
  buffer = b""
  if args.input:
    with open(args.input, "r") as f:
      buffer = f.read().encode('utf-8')
  return proc.run([qemu, f"temp/{basename}"], input=buffer if args.input else None)

def run(full_file: str, no_exec: bool):
  basename = os.path.splitext(os.path.basename(full_file))[0]

  command = [f"{BUILD_DIR}/sysc", full_file]

  if args.gdb:
    command = ["gdb", "--args", *command]
    no_exec = True
  
  if args.valgrind:
    command = ["valgrind", *command]
    no_exec = True
  
  if args.arm:
    command.append("--arm")
  
  if args.verbose:
    command.append("-v")
  
  if args.stats:
    command.append("--stats")

  if args.verify:
    command.append("--verify")

  if args.print_after:
    command.extend(["--print-after", args.print_after])

  if args.print_before:
    command.extend(["--print-before", args.print_before])

  if args.no_link:
    command.append("-S")

  if args.O1:
    command.append("-O1")

  if args.run:
    dir = Path(full_file).parent.absolute()
    input = dir / f"{basename}.in"
    command.extend(["--compare", f"{dir}/{basename}.out"])
    if input.exists():
      command.extend(["-i", str(input)])

  command.extend(["-o", f"temp/{basename}.s"])
  
  # Invoke SysY compiler.
  proc.run(command, check=True)
  print("Done.")

  if no_exec:
    return;
  return run_asm(f"temp/{basename}.s")

qemu = "qemu-aarch64-static" if args.arm else "qemu-riscv64-static"
gcc = "aarch64-linux-gnu-gcc" if args.arm else "riscv64-linux-gnu-gcc"

def run_test_case(sy_path: Path, in_path: Path, out_path: Path):
  test_name = f"{sy_path.parent.name}/{sy_path.name}"
  print(f"Testing {test_name}...", end='\r')
  
  with tempfile.TemporaryDirectory() as tmpdir:
    asm_path = Path(tmpdir) / "output.s"
    exe_path = Path(tmpdir) / "a.out"
    
    # Step 1: Compile .sy to .s using sysc
    commands = [f"{BUILD_DIR}/sysc", str(sy_path), "-o", str(asm_path)]
    if args.arm:
      commands.append("--arm")
    if args.verify:
      commands.append("--verify")
    
    try:
      proc.run(
        commands,
        check=True,
        stdout=proc.PIPE,
        stderr=proc.STDOUT,
        timeout=args.timeout
      )
    except proc.CalledProcessError as e:
      return (sy_path, f"Compile failed: {e.output.decode().strip()}")
    except proc.TimeoutExpired:
      return (sy_path, f"Compiler timeout ({args.timeout:.2f}s)")
      
    if not asm_path.exists():
      return (sy_path, "No assembly output generated")
    
    try:
      proc.run(
        [gcc, "-static", str(asm_path), "test/official/sylib.c", "-o", str(exe_path)],
        check=True,
        stdout=proc.PIPE,
        stderr=proc.STDOUT
      )
    except proc.CalledProcessError as e:
      return (sy_path, f"Linking failed: {e.output.decode().strip()}")
      
    if not exe_path.exists():
      return (sy_path, "No executable generated after linking")
      
    # Step 3: Run with QEMU
    input_data = None
    if in_path:
      with open(in_path, 'r') as f:
        input_data = f.read()

    try:
      start_time = time.perf_counter()
      result = proc.run(
        f"{qemu} {str(exe_path)}",
        input=None if not input_data else input_data.encode('utf-8'),
        stdout=proc.PIPE,
        stderr=proc.DEVNULL,
        timeout=args.timeout,
        shell=True
      )
      elapsed_time = time.perf_counter() - start_time
      # Record timing info
      with open(TIMES_FILE, "a") as f:
        f.write(f"{sy_path.name} {elapsed_time:.6f}\n")

      actual_out: str = result.stdout.decode('utf-8').strip()
      actual = f"{actual_out}\n{result.returncode}".strip()
    except proc.TimeoutExpired:
      return (sy_path, f"Timeout ({args.timeout:.2f}s)")
    except Exception as e:
      return (sy_path, f"Runtime error: {str(e)}")
    
    with open(out_path) as f:
      expected = f.read().strip()

    # Ignore leading/trailing spaces on each line.
    actual = '\n'.join([x.strip() for x in actual.split('\n')])
    expected = '\n'.join([x.strip() for x in expected.split('\n')])
      
    if actual != expected:
      return ((sy_path, f"Output mismatch:\n{actual}"))
    else:
      return None

def test_all():
  # Collect all test cases
  test_cases = []
  directory: str = args.directory
  for root, _, files in os.walk(directory):
    for file in files:
      if file.endswith('.sy'):
        sy_path: Path = Path(root) / file
        base = sy_path.stem
        in_path = sy_path.with_name(f"{base}.in")
        out_path = sy_path.with_name(f"{base}.out")
        
        if not out_path.exists():
          print(f"Warning: {sy_path} missing .out file, skipping")
          continue;
          
        test_cases.append((sy_path, in_path if in_path.exists() else None, out_path))
  
  total = len(test_cases)
  
  with mp.Pool() as pool:
    results = pool.starmap(run_test_case, test_cases)

  # In case there's still unfinished processes.
  proc.run(f"killall -9 {qemu}", shell=True)
  failures = [x for x in results if x is not None]
  failed = len(failures)
  passed = total - failed
  
  # Print results
  print("\nTest results:")
  print(f"Total:  {total}")
  print(f"Passed: {passed}")
  print(f"Failed: {failed}")
  
  if failed:
    print("\nFailed cases:")
    failures = sorted(failures, key=lambda x: x[0])
    for path, reason in failures:
      print(f"- {path}")
      lines = '\n'.join([x.strip() for x in reason.split('\n')[:10]])
      print(f"  Reason: {lines}")

def test_sat():
  files = os.listdir("test/cnf")
  for file in files:
    print(f"checking: {file}")
    out = proc.run(f"{BUILD_DIR}/sysc --sat < test/cnf/{file}", stdout=proc.PIPE, shell=True, text=True).stdout
    if file.startswith("uf") and not out.startswith("sat"):
      print(f"error: {file} out = '{out}'")
    elif file.startswith("uuf") and not out.startswith("unsat"):
      print(f"error: {file}")


if __name__ == "__main__":
  if args.asm:
    result = run_asm(args.asm)
    print("Return value: ", result.returncode)
    exit(0)

  build()

  if args.bv:
    commands = [f"{BUILD_DIR}/sysc", "--bv"]
    if args.stats:
      commands.append("-s")
    if args.verbose:
      commands.append("-v")
    proc.run(commands)
    exit(0)

  if args.sat:
    test_sat()

  if args.directory:
    test_all()

  if args.test:
    result = run(args.test, args.no_execute)
    if result is not None:
      print("Return value:", result.returncode)
