#!/usr/bin/env python3
"""Score Quodeq security runs against the OWASP Benchmark expected results.

Usage:
    score_owasp.py <run_dir> [<run_dir> ...] [--benchmark DIR] [--json out.json]

Two scorings per run, both at test-case (file) level:

  requirement-matched (primary)
      A test case is flagged only when a violation in its file was filed under
      a Quodeq security requirement that covers the test's CWE (CATEGORY_REQS
      below). Deterministic: it uses the requirement code, never the model's
      free-text label.

  any-finding (secondary)
      A test case is flagged when its file has any security violation at all.
      Penalises a model for real but unrelated findings in a safe test case.

Score = TPR - FPR, as in the Benchmark's own scorecards. The overall figure is
pooled over all test cases; `mean_of_categories` is the unweighted average of
the per-category scores, which is how the official scorecard reports it.

Quodeq requirements are coarser than CWEs (one requirement covers LDAP, XPath
and code injection; the crypto requirements overlap), so this is a
requirement-family match, not an exact CWE match.
"""
import argparse
import collections
import csv
import json
import os
import re

TEST_RE = re.compile(r"(BenchmarkTest\d{5})\.java")

# Benchmark category (CWE) -> Quodeq requirements that cover it.
CATEGORY_REQS = {
    "sqli": {"S-INT-2"},                                  # CWE-89
    "xss": {"S-INT-3"},                                   # CWE-79
    "cmdi": {"S-AUT-2"},                                  # CWE-78
    "pathtraver": {"S-AUT-3"},                            # CWE-22
    "ldapi": {"S-INT-7"},                                 # CWE-90
    "xpathi": {"S-INT-7"},                                # CWE-643
    "crypto": {"S-INT-5", "S-CON-11", "S-CON-15"},        # CWE-327
    "hash": {"S-INT-5", "S-CON-11", "S-CON-15"},          # CWE-328
    "weakrand": {"S-CON-11", "S-CON-15"},                 # CWE-330
    "securecookie": {"S-CON-12"},                         # CWE-614
    "trustbound": {"S-INT-19", "S-INT-1"},                # CWE-501
}


def load_expected(benchmark):
    exp = {}
    with open(os.path.join(benchmark, "expectedresults-1.2.csv")) as fh:
        for row in csv.reader(fh):
            if row and not row[0].startswith("#"):
                exp[row[0]] = (row[1], row[2] == "true", row[3])
    return exp


def confusion(tests, exp, flagged):
    tp = sum(1 for t in tests if exp[t][1] and t in flagged)
    fn = sum(1 for t in tests if exp[t][1] and t not in flagged)
    fp = sum(1 for t in tests if not exp[t][1] and t in flagged)
    tn = sum(1 for t in tests if not exp[t][1] and t not in flagged)
    tpr = tp / (tp + fn) if tp + fn else 0.0
    fpr = fp / (fp + tn) if fp + tn else 0.0
    return {"n": len(tests), "tp": tp, "fn": fn, "fp": fp, "tn": tn,
            "tpr": round(tpr, 4), "fpr": round(fpr, 4), "score": round(tpr - fpr, 4)}


def scorecard(tests_by_cat, exp, flagged):
    cats = {c: confusion(ts, exp, flagged) for c, ts in sorted(tests_by_cat.items())}
    overall = confusion(list(exp), exp, flagged)
    overall["mean_of_categories"] = round(sum(c["score"] for c in cats.values()) / len(cats), 4)
    return {"overall": overall, "categories": cats}


def score_run(run_dir, exp, tests_by_cat):
    report = json.load(open(os.path.join(run_dir, "evaluation", "security.json")))
    status = json.load(open(os.path.join(run_dir, "status.json")))
    any_flag, req_flag = set(), set()
    for f in report["violations"]:
        m = TEST_RE.search(str(f.get("file")))
        if not m or m.group(1) not in exp:
            continue
        test = m.group(1)
        any_flag.add(test)
        if f.get("req") in CATEGORY_REQS[exp[test][0]]:
            req_flag.add(test)
    return {
        "run_dir": run_dir,
        "provider": status.get("ai_provider"),
        "model": status.get("ai_model"),
        "commit": status.get("commit_sha"),
        "started_at": status.get("started_at"),
        "finalized_at": status.get("finalized_at"),
        "report_score": report.get("overallScore"),
        "report_grade": report.get("overallGrade"),
        "files_read": report.get("filesRead"),
        "source_files": report.get("sourceFileCount"),
        "totals": report.get("totals"),
        "requirement_matched": scorecard(tests_by_cat, exp, req_flag),
        "any_finding": scorecard(tests_by_cat, exp, any_flag),
    }


def print_run(out):
    t = out["totals"]
    print(f"\n## {out['provider']} / {out['model']}\n")
    print(f"Quodeq grade {out['report_score']} {out['report_grade']}, files read {out['files_read']},"
          f" violations {t['violationCount']}, compliance {t['complianceCount']}, severity {t['severity']}\n")
    for key in ("requirement_matched", "any_finding"):
        sc = out[key]
        o = sc["overall"]
        print(f"**{key}**: TPR {o['tpr']:.0%}, FPR {o['fpr']:.0%}, score {o['score']:+.2f}"
              f" (mean of categories {o['mean_of_categories']:+.2f})\n")
        print("| Category | Cases | TP | FN | FP | TN | TPR | FPR | Score |")
        print("|---|---|---|---|---|---|---|---|---|")
        for name, c in sorted(sc["categories"].items(), key=lambda kv: -kv[1]["score"]):
            print(f"| {name} | {c['n']} | {c['tp']} | {c['fn']} | {c['fp']} | {c['tn']} |"
                  f" {c['tpr']:.0%} | {c['fpr']:.0%} | {c['score']:+.2f} |")
        print()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("run_dirs", nargs="+")
    ap.add_argument("--benchmark", default=os.path.expanduser("~/quodeq/repos/BenchmarkJava"))
    ap.add_argument("--json")
    args = ap.parse_args()

    exp = load_expected(args.benchmark)
    tests_by_cat = collections.defaultdict(list)
    for test, (cat, _, _) in exp.items():
        tests_by_cat[cat].append(test)

    results = [score_run(r.rstrip("/"), exp, tests_by_cat) for r in args.run_dirs]
    for out in results:
        print_run(out)
    if args.json:
        with open(args.json, "w") as fh:
            json.dump(results if len(results) > 1 else results[0], fh, indent=2)


if __name__ == "__main__":
    main()