doc_index.py

import json
import re
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path

from aider.dap_paths import legacy_project_state_dir, project_state_dir


@dataclass(frozen=True)
class DocSearchResult:
    rel_path: str
    score: int
    line: int | None
    snippet: str


class DocIndex:
    VERSION = 1

    def __init__(self, root, *, encoding="utf-8"):
        self.root = Path(root)
        self.encoding = encoding
        self.preferred_store_dir = project_state_dir(self.root)
        self.legacy_store_dir = legacy_project_state_dir(self.root)
        preferred = self.preferred_store_dir / "doc_index.json"
        legacy = self.legacy_store_dir / "doc_index.json"
        self.index_path = preferred if preferred.exists() else (legacy if legacy.exists() else preferred)

    def rebuild(self):
        docs = {}
        for rel_path in self._iter_doc_rel_paths():
            abs_path = self.root / rel_path
            rec = self._read_doc(abs_path)
            if rec is not None:
                docs[rel_path] = rec

        payload = {
            "version": self.VERSION,
            "built_at": int(time.time()),
            "docs": docs,
        }
        self._write_index(payload)
        return len(docs)

    def update(self):
        payload = self._load_index_payload()
        existing = payload.get("docs", {}) if payload else {}

        current_paths = set(self._iter_doc_rel_paths())
        next_docs = {}

        for rel_path in sorted(current_paths):
            abs_path = self.root / rel_path
            try:
                stat = abs_path.stat()
            except OSError:
                continue

            prev = existing.get(rel_path)
            if prev and prev.get("mtime_ns") == stat.st_mtime_ns and prev.get("size") == stat.st_size:
                next_docs[rel_path] = prev
                continue

            rec = self._read_doc(abs_path)
            if rec is not None:
                next_docs[rel_path] = rec

        payload = {
            "version": self.VERSION,
            "built_at": int(time.time()),
            "docs": next_docs,
        }
        self._write_index(payload)
        return {"indexed": len(next_docs), "removed": len(set(existing) - current_paths)}

    def status(self):
        payload = self._load_index_payload()
        if not payload:
            return {"exists": False, "path": str(self.index_path)}
        return {
            "exists": True,
            "path": str(self.index_path),
            "built_at": payload.get("built_at"),
            "docs": len(payload.get("docs", {}) or {}),
        }

    def search(self, query, *, topk=5, context_lines=2):
        query = (query or "").strip()
        if not query:
            return []

        payload = self._load_index_payload()
        if not payload:
            self.rebuild()
            payload = self._load_index_payload() or {}

        docs = payload.get("docs", {}) or {}
        q_tokens = [t for t in re.split(r"\s+", query.casefold()) if t]
        if not q_tokens:
            return []

        scored = []
        for rel_path, rec in docs.items():
            text = (rec.get("text") or "").casefold()
            if not text:
                continue

            score = 0
            for tok in q_tokens:
                score += text.count(tok)
                score += 2 * rel_path.casefold().count(tok)
            score += 2 * text.count(query.casefold())

            if score <= 0:
                continue

            line, snippet = self._make_snippet(rec.get("text") or "", query, q_tokens, context_lines)
            scored.append(DocSearchResult(rel_path=rel_path, score=score, line=line, snippet=snippet))

        scored.sort(key=lambda r: (-r.score, r.rel_path))
        return scored[: max(1, int(topk))]

    def _iter_doc_rel_paths(self):
        paths = []

        overview = self.root / "PROJECT_OVERVIEW.md"
        if overview.is_file():
            paths.append("PROJECT_OVERVIEW.md")

        for abs_path in self.root.rglob("AI_ARCH.md"):
            try:
                rel = abs_path.relative_to(self.root).as_posix()
            except ValueError:
                continue
            if rel.startswith("DAP/") or rel.startswith(".aider/") or rel.startswith(".git/"):
                continue
            paths.append(rel)

        return sorted(set(paths))

    def _read_doc(self, abs_path):
        try:
            text = abs_path.read_text(encoding=self.encoding, errors="replace")
            stat = abs_path.stat()
        except OSError:
            return None

        return {"mtime_ns": stat.st_mtime_ns, "size": stat.st_size, "text": text}

    def _load_index_payload(self):
        try:
            raw = self.index_path.read_text(encoding="utf-8")
        except OSError:
            return None

        try:
            payload = json.loads(raw)
        except Exception:
            return None

        if payload.get("version") != self.VERSION:
            return None

        if not isinstance(payload.get("docs"), dict):
            return None

        return payload

    def _write_index(self, payload):
        self.preferred_store_dir.mkdir(parents=True, exist_ok=True)
        tmp_path = None
        try:
            with tempfile.NamedTemporaryFile(
                "w",
                encoding="utf-8",
                delete=False,
                dir=str(self.preferred_store_dir),
                prefix="doc_index_",
                suffix=".json",
            ) as f:
                json.dump(payload, f, ensure_ascii=False)
                f.write("\n")
                tmp_path = Path(f.name)
            tmp_path.replace(self.preferred_store_dir / "doc_index.json")
            self.index_path = self.preferred_store_dir / "doc_index.json"
        finally:
            if tmp_path and tmp_path.exists() and tmp_path != self.index_path:
                try:
                    tmp_path.unlink()
                except OSError:
                    pass

    def _make_snippet(self, text, query, q_tokens, context_lines):
        lines = text.splitlines()
        if not lines:
            return None, ""

        q_cf = query.casefold()
        first_line = None
        for i, line in enumerate(lines):
            line_cf = line.casefold()
            if q_cf in line_cf:
                first_line = i
                break
            if any(tok in line_cf for tok in q_tokens):
                first_line = i
                break

        if first_line is None:
            first_line = 0

        start = max(0, first_line - int(context_lines))
        end = min(len(lines), first_line + int(context_lines) + 1)
        snippet = "\n".join(lines[start:end]).rstrip()
        return first_line + 1, snippet