reuse-surface/reuse_surface/plan_check.py

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import yaml

from jsonschema import Draft202012Validator

from reuse_surface.federation import FEDERATED_INDEX_PATH
from reuse_surface.llm_bridge import execute_prompt, extract_json_object
from reuse_surface.overlaps import TOKEN_RE
from reuse_surface.registry import ROOT, load_index
from reuse_surface.statehub_bridge import file_capability_request

TELEMETRY_PATH = ROOT / "registry" / "telemetry" / "plan-check-events.jsonl"
RERANK_SCHEMA_PATH = ROOT / "schemas" / "plan-check-rerank.schema.json"

STALE_DAYS = 14
DEFAULT_REUSE_THRESHOLD = 0.45
DEFAULT_EXTEND_THRESHOLD = 0.22
DEFAULT_TIE_WINDOW = 0.05


@dataclass
class MatchQuery:
    source: str
    text: str
    workplan_id: str | None = None
    workplan_path: str | None = None
    tokens: set[str] = field(default_factory=set)


@dataclass
class Match:
    id: str
    score: float
    kind: str
    vector: str | None = None
    owner: str | None = None
    summary: str | None = None


def _tokens(text: str) -> set[str]:
    return set(TOKEN_RE.findall(text.lower()))


def load_query_from_workplan(path: Path) -> MatchQuery:
    text = path.read_text(encoding="utf-8")
    match = re.match(r"^---\n(.*?)\n---\n?(.*)$", text, re.DOTALL)
    if not match:
        raise ValueError(f"{path}: missing YAML front matter")
    front = yaml.safe_load(match.group(1)) or {}
    body = match.group(2)

    parts = [str(front.get("title") or front.get("id") or path.stem)]
    for heading in ("Core Idea", "Problem statement", "One-liner"):
        section = re.search(
            rf"^##\s+{re.escape(heading)}\s*\n(.*?)(?:\n##\s|\Z)",
            body,
            re.DOTALL | re.MULTILINE,
        )
        if section:
            parts.append(section.group(1).strip())
    if len(parts) == 1:
        intro = body.strip().split("\n\n", 1)[0]
        parts.append(intro)

    blob = "\n".join(p for p in parts if p)
    return MatchQuery(
        source="workplan",
        text=blob,
        workplan_id=front.get("id"),
        workplan_path=str(path),
        tokens=_tokens(blob),
    )


def load_query_from_intent(intent: str) -> MatchQuery:
    return MatchQuery(source="intent", text=intent, tokens=_tokens(intent))


def _federated_entry_blob(item: dict[str, Any]) -> str:
    parts = [
        item.get("name", ""),
        item.get("summary", ""),
        " ".join(item.get("tags", [])),
    ]
    return " ".join(str(p) for p in parts if p)


def load_federated_capabilities() -> tuple[list[dict[str, Any]], str | None]:
    """Returns (capabilities, updated_date). Falls back to the local index
    if the composed federated index doesn't exist yet."""
    if FEDERATED_INDEX_PATH.exists():
        data = yaml.safe_load(FEDERATED_INDEX_PATH.read_text(encoding="utf-8"))
        return data.get("capabilities", []), data.get("updated")
    data = load_index()
    return data.get("capabilities", []), data.get("updated")


def _staleness_warning(updated: str | None) -> str | None:
    if not updated:
        return None
    try:
        updated_date = datetime.strptime(updated, "%Y-%m-%d").replace(
            tzinfo=timezone.utc
        )
    except ValueError:
        return None
    age_days = (datetime.now(timezone.utc) - updated_date).days
    if age_days > STALE_DAYS:
        return f"federated index is {age_days} days old (last composed {updated})"
    return None


def _vector_rank(vector: str | None) -> tuple[int, int]:
    """Higher discovery/availability levels rank first among near-tied scores."""
    if not vector:
        return (0, 0)
    m = re.match(r"D(\d+)\s*/\s*A(\d+)", vector)
    if not m:
        return (0, 0)
    return (int(m.group(1)), int(m.group(2)))


