#!/usr/bin/env python3
"""
Generate small random C programs and compile them with the W65816
backend.  Catches crashes / lowering gaps / verifier failures.

Each generated program is small (~10-50 lines), uses combinations of
features the compiler should handle:
  - integer arithmetic (i8, i16, i32, i64)
  - control flow (if, while, for, switch)
  - structs and pointer derefs
  - function calls (recursive, multi-arg)
  - casts and bit operations
  - arrays (small)

For each program, we just compile to .o.  If clang exits non-zero or
crashes, we save the offending source for inspection.

Optionally MAME-runs each program for additional runtime checks (off
by default — slow).

Usage: fuzzCompile.py [-n COUNT] [-s SEED] [--keep-failures DIR]
"""

import argparse, os, random, subprocess, sys, tempfile, hashlib
from pathlib import Path

CLANG = Path(__file__).parent.parent / "tools/llvm-mos-build/bin/clang"

# --- generators ---

def gen_expr(rng, depth=0):
    """Generate a random arithmetic expression returning int."""
    if depth > 3 or rng.random() < 0.3:
        return rng.choice([
            str(rng.randint(0, 100)),
            f"({rng.randint(0, 5)} + {rng.randint(0, 5)})",
            "x",
        ])
    op = rng.choice(["+", "-", "*", "&", "|", "^", "<<", ">>"])
    lhs = gen_expr(rng, depth + 1)
    rhs = rng.choice(["1", "2", "3", "4", str(rng.randint(0, 10))])
    if op in ("<<", ">>"):
        rhs = str(rng.randint(0, 7))
    return f"({lhs} {op} {rhs})"


def gen_stmt(rng, varCount, depth=0):
    """Generate a random statement."""
    kind = rng.choice(["assign", "if", "while", "loop"])
    if depth > 2:
        kind = "assign"
    if kind == "assign":
        v = f"v{rng.randint(0, varCount - 1)}"
        return f"{v} = {gen_expr(rng)};"
    if kind == "if":
        cond = f"{gen_expr(rng)} {rng.choice(['<', '>', '==', '!='])} {rng.randint(0, 30)}"
        body = gen_stmt(rng, varCount, depth + 1)
        return f"if ({cond}) {{ {body} }}"
    if kind == "while":
        cnt = rng.randint(2, 5)
        body = gen_stmt(rng, varCount, depth + 1)
        return f"{{ int j = {cnt}; while (j-- > 0) {{ {body} }} }}"
    if kind == "loop":
        v = f"v{rng.randint(0, varCount - 1)}"
        return f"for (int i = 0; i < {rng.randint(2, 6)}; i++) {{ {v} += i; }}"
    return ";"


def gen_function(rng, name, varCount):
    """Generate a function `int name(int x)` with random body."""
    decls = "\n  ".join(f"int v{i} = {rng.randint(0, 50)};" for i in range(varCount))
    stmts = "\n  ".join(gen_stmt(rng, varCount) for _ in range(rng.randint(3, 8)))
    ret = "v0"
    if varCount > 1:
        ret = " + ".join(f"v{i}" for i in range(min(varCount, 3)))
    return f"""int {name}(int x) {{
  {decls}
  {stmts}
  return {ret};
}}"""


def gen_program(rng):
    funcCount = rng.randint(1, 3)
    parts = []
    for i in range(funcCount):
        varCount = rng.randint(1, 5)
        parts.append(gen_function(rng, f"f{i}", varCount))
    parts.append(f"int call_all(int x) {{ return " +
                 " + ".join(f"f{i}(x)" for i in range(funcCount)) + "; }")
    return "\n\n".join(parts) + "\n"


# --- driver ---

def compile_one(source, keepDir=None, idx=0):
    """Compile source bytes; return (ok, msg)."""
    with tempfile.NamedTemporaryFile(suffix=".c", delete=False, mode="w") as f:
        f.write(source); cFile = f.name
    oFile = cFile + ".o"
    try:
        r = subprocess.run(
            [str(CLANG), "-target", "w65816", "-O2",
             "-ffunction-sections", "-c", cFile, "-o", oFile],
            capture_output=True, timeout=60
        )
        if r.returncode != 0:
            if keepDir:
                tag = hashlib.sha256(source.encode()).hexdigest()[:8]
                kept = Path(keepDir) / f"fail_{idx:03d}_{tag}.c"
                kept.write_text(source)
                kept.with_suffix(".c.stderr").write_bytes(r.stderr)
            return False, r.stderr.decode("utf-8", errors="replace")
        return True, ""
    except subprocess.TimeoutExpired:
        return False, "timeout (60s)"
    finally:
        for p in (cFile, oFile):
            try: os.unlink(p)
            except FileNotFoundError: pass


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-n", "--count", type=int, default=20)
    ap.add_argument("-s", "--seed", type=int, default=42)
    ap.add_argument("--keep-failures", default=None,
                    help="directory to save sources of failing inputs")
    ap.add_argument("-q", "--quiet", action="store_true")
    args = ap.parse_args()

    if args.keep_failures:
        Path(args.keep_failures).mkdir(parents=True, exist_ok=True)

    rng = random.Random(args.seed)
    fails = 0
    for i in range(args.count):
        src = gen_program(rng)
        ok, msg = compile_one(src, args.keep_failures, i)
        if not ok:
            fails += 1
            if not args.quiet:
                print(f"[fuzz] FAIL #{i}: {msg.splitlines()[0] if msg else '?'}")
        elif not args.quiet:
            print(f"[fuzz] OK   #{i}")
    print(f"fuzz: {args.count - fails}/{args.count} passed ({fails} fails)")
    sys.exit(1 if fails else 0)


if __name__ == "__main__":
    main()