codex-py/codex/provenance.py

"""Provenance tracking: @cite-scan, code_links CRUD, references.bib export."""

from __future__ import annotations

import re
from pathlib import Path

from codex.db import get_conn
from codex.models import CodeLink

# ---------------------------------------------------------------------------
# @cite scanning
# ---------------------------------------------------------------------------

_CITE_RE = re.compile(r"@cite\s+(\S+)", re.IGNORECASE)
_FUNC_RE = re.compile(r"(\w+)\s*\(")
_CLASS_RE = re.compile(r"class\s+(\w+)")

_CPP_SUFFIXES = {".cpp", ".h", ".cxx", ".hpp"}


def scan_cite_tags(root_dir: str) -> list[dict[str, str]]:
    """Recursively scan C++ source files for @cite tags (Doxygen-style).

    Returns one dict per hit:
        file   (str) — relative path to the source file
        line   (int as str) — 1-based line number
        symbol (str) — function or class name (extracted from context, best-effort)
        bibkey (str) — the @cite argument

    Scans files matching: *.cpp, *.h, *.cxx, *.hpp
    Pattern matched: ``@cite <bibkey>`` (case-insensitive, optional whitespace)
    """
    root = Path(root_dir)
    results: list[dict[str, str]] = []

    for path in sorted(root.rglob("*")):
        if path.suffix.lower() not in _CPP_SUFFIXES:
            continue
        try:
            lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
        except OSError:
            continue

        rel_path = str(path.relative_to(root))

        for lineno, line in enumerate(lines, start=1):
            for match in _CITE_RE.finditer(line):
                bibkey = match.group(1)
                # Best-effort symbol extraction: scan backwards for func or class
                symbol = ""
                context_lines = lines[: lineno - 1]
                for prev_line in reversed(context_lines):
                    func_match = _FUNC_RE.search(prev_line)
                    class_match = _CLASS_RE.search(prev_line)
                    if func_match:
                        symbol = func_match.group(1)
                        break
                    if class_match:
                        symbol = class_match.group(1)
                        break

                results.append(
                    {
                        "file": rel_path,
                        "line": str(lineno),
                        "symbol": symbol,
                        "bibkey": bibkey,
                    }
                )

    return results


# ---------------------------------------------------------------------------
# code_links CRUD
# ---------------------------------------------------------------------------


def add_code_link(
    symbol: str,
    paper_id: str | None = None,
    role: str | None = None,
    note: str | None = None,
) -> CodeLink:
    """Insert a new code_link row and return the populated CodeLink (with db-assigned id)."""
    sql = """
        INSERT INTO code_links (symbol, paper_id, role, note)
        VALUES (%(symbol)s, %(paper_id)s, %(role)s, %(note)s)
        RETURNING id, added_at
    """
    with get_conn() as conn:
        row = conn.execute(
            sql,
            {"symbol": symbol, "paper_id": paper_id, "role": role, "note": note},
        ).fetchone()
        conn.commit()
    assert row is not None
    return CodeLink(
        symbol=symbol,
        paper_id=paper_id,
        role=role,
        note=note,
        id=row["id"],
        added_at=row["added_at"],
    )


def list_code_links(paper_id: str | None = None) -> list[CodeLink]:
    """Return all code_links, optionally filtered by paper_id."""
    if paper_id is not None:
        sql = (
            "SELECT id, symbol, paper_id, role, note, added_at"
            " FROM code_links WHERE paper_id = %(paper_id)s"
        )
        params: dict[str, str] = {"paper_id": paper_id}
    else:
        sql = "SELECT id, symbol, paper_id, role, note, added_at FROM code_links"
        params = {}

    with get_conn() as conn:
        rows = conn.execute(sql, params).fetchall()

    return [
        CodeLink(
            symbol=row["symbol"],
            paper_id=row["paper_id"],
            role=row["role"],
            note=row["note"],
            id=row["id"],
            added_at=row["added_at"],
        )
        for row in rows
    ]


# ---------------------------------------------------------------------------
# references.bib export
# ---------------------------------------------------------------------------


def export_bib(paper_ids: list[str] | None = None) -> str:
    """Export papers as a BibTeX string.

    If paper_ids is None, exports ALL papers with a non-NULL bibkey.
    Only papers with a non-NULL bibkey are included (bibkey is the BibTeX key).

    Entry format (article):
        @article{<bibkey>,
          title  = {<title>},
          author = {<authors joined with " and ">},
          year   = {<year>},
          note   = {<id>},
        }
    """
    if paper_ids is not None:
        sql = (
            "SELECT id, bibkey, title, authors, year"
            " FROM papers"
            " WHERE bibkey IS NOT NULL AND id = ANY(%(paper_ids)s)"
        )
        params: dict[str, list[str]] = {"paper_ids": paper_ids}
    else:
        sql = "SELECT id, bibkey, title, authors, year FROM papers WHERE bibkey IS NOT NULL"
        params = {}

    with get_conn() as conn:
        rows = conn.execute(sql, params).fetchall()

    entries: list[str] = []
    for row in rows:
        authors: list[str] = row["authors"] or []
        author_str = " and ".join(authors) if authors else ""
        entry = (
            f"@article{{{row['bibkey']},\n"
            f"  title  = {{{row['title']}}},\n"
            f"  author = {{{author_str}}},\n"
            f"  year   = {{{row['year']}}},\n"
            f"  note   = {{{row['id']}}},\n"
            f"}}"
        )
        entries.append(entry)

    return "\n\n".join(entries)