chat_search.py

"""Unified fuzzy search across Codex / Claude Code / Kimi conversation corpora.

Walks three converted-markdown roots and ranks matches by filename slug + body
content. Supports substring and fuzzy modes, AND semantics across tokens, plus
date and CLI filters. Designed to find conversations even when you only
half-remember the topic (e.g. "PDF 知识库").

Usage:
    python chat-search.py "pdf 知识库"
    python chat-search.py --cli kimi 抖音
    python chat-search.py --fuzzy "pf zhshk"
    python chat-search.py --name-only 飞书
    python chat-search.py --after 2026-04-15 --limit 5 ppt

Options:
    --cli {codex,claude,kimi}   restrict to one source
    --name-only                 match only against filename slug
    --content-only              match only against file body
    --fuzzy                     accept loose matches (difflib threshold)
    --limit N                   max results (default 20)
    --snippet-len N             chars of context around each match (default 100)
    --after YYYY-MM-DD          earliest start time
    --before YYYY-MM-DD         latest start time
    -v / --verbose              print scoring details
"""

from __future__ import annotations

import argparse
import difflib
import hashlib
import json
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path

HOME = Path.home()
ROOTS = [
    ("codex", HOME / ".codex" / "sessions_md"),
    ("claude", HOME / ".claude" / "sessions_md"),
    ("kimi", HOME / ".kimi" / "sessions_md"),
]


def kimi_hash_to_cwd() -> dict[str, str]:
    """Reverse-map Kimi project_hash (= md5(cwd)) → cwd via kimi.json work_dirs."""
    out: dict[str, str] = {}
    kj_path = HOME / ".kimi" / "kimi.json"
    if not kj_path.exists():
        return out
    try:
        kj = json.loads(kj_path.read_text(encoding="utf-8"))
    except Exception:
        return out
    for w in kj.get("work_dirs", []) or []:
        path = w.get("path") if isinstance(w, dict) else None
        if path:
            out[hashlib.md5(path.encode("utf-8")).hexdigest()] = path
    return out


_CWD_LINE_RE = re.compile(r"^\|\s*(?:CWD|Source dir)\s*\|\s*`([^`]+)`\s*\|", re.MULTILINE)


def cwd_from_markdown_header(path: Path) -> str:
    """Read just the header table to extract `| CWD |` (codex/claude) or
    `| Source dir |` (kimi). Reads ~3 KB max."""
    try:
        with path.open("r", encoding="utf-8", errors="replace") as fh:
            chunk = fh.read(3072)
    except Exception:
        return ""
    m = _CWD_LINE_RE.search(chunk)
    return m.group(1).strip() if m else ""

# Filenames + directories that are part of the conversion infrastructure, not
# actual conversation markdown.
SKIP_NAMES = {"README.md", "_convert.py"}
SKIP_DIRS = {"images", "outputs", "__pycache__"}

# Filename pattern: <prefix>-<TIMESTAMP>-<first8>--<slug>.md
# prefix is one of: rollout (codex), claude, kimi
_NAME_RE = re.compile(
    r"^(?P<prefix>rollout|claude|kimi)-"
    r"(?P<ts>\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2})"
    r"-(?P<first8>[0-9a-f]{8})"
    r"(?:--(?P<slug>.+))?\.md$"
)


@dataclass
class Hit:
    cli: str
    path: Path
    timestamp: str          # YYYY-MM-DDTHH-MM-SS
    first8: str
    slug: str
    score: float = 0.0
    name_hits: list[str] = field(default_factory=list)
    content_hits: list[tuple[str, int]] = field(default_factory=list)  # (snippet, position)
    cwd: str = ""


def parse_filename(p: Path) -> tuple[str, str, str] | None:
    m = _NAME_RE.match(p.name)
    if not m:
        return None
    return m.group("ts"), m.group("first8"), (m.group("slug") or "")


def iter_md_files(cli_filter: str | None = None):
    for cli, root in ROOTS:
        if cli_filter and cli != cli_filter:
            continue
        if not root.exists():
            continue
        for p in root.rglob("*.md"):
            if p.name in SKIP_NAMES:
                continue
            if any(part in SKIP_DIRS for part in p.relative_to(root).parts):
                continue
            yield cli, p


def _token_in_slug(tok: str, slug_lower: str, fuzzy: bool) -> float:
    """Return per-token slug-match score (0 if no match)."""
    tl = tok.lower()
    if tl in slug_lower:
        return 1.0
    if not fuzzy:
        return 0.0
    ratio = difflib.SequenceMatcher(None, tl, slug_lower).ratio()
    if ratio >= 0.4:
        return ratio
    best = 0.0
    for i in range(0, max(0, len(slug_lower) - len(tl) + 1)):
        window = slug_lower[i:i + len(tl)]
        r = difflib.SequenceMatcher(None, tl, window).ratio()
        if r > best:
            best = r
            if best > 0.85:
                break
    return best if best >= 0.55 else 0.0


def _token_in_text(tok: str, text_lower: str) -> int:
    """Return number of substring hits for tok in text (case-insensitive)."""
    return text_lower.count(tok.lower())