def match_query(
    query: MatchQuery,
    capabilities: list[dict[str, Any]],
    *,
    tie_window: float = DEFAULT_TIE_WINDOW,
) -> list[Match]:
    if not query.tokens or not capabilities:
        return []

    scored: list[Match] = []
    for item in capabilities:
        blob = _federated_entry_blob(item)
        tokens = _tokens(blob)
        if not tokens:
            continue
        score = len(query.tokens & tokens) / len(query.tokens | tokens)
        if score <= 0:
            continue
        scored.append(
            Match(
                id=item["id"],
                score=round(score, 4),
                kind="deterministic",
                vector=item.get("vector"),
                owner=item.get("owner"),
                summary=item.get("summary"),
            )
        )

    if not scored:
        return []

    scored.sort(key=lambda m: m.score, reverse=True)
    top_score = scored[0].score
    tied = [m for m in scored if top_score - m.score <= tie_window]
    rest = [m for m in scored if top_score - m.score > tie_window]
    tied.sort(key=lambda m: (_vector_rank(m.vector), m.score), reverse=True)
    return tied + rest


def load_rerank_schema() -> dict[str, Any]:
    return json.loads(RERANK_SCHEMA_PATH.read_text(encoding="utf-8"))


def build_rerank_prompt(query: MatchQuery, matches: list[Match]) -> str:
    candidates = [
        {"id": m.id, "score": m.score, "vector": m.vector, "summary": m.summary}
        for m in matches
    ]
    return (
        "You are reranking capability-registry search candidates for a "
        "query-before-build check. Score how well each candidate semantically "
        "matches the query intent, on a 0-1 confidence scale. You may ONLY "
        "return ids from the candidate list below -- do not invent new ones.\n\n"
        f"Query intent:\n{query.text}\n\n"
        f"Candidates (JSON):\n{json.dumps(candidates, indent=2)}\n\n"
        "Respond with a JSON object: "
        '{"candidates": [{"id": "...", "confidence": 0.0-1.0, "rationale": "..."}]}. '
        "Include every candidate id from the list, in any order."
    )


def request_rerank(
    query: MatchQuery,
    matches: list[Match],
    *,
    llm_url: str | None = None,
) -> list[dict[str, Any]]:
    prompt = build_rerank_prompt(query, matches)
    content = execute_prompt(prompt, base_url=llm_url, config={"temperature": 0.0, "max_tokens": 1500})
    payload = extract_json_object(content)
    validator = Draft202012Validator(load_rerank_schema())
    errors = sorted(validator.iter_errors(payload), key=lambda err: list(err.path))
    if errors:
        messages = "; ".join(error.message for error in errors[:3])
        raise ValueError(f"rerank schema validation failed: {messages}")
    known_ids = {m.id for m in matches}
    candidates = [c for c in payload["candidates"] if c["id"] in known_ids]
    if not candidates:
        raise ValueError("rerank response contained no known candidate ids")
    return candidates


def apply_rerank(matches: list[Match], rerank_result: list[dict[str, Any]]) -> list[Match]:
    """Per spec design principle 3: deterministic matches always rank first.
    LLM rerank results are appended after as separately-labeled (kind='llm')
    entries carrying the semantic confidence score -- never reordering or
    replacing the deterministic base result, so the output is trustworthy
    even when a caller ignores the LLM-kind entries entirely (e.g. --no-llm
    runs never produce them in the first place)."""
    by_id = {m.id: m for m in matches}
    llm_matches = [
        Match(
            id=c["id"],
            score=c["confidence"],
            kind="llm",
            vector=by_id[c["id"]].vector,
            owner=by_id[c["id"]].owner,
            summary=c.get("rationale") or by_id[c["id"]].summary,
        )
        for c in rerank_result
        if c["id"] in by_id
    ]
    llm_matches.sort(key=lambda m: m.score, reverse=True)
    return matches + llm_matches


def verdict_for_score(
    top_score: float,
    *,
    reuse_threshold: float = DEFAULT_REUSE_THRESHOLD,
    extend_threshold: float = DEFAULT_EXTEND_THRESHOLD,
) -> str:
    if top_score >= reuse_threshold:
        return "reuse"
    if top_score >= extend_threshold:
        return "extend"
    return "new"


