#!/usr/bin/env python3 -I
"""Emit a Kani / proptest / cargo-fuzz project for one bank. Usage: gen_harness.py BANK ARM OUTDIR [pre|post]"""
import json, os, shutil, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import property_bank as pb


def kani(bank, out):
    lib = bank["lib_name"]
    body = ["use %s::*;" % lib]
    for p in bank["properties"]:
        if any(t not in pb.SCALARS for _, t in p["args"]):
            continue  # Kani arm: scalar-argument properties only
        decl = "\n    ".join(f"let {n}: {t} = kani::any();" for n, t in p["args"])
        call = ", ".join(n for n, _ in p["args"])
        body.append(f"#[kani::proof]\nfn proof_{p['id']}() {{\n    {decl}\n    kani::assume({p['pre']});\n"
                    f"    let r: {p['ret']} = {lib}::{p['fn']}({call});\n    assert!({p['post']});\n}}")
    os.makedirs(out + "/tests", exist_ok=True)
    open(out + "/tests/kani_props.rs", "w").write("\n\n".join(body) + "\n")
    open(out + "/Cargo.toml", "w").write(_cargo(bank, "kani", ""))


def proptest(bank, out):
    lib = bank["lib_name"]
    body = ["use proptest::prelude::*;", "use %s::*;" % lib, "proptest! {", "#![proptest_config(ProptestConfig { cases: u32::MAX, max_shrink_iters: 0, failure_persistence: None, .. ProptestConfig::default() })]"]
    for p in bank["properties"]:
        sig = ", ".join(f"{n} in any::<{pb.rust_type(t)}>()" for n, t in p["args"])
        call = ", ".join(pb.call_arg(n, t) for n, t in p["args"])
        body.append(f"#[test]\nfn prop_{p['id']}({sig}) {{\n    prop_assume!({p['pre']});\n    let r: {p['ret']} = {lib}::{p['fn']}({call});\n    let ok: bool = {{ {p['post']} }};\n    prop_assert!(ok, \"property {p['id']} violated\");\n}}")
    body.append("}")
    os.makedirs(out + "/tests", exist_ok=True)
    open(out + "/tests/pt_props.rs", "w").write("\n".join(body) + "\n")
    open(out + "/Cargo.toml", "w").write(_cargo(bank, "proptest", 'proptest = "1"\n'))


def fuzz(bank, out):
    lib = bank["lib_name"]
    os.makedirs(out + "/fuzz/fuzz_targets", exist_ok=True)
    names = []
    open(out + "/Cargo.toml", "w").write('[package]\nname = "fuzz-root"\nversion = "0.0.0"\nedition = "2021"\n\n[lib]\npath = "empty.rs"\n\n[workspace]\n')
    for p in bank["properties"]:
        fields = "\n".join(f"    {n}: {pb.rust_type(t)}," for n, t in p["args"])
        call = ", ".join(pb.call_arg(n, t) for n, t in p["args"])
        binds = "\n    ".join(f"let {n} = i.{n};" for n, _ in p["args"])
        src = (f"#![no_main]\nuse libfuzzer_sys::fuzz_target;\nuse arbitrary::Arbitrary;\nuse {lib}::*;\n"
               f"#[derive(Arbitrary, Debug)]\nstruct In {{\n{fields}\n}}\n"
               f"fuzz_target!(|i: In| {{\n    {binds}\n    if !({p['pre']}) {{ return; }}\n"
               f"    let r: {p['ret']} = {lib}::{p['fn']}({call});\n    let ok: bool = {{ {p['post']} }};\n    assert!(ok);\n}});\n")
        open(f"{out}/fuzz/fuzz_targets/{p['id']}.rs", "w").write(src); names.append(p["id"])
    bins = "".join(f'\n[[bin]]\nname = "{n}"\npath = "fuzz_targets/{n}.rs"\ntest = false\ndoc = false\nbench = false\n' for n in names)
    open(out + "/fuzz/Cargo.toml", "w").write(
        f'[package]\nname = "fuzz-harness"\nversion = "0.0.0"\nedition = "2021"\npublish = false\n\n[package.metadata]\ncargo-fuzz = true\n\n'
        f'[dependencies]\nlibfuzzer-sys = "0.4"\narbitrary = {{ version = "1", features = ["derive"] }}\n{lib} = {{ path = "{os.path.abspath(bank["crate_dir"])}" }}\n{bins}\n[workspace]\n')


def _cargo(bank, name, deps):
    return (f'[package]\nname = "{name}-harness"\nversion = "0.0.0"\nedition = "2021"\n\n[lib]\npath = "empty.rs"\n\n[dependencies]\n{deps}'
            f'{bank["lib_name"]} = {{ path = "{os.path.abspath(bank["crate_dir"])}" }}\n\n[workspace]\n')


def main():
    bank_path, arm, out = sys.argv[1:4]
    bank = pb.load(bank_path)
    rev = sys.argv[4] if len(sys.argv) > 4 else "pre"
    bank["crate_dir"] = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(bank_path)), bank["crate_dir"]))
    bank = pb.revision(bank, rev)
    os.makedirs(out, exist_ok=True)
    open(out + "/empty.rs", "w").write("")
    {"kani": kani, "proptest": proptest, "fuzz": fuzz}[arm](bank, out)
    print("generated", arm, "->", out)

if __name__ == "__main__":
    main()