#!/usr/bin/env python3
"""Rebuild <slice>/catalog.toml and <slice>/training/manifest.jsonl from <slice>/recipes.toml + excerpts.

Shared by every harvested slice (vfc, vibe-py, vfc-js, agent-vfc). Recipes carry the
labels; the catalog is generated data. Do not hand-edit generated catalogs.
"""

from __future__ import annotations

import argparse
import json
import re
import shutil
import tomllib
from pathlib import Path

try:
    from openultrasast.semantic.functions import ANON_FUNCTION, named_function_ranges
except ImportError as exc:  # never degrade into "no named function": that path deletes corpora under --prune
    raise SystemExit(f"catalog_gen needs the openultrasast package (run it with the project venv): {exc}") from exc

ROOT = Path(__file__).resolve().parent
LANG_TO_PAIR = {
    "c": "c_cpp",
    "cpp": "c_cpp",
    "perl": "unknown",
    "javascript": "javascript",
    "typescript": "typescript",
    "python": "python",
    "java": "java",
}


def excerpt_rel(recipe: dict[str, object]) -> tuple[str, str]:
    name = str(recipe["name"])
    ext = Path(str(recipe.get("path", "a.c"))).suffix or ".c"
    return f"{name}/vuln{ext}", f"{name}/fixed{ext}"


def _q(value: object) -> str:
    return str(value).replace("\\", "\\\\").replace('"', '\\"')


def catalog_text(slice_name: str, recipes: list[dict[str, object]]) -> str:
    lines = [
        f"# {slice_name} pair slice. Reviewed isolated functions with provenance.",
        "# Honesty dashboard — not the local 100% CI gate.",
        "# Do not add these names to LANGUAGE_MANIFESTS.",
        "# Regenerated by benchmarks/pairs/catalog_gen.py from recipes.toml.",
        "",
    ]
    for recipe in recipes:
        vuln, fixed = excerpt_rel(recipe)
        lang = LANG_TO_PAIR.get(str(recipe.get("language", "c")), "c_cpp")
        pointer = recipe.get("vendored") is False
        lines += [
            "[[pair]]",
            f'name = "{_q(recipe["name"])}"',
            f'slice = "{slice_name}"',
            f'language = "{lang}"',
            f'origin = "{_q(recipe.get("origin", ""))}"',
            f'license = "{_q(recipe.get("license", ""))}"',
            f'repo = "{_q(recipe.get("repo", ""))}"',
            f'cve = "{_q(recipe.get("cve", ""))}"',
            f'cwe = "{_q(recipe.get("cwe", ""))}"',
            f'commit = "{_q(recipe.get("commit", ""))}"',
            f'parent = "{_q(recipe.get("parent", ""))}"',
            f'commit_url = "{_q(recipe.get("commit_url", ""))}"',
        ]
        if pointer:
            lines.append("vendored = false")  # Req 10.1: the recipe travels in the row, no excerpt in the repository
            lines += _pointer_lines(recipe)
        else:
            lines += [f'vuln = "{vuln}"', f'fixed = "{fixed}"']
        lines += [
            f'relpath = "{_q(recipe.get("relpath", recipe.get("path", "")))}"',
            "min_recall = 1.0",
            'fix_policy = "silent"',
            f'provenance = "{_q(recipe.get("provenance", "human"))}"',
            f'split = "{_q(recipe.get("split", "train"))}"',
            f'review_tier = "{_review_tier(slice_name, recipe)}"',
        ]
        reviewer = str(recipe.get("reviewer", "") or "")
        if reviewer and reviewer.lower() != "pending":
            lines.append(f'reviewer = "{_q(reviewer)}"')
        if recipe.get("known_limit"):
            lines.append(f'known_limit = "{_q(recipe["known_limit"])}"')
        lines += [
            "",
            "[[pair.expected]]",
            f'cwe = "{_q(recipe.get("cwe", ""))}"',
            f'class = "{_q(recipe.get("class", ""))}"',
            f'path = "{_q(recipe.get("relpath", recipe.get("path", "")))}"',
        ]
        if recipe.get("function"):
            lines.append(f'function = "{_q(recipe["function"])}"')
        lines.append(f'mechanism = "{_q(recipe.get("mechanism", "other"))}"')
        if recipe.get("rule_id"):
            lines.append(f'rule_id = "{_q(recipe["rule_id"])}"')
        if recipe.get("sink"):
            lines.append(f'sink = "{_q(recipe["sink"])}"')
        if recipe.get("evidence"):
            lines.append(f'evidence = "{_q(recipe["evidence"])}"')
        lines.append("")
    return "\n".join(lines)


