#!/usr/bin/env python3
"""Differential soundness fuzzer for VeriCoding (keystone/vericoding.check_contract).

Generates random functions in VeriCoding's supported subset (int/bool parameters, + - * // %, abs/min/max, if/else,
conditional expressions) and random pre/postconditions. Every verdict is checked by running the real Python on
every input in a grid:

  FALSE_PROVED          Z3 said "proved for all inputs" but Python breaks the postcondition on some input
  PROVED_BUT_CRASHES    Z3 said "proved" but Python raises (e.g. ZeroDivisionError) on an input meeting the precondition
  FALSE_COUNTEREXAMPLE  the counterexample does not break the postcondition when run in Python
  CRASH_COUNTEREXAMPLE  the counterexample makes Python raise, but the verdict calls it an ordinary violation
  WRONG_FAILURE_KIND    the verdict names a failure kind (violated / which part raises) that Python does not show
  ENGINE_ERROR          check_contract raised instead of answering

    python3 scripts/vericoding/fuzz_soundness.py [--n 3000] [--seed 1] [--json out.json]
"""
from __future__ import annotations

import argparse
import itertools
import json
import random
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from keystone import vericoding as V  # noqa: E402

GRID = range(-5, 6)


def gen_expr(r: random.Random, names, depth=0):
    if depth > 2 or r.random() < 0.3:
        return r.choice(names + [str(r.randint(-3, 3))])
    k = r.random()
    a, b = gen_expr(r, names, depth + 1), gen_expr(r, names, depth + 1)
    if k < 0.55:
        return f"({a} {r.choice(['+', '-', '*', '//', '%'])} {b})"
    if k < 0.75:
        return f"{r.choice(['abs'])}({a})" if r.random() < 0.4 else f"{r.choice(['min', 'max'])}({a}, {b})"
    c = gen_cond(r, names, depth + 1)
    return f"({a} if {c} else {b})"


def gen_cond(r: random.Random, names, depth=0):
    a, b = gen_expr(r, names, depth + 1), gen_expr(r, names, depth + 1)
    c = f"{a} {r.choice(['<', '<=', '>', '>=', '==', '!='])} {b}"
    if depth < 2 and r.random() < 0.25:
        c = f"({c}) {r.choice(['and', 'or'])} ({gen_cond(r, names, depth + 1)})"
    return c


def gen_case(r: random.Random):
    params = ["x", "y"][: r.choice([1, 2])]
    names = list(params)
    body = []
    if r.random() < 0.5:
        body.append(f"    t = {gen_expr(r, names)}")
        names = names + ["t"]
    if r.random() < 0.5:
        body.append(f"    if {gen_cond(r, names)}:\n        return {gen_expr(r, names)}\n    else:\n        return {gen_expr(r, names)}")
    else:
        body.append(f"    return {gen_expr(r, names)}")
    src = f"def f({', '.join(p + ': int' for p in params)}) -> int:\n" + "\n".join(body) + "\n"
    pre = "True" if r.random() < 0.5 else gen_cond(r, params)
    post = gen_cond(r, params + ["result"])
    return src, pre, post, params


def run_py(src, params, values):
    env = {}
    exec(src, env)  # noqa: S102 (generated code in the supported subset only)
    return env["f"](*[values[p] for p in params])


def ev(expr, env):
    return eval(expr, {"__builtins__": {"abs": abs, "min": min, "max": max, "True": True, "False": False}}, dict(env))  # noqa: S307


def judge(src, pre, post, params, verdict):
    if not verdict.get("ok"):
        return "REFUSED", None
    if verdict["contract_holds"]:
        for vals in itertools.product(GRID, repeat=len(params)):
            env = dict(zip(params, vals))
            try:
                if not ev(pre, env):
                    continue
            except ZeroDivisionError:
                continue
            try:
                res = run_py(src, params, env)
            except ZeroDivisionError:
                return "PROVED_BUT_CRASHES", env
            try:
                if not ev(post, {**env, "result": res}):
                    return "FALSE_PROVED", env
            except ZeroDivisionError:
                return "PROVED_BUT_CRASHES", env
        return "PROVED_OK", None
    ce = verdict["counterexample"]
    env = {p: ce[p] for p in params}
    failure = verdict.get("failure", "postcondition_violated")
    stage = "pre"
    try:
        okpre = ev(pre, env)
        stage = "function"
        if not okpre:
            return "FALSE_COUNTEREXAMPLE", env
        res = run_py(src, params, env)
        stage = "post"
        held = ev(post, {**env, "result": res})
    except ZeroDivisionError:
        want = {"pre": "precondition_raises", "function": "function_raises", "post": "postcondition_raises"}[stage]
        if failure == want:
            return "RAISES_COUNTEREXAMPLE_OK", None
        return ("CRASH_COUNTEREXAMPLE" if failure == "postcondition_violated" else "WRONG_FAILURE_KIND"), env
    if failure != "postcondition_violated":
        return "WRONG_FAILURE_KIND", env
    if held:
        return "FALSE_COUNTEREXAMPLE", env
    return "COUNTEREXAMPLE_OK", None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=3000)
    ap.add_argument("--seed", type=int, default=1)
    ap.add_argument("--json")
    a = ap.parse_args()
    r = random.Random(a.seed)
    counts, examples = {}, {}
    for _ in range(a.n):
        src, pre, post, params = gen_case(r)
        try:
            v = V.check_contract(src, pre, post)
            kind, at = judge(src, pre, post, params, v)
        except Exception as e:  # noqa: BLE001
            kind, at = "ENGINE_ERROR", f"{type(e).__name__}: {e}"[:200]
        counts[kind] = counts.get(kind, 0) + 1
        if kind not in ("PROVED_OK", "COUNTEREXAMPLE_OK", "RAISES_COUNTEREXAMPLE_OK", "REFUSED") and len(examples.setdefault(kind, [])) < 3:
            examples[kind].append({"source": src, "pre": pre, "post": post, "at": at})
    print(json.dumps({"n": a.n, "seed": a.seed, "counts": dict(sorted(counts.items()))}, indent=1))
    for k, ex in examples.items():
        print(f"\n== {k}")
        for e in ex[:2]:
            print(e["source"] + f"   pre: {e['pre']}\n   post: {e['post']}\n   at: {e['at']}")
    if a.json:
        json.dump({"n": a.n, "seed": a.seed, "grid": [GRID.start, GRID.stop - 1], "counts": counts, "examples": examples},
                  open(a.json, "w"), indent=1)
    bad = sum(counts.get(k, 0) for k in ("FALSE_PROVED", "FALSE_COUNTEREXAMPLE", "ENGINE_ERROR", "PROVED_BUT_CRASHES",
                                          "CRASH_COUNTEREXAMPLE", "WRONG_FAILURE_KIND"))
    return 1 if bad else 0


if __name__ == "__main__":
    sys.exit(main())