#!/usr/bin/env python3
"""Generate gold-standard files for all fold-relevant Juliet CWE categories.

Reads bench/juliet/expansion/cwe_catalog.json to determine which CWEs map to
fold detectors. For each mapped CWE that has a filtered corpus directory,
generates a bench/gold/juliet-<cwe>.json file in the same schema used by
bench/gold/_generate_juliet.py.

Usage:
  # Generate gold for the 6 target detectors only
  python3 bench/juliet/expansion/_generate_gold.py --detectors null-deref,div-by-zero,int-overflow,int-truncation,sign-conversion,oob-access

  # Generate gold for ALL fold-mapped CWEs (including uninit, UAF, etc.)
  python3 bench/juliet/expansion/_generate_gold.py --all-detectors

  # Generate gold for a specific set of CWEs
  python3 bench/juliet/expansion/_generate_gold.py --cwes 476,690,369
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
HERE = Path(__file__).resolve().parent
CATALOG_PATH = HERE / "cwe_catalog.json"
FILTERED_ROOT = REPO_ROOT / "bench" / "corpora" / "juliet" / "filtered"
GOLD_ROOT = REPO_ROOT / "bench" / "gold"

# ── CWE kind classification (same logic as bench/gold/_generate_juliet.py) ──

CWE_KIND = {
    "476": "null-deref",
    "369": "div-by-zero",
    "190": "int-overflow",
    "191": "int-overflow",
    "121": "oob-access",
    "122": "oob-access",
    "124": "oob-access",
    "126": "oob-access",
    "127": "oob-access",
    "194": "sign-conversion",
    "195": "sign-conversion",
    "196": "sign-conversion",
    "197": "int-truncation",
    "561": "dead",
    # ── Expanded CWEs ──
    "690": "null-deref",
    "252": "null-deref",
    "588": "null-deref",
    "562": "null-deref",
    "680": "int-overflow",
    "789": "int-overflow",
    "681": "int-truncation",
    "467": "sign-conversion",
    "123": "oob-access",
    "468": "oob-access",
    "785": "oob-access",
    # ── Other fold detectors (not the 6 target detectors) ──
    "457": "uninit-local-read",
    "665": "uninit-local-read",
    "415": "use-after-free",
    "416": "use-after-free",
    "401": "use-after-free",
    "590": "use-after-free",
    "672": "use-after-free",
    "675": "use-after-free",
    "761": "use-after-free",
    "762": "use-after-free",
    "242": "deprecated-api",
    "676": "deprecated-api",
    "78": "taint-track",
    "114": "taint-track",
    "134": "taint-track",
    "426": "taint-track",
    # CWE-561 special cases handled by kind_for_path
    "478": "dead",
    "563": "dead",
    "570": "dead",
    "571": "dead",
}

# ── Specialized kind routing for CWE-561 (from _generate_juliet.py) ──

CWE_561_SUB_KINDS = {
    "return_before_code": "dead-statement",
    "unused_function": "unused-function",
}

BAD_NAME_RE = re.compile(r"(?:__bad|_bad_|_bad\.)", re.IGNORECASE)
GOOD_NAME_RE = re.compile(r"(?:__good|_good_|_good\.)", re.IGNORECASE)
BAD_CONTENT_RE = re.compile(r"\bbad\s*\(")
GOOD_CONTENT_RE = re.compile(r"\bgood(?:[A-Z0-9_]+)?\s*\(")


def is_driver_file(path: Path) -> bool:
    lower = path.name.lower()
    return lower.endswith("_main.c") or lower.endswith("_main.cpp")


def classify_file(path: Path) -> str | None:
    """Return 'bad', 'good', or None for a Juliet testcase file."""
    if is_driver_file(path):
        return None
    name = path.name
    if BAD_NAME_RE.search(name):
        return "bad"
    if GOOD_NAME_RE.search(name):
        return "good"
    try:
        text = path.read_text(encoding="utf-8", errors="ignore")
    except OSError:
        return None
    has_bad = bool(BAD_CONTENT_RE.search(text))
    has_good = bool(GOOD_CONTENT_RE.search(text))
    if has_bad:
        return "bad"
    if has_good:
        return "good"
    return None


def kind_for_path(cwe: str, path: Path) -> str:
    """Pick the per-entry kind for a Juliet bad testcase."""
    base_kind = CWE_KIND.get(cwe, "unknown")
    if cwe == "561":
        name = path.name
        if "__return_before_code" in name:
            return "dead-statement"
        if "__unused_function" in name:
            return "unused-function"
        return "dead"
    return base_kind


def build_expected(cwe: str, corpus_dir: Path) -> list[dict]:
    """Build gold entries for a single CWE from its filtered corpus directory."""
    expected = []
    counter = 1
    for path in sorted(corpus_dir.rglob("*")):
        if path.suffix not in {".c", ".cpp"}:
            continue
        classification = classify_file(path)
        if classification != "bad":
            continue
        relative = path.relative_to(corpus_dir).as_posix()
        expected.append({
            "kind": kind_for_path(cwe, path),
            "files": [relative],
            "id": f"juliet-{cwe}-{counter:04d}",
        })
        counter += 1
    return expected


def write_gold(cwe: str, expected: list[dict], out_dir: Path) -> Path:
    """Write a gold-standard JSON file."""
    cwe_desc = CATALOG["cwe_entries"].get(cwe, {}).get("title", f"CWE-{cwe}")
    out_path = out_dir / f"juliet-{cwe}.json"
    payload = {
        "name": f"juliet-cwe-{cwe}",
        "description": f"Juliet CWE-{cwe} ({cwe_desc}) test cases",
        "expected": expected,
    }
    out_path.write_text(json.dumps(payload, indent=2) + "\n")
    return out_path


def load_catalog() -> dict:
    return json.loads(CATALOG_PATH.read_text())


CATALOG = load_catalog()


def resolve_cwes(
    detectors: list[str] | None,
    all_detectors: bool,
    cwe_list: list[str] | None,
) -> list[str]:
    """Resolve which CWEs to generate gold for."""
    if cwe_list:
        return sorted(set(cwe_list))

    cwe_entries = CATALOG["cwe_entries"]
    if all_detectors:
        # All CWEs with any fold detector mapping
        return sorted(
            cwe
            for cwe, entry in cwe_entries.items()
            if entry.get("fold_detector") is not None
        )

    if detectors:
        target_detectors = set(detectors)
        # Map detector aliases
        if "oob" in target_detectors:
            target_detectors.add("oob-access")
        if "oob-access" in target_detectors:
            target_detectors.add("oob")
        return sorted(
            cwe
            for cwe, entry in cwe_entries.items()
            if entry.get("fold_detector") in target_detectors
        )

    # Default: the 6 target detectors
    target_detectors = set(CATALOG["target_detectors"])
    return sorted(
        cwe
        for cwe, entry in cwe_entries.items()
        if entry.get("fold_detector") in target_detectors
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    cwe_group = parser.add_mutually_exclusive_group()
    cwe_group.add_argument(
        "--detectors",
        help="Comma-separated detector names (default: 6 target detectors)",
    )
    cwe_group.add_argument(
        "--all-detectors",
        action="store_true",
        help="Generate gold for all fold-mapped CWEs, not just the 6 targets",
    )
    cwe_group.add_argument(
        "--cwes",
        help="Comma-separated CWE numbers (e.g., 476,690)",
    )
    parser.add_argument(
        "--out-dir",
        default=str(GOLD_ROOT),
        help="Output directory for gold files (default: bench/gold/)",
    )
    args = parser.parse_args()

    detector_list = (
        [d.strip() for d in args.detectors.split(",") if d.strip()]
        if args.detectors
        else None
    )
    cwe_list = (
        [c.strip() for c in args.cwes.split(",") if c.strip()]
        if args.cwes
        else None
    )

    cwes = resolve_cwes(detector_list, args.all_detectors, cwe_list)

    if not cwes:
        print("No CWEs resolved. Check detector names or CWE numbers.", file=sys.stderr)
        return 1

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    total_entries = 0
    skipped = 0
    for cwe in cwes:
        corpus_dir = FILTERED_ROOT / cwe
        if not corpus_dir.is_dir():
            print(f"CWE-{cwe}: skipped — filtered corpus not found at {corpus_dir}",
                  file=sys.stderr)
            print(f"  Run: python3 bench/corpora/juliet/filter_cwes.py first",
                  file=sys.stderr)
            skipped += 1
            continue

        expected = build_expected(cwe, corpus_dir)
        out_path = write_gold(cwe, expected, out_dir)
        entry = CATALOG["cwe_entries"].get(cwe, {})
        detector = entry.get("fold_detector", "-")
        print(f"CWE-{cwe}: {len(expected):5d} gold entries → {out_path.name} "
              f" (detector: {detector})")
        total_entries += len(expected)

    print(f"\nTotal: {total_entries} gold entries across {len(cwes) - skipped} CWEs "
          f"({skipped} skipped — run filter_cwes.py first)")
    print(f"Output directory: {out_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())