_DEFAULT_TIER = {"vibe-py": "seeded", "vfc-js": "advisory", "agent-vfc": "title", "vfc": "advisory", "sast": "advisory"}


def _review_tier(slice_name: str, recipe: dict[str, object]) -> str:
    """Req 9, same precedence as the loader: explicit `review_tier`, else `reviewer` (pending -> title, a name -> reviewed), else the slice default."""
    if recipe.get("review_tier"):
        return str(recipe["review_tier"])
    reviewer = str(recipe.get("reviewer", "") or "")
    if reviewer.lower() == "pending":
        return "title"
    if reviewer:
        return "reviewed"
    return _DEFAULT_TIER.get(slice_name, "advisory")


_POINTER_STR = ("host", "path", "fix_repo", "fix_path", "mode")
_POINTER_INT = ("line", "fix_line", "line_start", "line_end")


def _pointer_lines(recipe: dict[str, object]) -> list[str]:
    out = [f'{key} = "{_q(recipe[key])}"' for key in _POINTER_STR if recipe.get(key)]
    out += [f"{key} = {int(recipe[key])}" for key in _POINTER_INT if isinstance(recipe.get(key), int)]
    return out


def training_rows(slice_name: str, recipes: list[dict[str, object]]) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    for recipe in recipes:
        if recipe.get("vendored") is False:
            continue  # nothing vendored to train on
        vuln, fixed = excerpt_rel(recipe)
        base = {
            "pair": recipe["name"],
            "cve": recipe.get("cve", ""),
            "year": int(recipe.get("year", 0) or 0),
            "cwe": recipe.get("cwe", ""),
            "class": recipe.get("class", ""),
            "project": recipe.get("project", recipe.get("origin", "")),
            "repo": recipe.get("repo", ""),
            "license": recipe.get("license", ""),
            "function": recipe.get("function", ""),
            "mechanism": recipe.get("mechanism", "other"),
            "provenance": recipe.get("provenance", "human"),
            "split": recipe.get("split", "train"),
            "review_tier": _review_tier(slice_name, recipe),
            "relpath": recipe.get("relpath", recipe.get("path", "")),
            "commit_url": recipe.get("commit_url", ""),
        }
        rows.append({**base, "id": f"{recipe['name']}:vuln", "label": "vuln", "file": vuln, "commit": recipe.get("parent", "")})
        rows.append({**base, "id": f"{recipe['name']}:fixed", "label": "fixed", "file": fixed, "commit": recipe.get("commit", "")})
    return rows


