reuse-surface/reuse_surface/plan_check.py

from __future__ import annotations

import json
import re
import urllib.error
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 import hub_client
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"
REUSE_EVENT_SCHEMA_PATH = ROOT / "schemas" / "reuse-event.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 load_reuse_event_schema() -> dict[str, Any]:
    return json.loads(REUSE_EVENT_SCHEMA_PATH.read_text(encoding="utf-8"))


def validate_reuse_event(event: dict[str, Any]) -> None:
    validator = Draft202012Validator(load_reuse_event_schema())
    errors = sorted(validator.iter_errors(event), key=lambda err: list(err.path))
    if errors:
        messages = "; ".join(error.message for error in errors[:3])
        raise ValueError(f"reuse event schema validation failed: {messages}")


def append_local_reuse_event(event: dict[str, Any]) -> Path:
    TELEMETRY_PATH.parent.mkdir(parents=True, exist_ok=True)
    with TELEMETRY_PATH.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(event, sort_keys=True) + "\n")
    return TELEMETRY_PATH


def post_or_fallback_reuse_event(
    event: dict[str, Any], *, hub_url: str | None = None
) -> dict[str, Any]:
    """Records one reuse event: the shared schema (schemas/reuse-event.schema.json)
    is the same whether the fact ends up in the hub's reuse_events table or the
    local JSONL fallback. Tries POST /v1/reuse-events first; on any
    misconfiguration or unreachability, falls back to the local file rather
    than losing the fact -- never both, per the design ("+ local JSONL
    fallback when hub unreachable", not a dual-write)."""
    validate_reuse_event(event)

    try:
        status, _payload = hub_client.hub_record_reuse_event(event, hub_url)
    except (ValueError, urllib.error.URLError, TimeoutError, OSError):
        status = None

    if status == 201:
        return {"event": event, "recorded_to": "hub", "path": None}

    path = append_local_reuse_event(event)
    return {"event": event, "recorded_to": "local", "path": str(path)}


def record_outcome(
    result: dict[str, Any],
    outcome: str,
    *,
    consumer_repo: str = "reuse-surface",
    capability_id: str | None = None,
    verdict: str | None = None,
    source: str = "plan-check",
    hub_url: str | None = None,
) -> dict[str, Any]:
    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": capability_id if capability_id is not None else top_match,
        "verdict": verdict if verdict is not None else result["verdict"],
        "outcome": outcome,
        "source": source,
    }
    return post_or_fallback_reuse_event(event, hub_url=hub_url)


def record_manual_reuse_event(
    *,
    consumer_repo: str,
    capability_id: str | None,
    verdict: str,
    outcome: str | None = None,
    hub_url: str | None = None,
) -> dict[str, Any]:
    """Retroactive manual fact recording (`reuse-surface record-reuse`) --
    for reuse decisions made outside plan-check, e.g. discovered after the
    fact during a review."""
    event = {
        "ts": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "consumer_repo": consumer_repo,
        "capability_id": capability_id,
        "verdict": verdict,
        "outcome": outcome,
        "source": "manual",
    }
    return post_or_fallback_reuse_event(event, hub_url=hub_url)