import argparse
import json

from laya_x64 import DEFAULT_MODEL, Analyzer


def describe(value):
    if isinstance(value, dict):
        return f"{value['answer']}  ({value['p']:.0%} sure)"
    answer = "yes" if value >= 0.5 else "no"
    confidence = max(value, 1 - value)
    note = "   <- unsure" if confidence < 0.8 else ""
    return f"{answer}  ({confidence:.0%} sure){note}"


def main():
    ap = argparse.ArgumentParser(description="Ask the x64-trained Laya model about instructions.")
    ap.add_argument("instructions", nargs="+", help='Intel syntax, e.g. "pdep rax, rbx, rcx"')
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--full", action="store_true", help="also ask per-flag and per-register questions")
    ap.add_argument("--json", action="store_true", help="print raw probabilities (noul = probability of yes)")
    args = ap.parse_args()
    analyzer = Analyzer(args.model)
    for text in args.instructions:
        answers = analyzer.ask(text, full=args.full)
        print(f"\n{text}")
        if args.json:
            print(json.dumps(answers, indent=2))
            continue
        width = max(len(k) for k in answers)
        for question, value in answers.items():
            print(f"  {question:<{width}}  {describe(value)}")


if __name__ == "__main__":
    main()