# Sink call tokens per language, keyed by the CWE family the recipe declares, so a stray
# `re.exec(` in a code-injection excerpt is never mistaken for the labeled sink.
_SINK_TOKENS: dict[str, dict[str, dict[str, str]]] = {
    "javascript": {
        "CWE-94": {"eval": "eval", "Function": "eval"},
        "CWE-95": {"eval": "eval", "Function": "eval"},
        "CWE-78": {"exec": "exec", "execSync": "exec", "execFile": "exec", "execFileSync": "exec", "spawn": "exec", "spawnSync": "exec"},
        "CWE-22": {
            "join": "path",
            "resolve": "path",
            "readFile": "path",
            "readFileSync": "path",
            "createReadStream": "path",
            "sendFile": "path",
            "normalize": "path",
        },
        "CWE-89": {"query": "query", "execute": "query"},
    },
    "python": {
        "CWE-94": {"eval": "eval", "exec": "eval"},
        "CWE-95": {"eval": "eval", "exec": "eval"},
        "CWE-78": {
            "system": "os.system",
            "popen": "os.system",
            "run": "os.system",
            "Popen": "os.system",
            "call": "os.system",
            "check_output": "os.system",
        },
        "CWE-89": {"execute": "execute", "executemany": "execute"},
        "CWE-502": {"loads": "pickle.loads", "load": "pickle.loads"},
        "CWE-22": {"open": "open", "join": "path", "send_file": "path"},
    },
    "c": {
        "CWE-119": {"memcpy": "memcpy", "memmove": "memcpy", "strcpy": "strcpy", "strcat": "strcpy", "sprintf": "printf"},
        "CWE-120": {"strcpy": "strcpy", "strcat": "strcpy", "memcpy": "memcpy"},
        "CWE-78": {"system": "system", "popen": "system"},
        "CWE-134": {"printf": "printf", "sprintf": "printf", "fprintf": "printf"},
    },
    "java": {
        "CWE-78": {"exec": "exec"},
        "CWE-89": {
            "executeQuery": "executeQuery",
            "execute": "executeQuery",
            "prepareCall": "executeQuery",
            "prepareStatement": "executeQuery",
        },
        "CWE-502": {"readObject": "readObject"},
        "CWE-327": {"getInstance": "MessageDigest"},
    },
}
_SINK_TOKENS["typescript"] = _SINK_TOKENS["javascript"]
_SINK_TOKENS["cpp"] = _SINK_TOKENS["c"]
_REGEX_EXEC = re.compile(r"(?:/[^/\n]*/[gimsuy]*|\)|\b(?:re|regex|regexp|pattern|matcher)[A-Za-z_]*)\.exec\s*\(")
_HEADER_LINE = re.compile(r"^\s*(//|#|/\*|\*|\*/)")
_UPSTREAM_START = re.compile(r"^(?://|#| \*)\s*upstream_start:\s*(\d+)\s*$", re.M)


def _body_start(lines: list[str]) -> int:
    """0-based index of the first line after the provenance header block (leading comment lines plus blanks)."""
    index = 0
    while index < len(lines) and (_HEADER_LINE.match(lines[index]) or not lines[index].strip()):
        index += 1
    return index


def _language(recipe: dict[str, object]) -> str:
    return {"python": "python", "typescript": "typescript", "javascript": "javascript", "java": "java", "cpp": "cpp"}.get(
        str(recipe.get("language", "c")), "c"
    )


def _ranges(path: Path, language: str) -> tuple[tuple[str, int, int], ...] | None:
    ranges = named_function_ranges(str(path), path.read_text(errors="ignore"), language)
    if ranges is None:
        return None  # parser rejected the excerpt; a regex label would not scope
    return tuple(r for r in ranges if r[0] != ANON_FUNCTION)


def _innermost(ranges: tuple[tuple[str, int, int], ...], line: int) -> str | None:
    inside = [r for r in ranges if r[1] <= line <= r[2]]
    return min(inside, key=lambda r: r[2] - r[1])[0] if inside else None


def _sink_tokens(recipe: dict[str, object]) -> dict[str, str]:
    by_cwe = _SINK_TOKENS.get(_language(recipe), {})
    tokens = by_cwe.get(str(recipe.get("cwe", "")), {})
    if not tokens:  # unknown CWE family: accept any token of the language
        for family in by_cwe.values():
            tokens = {**tokens, **family}
    return tokens


def _sink_on_line(line: str, tokens: dict[str, str]) -> str | None:
    if _REGEX_EXEC.search(line):
        return None
    for token, sink_id in tokens.items():
        if re.search(r"(?<![\w$])" + re.escape(token) + r"\s*\(", line):
            return sink_id
    return None


