#!/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 context_rel(recipe: dict[str, object]) -> list[tuple[str, str, str]]:
    """`(relpath, vuln document, fixed document)` for every document that registers the handler from outside its module."""
    name = str(recipe["name"])
    value = recipe.get("context") or ()
    relpaths = (value,) if isinstance(value, str) else tuple(str(item) for item in value)
    return [(relpath, f"{name}/context/vuln/{relpath}", f"{name}/context/fixed/{relpath}") for relpath in relpaths]


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"])}"')
        if recipe.get("fix_function"):
            # Req 3.2: the fixed side of a trap pair is a different handler, and the scorer has to know which one
            lines.append(f'fix_function = "{_q(recipe["fix_function"])}"')
        fix_date = _fix_date(recipe)
        if fix_date:
            lines.append(f'fix_date = "{_q(fix_date)}"')  # Req 10.5: the post-cutoff slice needs a date to sort by
        for relpath, context_vuln, context_fixed in context_rel(recipe):
            lines += [
                "",
                "[[pair.context]]",
                f'relpath = "{_q(relpath)}"',
                f'vuln = "{_q(context_vuln)}"',
                f'fixed = "{_q(context_fixed)}"',
            ]
        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("family"):
            lines.append(f'family = "{_q(recipe["family"])}"')  # Req 1: the family routes the work, never a CWE lookup
        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)


def _fix_date(recipe: dict[str, object]) -> str:
    """The date of the fix at the precision the recipe recorded it: a full ISO day, or the year, or nothing.

    Every harvested recipe carries a year; almost none carries a day, and inventing `-01-01` would turn a year into
    a claim about a month. The post-cutoff comparison reads the earliest possible date of whatever precision is
    here, so a coarse date can only ever under-claim (learning-harness Req 10.5)."""
    recorded = str(recipe.get("fix_date", "") or "")
    if recorded:
        return recorded
    year = recipe.get("year")
    return str(year) if isinstance(year, int) and 1990 <= year <= 2100 else ""


_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 derive_fix_function(slice_root: Path, recipe: dict[str, object]) -> str:
    """The function the fixed excerpt actually holds, when it is not the labeled one.

    Real-Vuln-Benchmark traps take their fixed side from a different, correctly guarded handler of the same
    snapshot. Without this the fixed side resolves no span, every finding there counts as outside the labeled
    function, and the pair scores as silence instead of as a leak — always in the direction that flatters the
    number (learning-harness Req 3.2)."""
    _vuln, fixed = excerpt_rel(recipe)
    path = slice_root / fixed
    if not path.is_file():
        return ""
    ranges = _ranges(path, _language(recipe))
    if not ranges:
        return ""
    labeled = str(recipe.get("function", "") or "")
    names = [name for name, _start, _end in ranges]
    if labeled in names or len(set(names)) != 1:
        return ""  # the fixed side holds the labeled function, or holds several and the recipe must say which
    return names[0]


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)
            if not recipe.get("fix_function"):
                derived = derive_fix_function(slice_root, recipe)
                if derived:
                    recipe["fix_function"] = 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())