def run_plan_check(
    query: MatchQuery,
    *,
    reuse_threshold: float = DEFAULT_REUSE_THRESHOLD,
    extend_threshold: float = DEFAULT_EXTEND_THRESHOLD,
    tie_window: float = DEFAULT_TIE_WINDOW,
    top_n: int = 5,
    use_llm: bool = True,
    llm_url: str | None = None,
) -> dict[str, Any]:
    capabilities, updated = load_federated_capabilities()
    matches = match_query(query, capabilities, tie_window=tie_window)
    top_score = matches[0].score if matches else 0.0
    verdict = verdict_for_score(
        top_score, reuse_threshold=reuse_threshold, extend_threshold=extend_threshold
    )

    top_matches = matches[:top_n]
    notes: list[str] = []
    if use_llm and top_matches:
        try:
            rerank_result = request_rerank(query, top_matches, llm_url=llm_url)
            top_matches = apply_rerank(top_matches, rerank_result)
        except ValueError as exc:
            if "LLM backend not configured" in str(exc):
                notes.append("LLM rerank skipped: LLM_CONNECT_URL not set")
            else:
                notes.append(f"LLM rerank skipped: {exc}")

    result = {
        "query": {
            "source": query.source,
            "text": query.text,
            "workplan_id": query.workplan_id,
            "workplan_path": query.workplan_path,
        },
        "verdict": verdict,
        "top_score": top_score,
        "matches": [
            {
                "id": m.id,
                "score": m.score,
                "vector": m.vector,
                "owner": m.owner,
                "summary": m.summary,
                "kind": m.kind,
            }
            for m in top_matches
        ],
        "federated_index_updated": updated,
        "federated_index_stale_warning": _staleness_warning(updated),
    }
    if notes:
        result["notes"] = notes
    return result


def format_plan_check_markdown(result: dict[str, Any]) -> str:
    lines = [f"# Plan check: {result['verdict']}", ""]
    query = result["query"]
    label = query.get("workplan_id") or query["text"][:80]
    lines.append(f"**Query ({query['source']}):** {label}")
    lines.append("")

    matches = result.get("matches", [])
    if matches:
        lines.append("## Matches")
        for m in matches:
            vec = f" ({m['vector']})" if m.get("vector") else ""
            owner = f" — {m['owner']}" if m.get("owner") else ""
            lines.append(f"- `{m['id']}`{vec}{owner} — score {m['score']:.2f} [{m['kind']}]")
            if m.get("summary"):
                lines.append(f"  > {m['summary']}")
    else:
        lines.append("_No matches found in the federated index._")
    lines.append("")

    verdict = result["verdict"]
    if verdict == "reuse":
        lines.append("**Verdict: REUSE** — an existing capability already covers this; link it instead of building new.")
    elif verdict == "extend":
        lines.append("**Verdict: EXTEND** — scope overlaps closely enough that extending the top match is likely cheaper than a new capability.")
    else:
        lines.append("**Verdict: NEW** — no close match in the federated index; proceed.")

    warning = result.get("federated_index_stale_warning")
    if warning:
        lines.append("")
        lines.append(f"⚠ {warning}")

    for note in result.get("notes", []):
        lines.append("")
        lines.append(f"ℹ {note}")

    return "\n".join(lines) + "\n"


def format_plan_check_json(result: dict[str, Any]) -> str:
    return json.dumps(result, indent=2, sort_keys=True)


def maybe_file_capability_request(
    result: dict[str, Any],
    *,
    requesting_domain: str,
    requesting_agent: str,
) -> dict[str, Any] | None:
    """On a 'new' verdict, file a State Hub capability request for the gap.
    Returns None (and does nothing) for any other verdict, or if the hub is
    unreachable -- this never blocks plan-check's primary output."""
    if result["verdict"] != "new":
        return None
    query = result["query"]
    title = (query.get("workplan_id") or query["text"])[:120]
    return file_capability_request(
        title=f"plan-check gap: {title}",
        description=query["text"],
        requesting_domain=requesting_domain,
        requesting_agent=requesting_agent,
        requesting_workplan_id=query.get("workplan_id"),
    )


def record_outcome(
    result: dict[str, Any],
    outcome: str,
    *,
    consumer_repo: str = "reuse-surface",
) -> Path:
    TELEMETRY_PATH.parent.mkdir(parents=True, exist_ok=True)
    top_match = result["matches"][0]["id"] if result.get("matches") else None
    event = {
        "ts": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "consumer_repo": consumer_repo,
        "capability_id": top_match,
        "verdict": result["verdict"],
        "outcome": outcome,
        "source": "plan-check",
    }
    with TELEMETRY_PATH.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(event, sort_keys=True) + "\n")
    return TELEMETRY_PATH