import argparse
import json
from pathlib import Path

import laya

ROOT = Path(__file__).parent
DEFAULT_MODEL = str(ROOT / "models" / "laya-x64-func")

from build_labeling import (ALGO_OPTIONS, CATEGORY_OPTIONS, BUG_OPTIONS, ALGO_Q, CAT_Q, VULN_Q, BUG_Q)

QUESTIONS = {
    "algorithm": {"type": "choice", "instructions": ALGO_Q, "criteria": {o: None for o in ALGO_OPTIONS}},
    "category": {"type": "choice", "instructions": CAT_Q, "criteria": {o: None for o in CATEGORY_OPTIONS}},
    "vulnerable": {"type": "noul", "instructions": VULN_Q},
    "bug_type": {"type": "choice", "instructions": BUG_Q, "criteria": {o: None for o in BUG_OPTIONS}},
}


def read_asm(source):
    """source is a path to a text file of assembly, or raw hex bytes to disassemble with Capstone."""
    p = Path(source)
    if p.exists():
        return p.read_text(encoding="utf-8").strip()
    import capstone
    from laya_x64 import normalize
    md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
    data = bytes.fromhex(source.replace(" ", ""))
    return "\n".join(normalize(i.mnemonic, i.op_str) for i in md.disasm(data, 0x1000))


def main():
    ap = argparse.ArgumentParser(description="Ask the v2 model about a whole x86-64 function.")
    ap.add_argument("source", help="path to an .asm text file, or a hex byte string to disassemble")
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--max-len", type=int, default=1024)
    ap.add_argument("--json", action="store_true")
    args = ap.parse_args()

    asm = read_asm(args.source)
    agent = laya.load(args.model, device="cuda")
    result = agent.predict({"function": asm}, QUESTIONS, max_len=args.max_len)["answers"]

    n_insns = sum(1 for l in asm.splitlines() if l.strip() and not l.rstrip().endswith(":"))
    print(f"\nfunction: {n_insns} instructions\n" + "-" * 40)
    if args.json:
        print(json.dumps(result, indent=2, default=str)); return
    a = result["algorithm"]
    print(f"algorithm : {a['choice']:24} ({a['probabilities'][a['choice']]:.0%})")
    c = result["category"]
    print(f"category  : {c['choice']:24} ({c['probabilities'][c['choice']]:.0%})")
    v = result["vulnerable"]["noul"]
    print(f"vulnerable: {'YES' if v >= 0.5 else 'no':24} ({max(v, 1 - v):.0%})")
    if v >= 0.5:
        b = result["bug_type"]
        print(f"bug_type  : {b['choice']:24} ({b['probabilities'][b['choice']]:.0%})")


if __name__ == "__main__":
    main()