def _sink_line_in_excerpt(recipe: dict[str, object], text: str, ranges: tuple[tuple[str, int, int], ...]) -> tuple[int | None, str | None]:
    """Excerpt line of the labeled sink: mapped from ``sink_location`` + ``upstream_start`` when possible, else the
    unique named function that contains a sink token of the recipe's CWE family."""
    lines = text.splitlines()
    tokens = _sink_tokens(recipe)
    if not tokens:
        return None, None
    location = str(recipe.get("sink_location") or "")
    parts = location.rsplit(":", 2)
    relpath = str(recipe.get("relpath", recipe.get("path", "")))
    match = _UPSTREAM_START.search(text)
    if (
        len(parts) >= 2
        and parts[1].isdigit()
        and match
        and (parts[0] == relpath or relpath.endswith(parts[0]) or Path(parts[0]).name == Path(relpath).name)
    ):
        mapped = int(parts[1]) - int(match.group(1)) + _body_start(lines) + 1
        for delta in sorted(range(-10, 11), key=abs):  # tarball-vs-blob drift is usually a few lines
            line_no = mapped + delta
            if 1 <= line_no <= len(lines):
                sink_id = _sink_on_line(lines[line_no - 1], tokens)
                if sink_id:
                    return line_no, sink_id
        mapped_fallback = mapped if 1 <= mapped <= len(lines) and re.search(r"[A-Za-z_$][\w$.]*\s*\(", lines[mapped - 1]) else None
    else:
        mapped_fallback = None
    # Drift (SecBench.js locations come from the npm tarball, not the commit blob): fall back to the
    # sink tokens of the CWE family anywhere in the excerpt, accepted only when they name one function.
    all_tokens: dict[str, str] = {}
    for family in _SINK_TOKENS.get(_language(recipe), {}).values():
        all_tokens.update(family)
    for candidate_tokens in (tokens, all_tokens):
        hits = [(index, sink) for index, line in enumerate(lines, start=1) for sink in [_sink_on_line(line, candidate_tokens)] if sink]
        owners = {name for index, _ in hits for name in [_innermost(ranges, index)] if name}
        if len(owners) == 1:
            line_no, sink_id = next((index, sink) for index, sink in hits if _innermost(ranges, index) in owners)
            return line_no, sink_id
    # The upstream location is ground truth even when the sink is a project wrapper (`proc(cmd)`):
    # label the function enclosing the mapped line and emit no sink id.
    if mapped_fallback is not None:
        return mapped_fallback, None
    return None, None


def derive_function(slice_root: Path, recipe: dict[str, object]) -> str | None:
    """Label = innermost named function enclosing the sink (from ``sink_location`` + ``upstream_start``).

    Falls back to the function enclosing the first changed body line only when the
    recipe has no usable sink location, and says so. Never returns ``<anon>`` or a
    reserved word; returns None (skip the pair) when nothing names a function.
    """
    vuln, fixed = excerpt_rel(recipe)
    vuln_path, fixed_path = slice_root / vuln, slice_root / fixed
    language = _language(recipe)
    ranges = _ranges(vuln_path, language)
    if not ranges:
        return None
    text = vuln_path.read_text(errors="ignore")
    sink_line, sink_id = _sink_line_in_excerpt(recipe, text, ranges)
    if sink_line is not None:
        name = _innermost(ranges, sink_line)
        if name:
            if sink_id and not recipe.get("sink"):
                recipe["sink"] = sink_id
            return name
        print(f"weak: {recipe.get('name')}: sink line {sink_line} is not inside a named function; falling back to the diff")
    elif recipe.get("sink_location"):
        print(
            f"weak: {recipe.get('name')}: no sink token of the CWE family near the mapped sink_location; label from diff, no sink emitted"
        )
    a = text.splitlines()
    b = fixed_path.read_text(errors="ignore").splitlines()
    a_body, b_body = a[_body_start(a) :], b[_body_start(b) :]
    offset = _body_start(a)

    first_diff = _first_semantic_change(a_body, b_body)
    return _innermost(ranges, offset + first_diff + 1) or ranges[0][0]