def search_one(
    cli: str,
    path: Path,
    slug: str,
    tokens: list[str],
    *,
    fuzzy: bool,
    name_only: bool,
    content_only: bool,
    snippet_len: int,
) -> tuple[float, list[str], list[tuple[str, int]]]:
    """Score one file. Returns (score, slug-matched tokens, content snippets).

    Per-token AND:  every token must hit somewhere (slug OR content), unless
    --name-only / --content-only restricts the scope. Per-token weight:
        slug substring   = 5.0
        slug fuzzy       = 3.0..5.0 * ratio
        content (capped) = min(count, 5)
    Slug coverage bonus when ALL tokens hit slug: +5.0.
    """
    slug_lower = slug.lower() if slug else ""
    text: str | None = None
    text_lower: str | None = None

    score = 0.0
    name_matched: list[str] = []
    snippets: list[tuple[str, int]] = []

    for tok in tokens:
        slug_w = 0.0
        if slug_lower and not content_only:
            slug_w = _token_in_slug(tok, slug_lower, fuzzy)
        content_n = 0
        if not name_only:
            if text_lower is None:
                try:
                    text = path.read_text(encoding="utf-8", errors="replace")
                    text_lower = text.lower()
                except Exception:
                    text = ""
                    text_lower = ""
            content_n = _token_in_text(tok, text_lower)

        if slug_w == 0.0 and content_n == 0:
            return 0.0, [], []

        if slug_w > 0:
            score += 5.0 * slug_w
            name_matched.append(tok)
        if content_n > 0:
            score += min(content_n, 5) * 1.0

        if content_n > 0 and text is not None:
            tl = tok.lower()
            idx = text_lower.find(tl)
            if idx >= 0:
                start = max(0, idx - snippet_len // 2)
                end = min(len(text), idx + len(tl) + snippet_len // 2)
                snip = text[start:end].replace("\n", " ").strip()
                if start > 0:
                    snip = "…" + snip
                if end < len(text):
                    snip = snip + "…"
                snippets.append((snip, idx))

    # All-slug-tokens bonus (rewards conversations whose TOPIC is the query).
    if name_matched and len(name_matched) == len(tokens):
        score += 5.0

    return score, name_matched, snippets


def filter_by_date(ts: str, after: str | None, before: str | None) -> bool:
    # ts is YYYY-MM-DDTHH-MM-SS; compare lexicographically against YYYY-MM-DD.
    ts_date = ts[:10]
    if after and ts_date < after:
        return False
    if before and ts_date > before:
        return False
    return True


def search(args: argparse.Namespace) -> list[Hit]:
    tokens: list[str] = args.query
    if not tokens:
        return []
    hits: list[Hit] = []
    for cli, p in iter_md_files(args.cli):
        parsed = parse_filename(p)
        if not parsed:
            continue
        ts, first8, slug = parsed
        if not filter_by_date(ts, args.after, args.before):
            continue

        total, n_matched, c_snippets = search_one(
            cli,
            p,
            slug,
            tokens,
            fuzzy=args.fuzzy,
            name_only=args.name_only,
            content_only=args.content_only,
            snippet_len=args.snippet_len,
        )
        if total <= 0:
            continue

        hits.append(
            Hit(
                cli=cli,
                path=p,
                timestamp=ts,
                first8=first8,
                slug=slug,
                score=total,
                name_hits=n_matched,
                content_hits=c_snippets,
            )
        )

    hits.sort(key=lambda h: (-h.score, -ord_ts(h.timestamp)))
    top = hits[: args.limit]

    # Resolve cwd only for the displayed slice — keeps the scan cheap.
    kimi_map = kimi_hash_to_cwd()
    for h in top:
        if h.cli == "kimi":
            try:
                project_hash = h.path.parent.name
            except Exception:
                project_hash = ""
            h.cwd = kimi_map.get(project_hash, "")
        else:
            h.cwd = cwd_from_markdown_header(h.path)
    return top


def ord_ts(ts: str) -> int:
    """Cheap monotonic ordering for timestamp strings — bigger = later."""
    digits = re.sub(r"\D", "", ts)[:14]
    try:
        return int(digits) if digits else 0
    except Exception:
        return 0


CLI_BADGE = {"codex": "🟦 codex", "claude": "🟧 claude", "kimi": "🟩 kimi  "}


def render(hit: Hit, root_for: dict[str, Path]) -> str:
    badge = CLI_BADGE.get(hit.cli, hit.cli)
    rel = hit.path.relative_to(root_for[hit.cli]).as_posix()
    score_str = f"{hit.score:5.1f}"
    header = (
        f"{badge}  {hit.timestamp}  ({score_str})  "
        f"{hit.slug or '(no slug)'}"
    )
    parts = [header, f"   📁 {hit.path}"]
    if hit.cwd:
        parts.append(f"   📂 cwd: {hit.cwd}")
    if hit.name_hits:
        parts.append(f"   🔖 name match: {', '.join(hit.name_hits)}")
    for snip, _ in hit.content_hits[:2]:
        parts.append(f"   ✏️  {snip[:240]}")
    return "\n".join(parts)


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(
        prog="chat-search",
        description="Search across Codex / Claude / Kimi converted conversations.",
    )
    parser.add_argument("query", nargs="+", help="Search tokens (AND semantics).")
    parser.add_argument("--cli", choices=("codex", "claude", "kimi"))
    parser.add_argument("--name-only", action="store_true")
    parser.add_argument("--content-only", action="store_true")
    parser.add_argument("--fuzzy", action="store_true")
    parser.add_argument("--limit", type=int, default=20)
    parser.add_argument("--snippet-len", type=int, default=100)
    parser.add_argument("--after", help="YYYY-MM-DD")
    parser.add_argument("--before", help="YYYY-MM-DD")
    parser.add_argument("-v", "--verbose", action="store_true")
    args = parser.parse_args(argv)

    root_for = {cli: root for cli, root in ROOTS}
    hits = search(args)
    if not hits:
        print("(no matches)")
        return 1

    print(f"Found {len(hits)} match(es) for: {' '.join(args.query)}\n")
    for h in hits:
        print(render(h, root_for))
        print()
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))