def _first_semantic_change(a: list[str], b: list[str]) -> int:
    """0-based index in ``a`` of the first hunk that changes more than whitespace (formatting hunks never label a pair)."""
    import difflib

    def norm(chunk: list[str]) -> str:
        return re.sub(r"\s+", "", "".join(chunk))

    matcher = difflib.SequenceMatcher(a=a, b=b, autojunk=False)
    for tag, i1, i2, j1, j2 in matcher.get_opcodes():
        if tag == "equal":
            continue
        if norm(a[i1:i2]) == norm(b[j1:j2]):
            continue
        return i1 if i2 > i1 else max(i1 - 1, 0)
    return min(len(a), len(b))


def _stamp_function(path: Path, function: str) -> None:
    """Fill the empty ``function:`` header line of a hunk-harvested excerpt."""
    text = path.read_text(errors="ignore")
    updated = re.sub(r"^((?://|#| \*) function:).*$", rf"\1 {function}", text, count=1, flags=re.M)
    if updated != text:
        path.write_text(updated)


def generate(slice_name: str, *, prune: bool = False) -> tuple[Path, int]:
    slice_root = ROOT / slice_name
    loaded = [dict(item) for item in tomllib.loads((slice_root / "recipes.toml").read_text()).get("recipe", [])]
    recipes: list[dict[str, object]] = []
    for recipe in loaded:
        vuln, fixed = excerpt_rel(recipe)
        if recipe.get("vendored") is False:
            if not recipe.get("function"):
                print(f"skip {recipe['name']}: pointer recipe without a function label (nothing to derive from)")
                continue
            recipes.append(recipe)
        elif (slice_root / vuln).is_file() and (slice_root / fixed).is_file():
            if not recipe.get("function"):
                derived = derive_function(slice_root, recipe)
                if not derived:
                    print(f"skip {recipe['name']}: no named function in the excerpt (label would not scope)")
                    continue
                recipe["function"] = derived
                _stamp_function(slice_root / vuln, derived)
                _stamp_function(slice_root / fixed, derived)
            recipes.append(recipe)  # `reviewer = "pending"` rows load at tier `title` (Req 9.4); no separate queue file
        else:
            print(f"skip {recipe['name']}: excerpt missing (harvest failed or not fetched)")
    catalog = slice_root / "catalog.toml"
    catalog.write_text(catalog_text(slice_name, recipes), encoding="utf-8")
    candidates = slice_root / "catalog-candidates.toml"
    if prune:
        if not recipes:
            raise SystemExit(f"refusing --prune for {slice_name}: no recipe derived a label, nothing would survive")
        keep = {str(r["name"]) for r in recipes}
        for entry in sorted(slice_root.iterdir()):
            if entry.is_dir() and entry.name not in keep and entry.name not in {"training", "__pycache__"}:
                shutil.rmtree(entry)
                print(f"pruned {entry.name}: excerpt not in the catalog")
    if candidates.exists():
        candidates.unlink()  # retired by Req 9: the `title` tier is the review queue
    training = slice_root / "training" / "manifest.jsonl"
    training.parent.mkdir(parents=True, exist_ok=True)
    with training.open("w", encoding="utf-8") as handle:
        for row in training_rows(slice_name, recipes):
            handle.write(json.dumps(row, sort_keys=True) + "\n")
    return catalog, len(recipes)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--slice", default="vfc")
    parser.add_argument("--prune", action="store_true", help="delete excerpt directories that ended up outside the catalog")
    args = parser.parse_args(argv)
    catalog, count = generate(args.slice, prune=args.prune)
    print(f"wrote {catalog} ({count} pairs)")
    return 0


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