#!/usr/bin/env python3
"""fzf-ai-index: walk every AI CLI's session store and emit one TAB-separated
record per session on stdout.

Output columns (TAB separated, 1-indexed to match fzf's {n}):
    1  agent-raw    machine tag (claude, codex, opencode, droid, pi)
    2  session_id   native session id used to resume
    3  source       absolute path or 'sqlite:<db>' locator
    4  agent-badge  padded + ANSI-colored agent tag (visible)
    5  updated_iso  ISO-8601 UTC timestamp of last activity
    6  msgs         message count (approx, best-effort)
    7  cwd          working directory / project root
    8  title        first real user prompt or stored title, truncated
    9  search_blob  concatenated user-prompt text for fuzzy matching
                    (hidden from view, searched via --nth)

Records are sorted with the most recently updated first so fzf shows them
at the top.
"""

from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor, as_completed
import json as stdjson
import os
import re
import sqlite3
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Iterator, Optional

# Fast JSON via orjson (hard dependency). Accepts bytes to skip a decode step.
import orjson

_JSON_LOADS = orjson.loads
_JSON_DUMPS = orjson.dumps

HOME = Path(os.path.expanduser("~"))
HOME_STR = str(HOME)
CACHE_DIR = Path(os.environ.get("XDG_CACHE_HOME", str(HOME / ".cache"))) / "fzf-ai"
INDEX_CACHE_VERSION = 1
INDEX_JOBS = max(
    1,
    int(os.environ.get("FZFAI_INDEX_JOBS") or min(32, (os.cpu_count() or 4) + 4)),
)
_RE_WS = re.compile(r"\s+")
_RE_WS_TAB_NL = re.compile(r"[\t\n]+")

MAX_TITLE = 120
DISPLAY_CWD = 28
DISPLAY_TITLE = 72
# Keep the search blob small so that fuzzy matching stays high-signal.
# Long soups of text are easy to accidentally match against, which floods
# the picker with irrelevant rows whenever the query shares common letters
# with filler content. 1200 chars is enough to hold ~10 distinctive prompt
# openings per session.
MAX_SEARCH_BLOB = 4000
MAX_SEARCH_CHUNK = 140  # first part of each message only
MAX_PROMPTS = 30       # stop collecting search text after this many chunks

# Colors for the agent tag in the main list (fzf with --ansi renders these).
_AGENT_COLOR = {
    "claude":   "38;5;208",  # orange
    "codex":    "38;5;39",   # cyan
    "opencode": "38;5;141",  # purple
    "droid":    "38;5;76",   # green
    "pi":       "38;5;220",  # yellow
}

NO_COLOR = bool(os.environ.get("NO_COLOR"))
_SEARCH_NOISE_PREFIXES = (
    "<environment_context>",
    "<permissions instructions>",
)


def colorise(agent: str) -> str:
    if NO_COLOR:
        return f"{agent:<8}"
    code = _AGENT_COLOR.get(agent, "1")
    return f"\033[{code}m{agent:<8}\033[0m"


# --------------------------------------------------------------------------
# Index cache: avoid re-parsing unchanged session files on every startup.
# Cache is keyed by source path (or "sqlite:<db>" for opencode).
# Each entry stores the file's mtime + size for validation.
# --------------------------------------------------------------------------
def _get_index_cache_path() -> Path:
    return CACHE_DIR / "index.json"


def _load_index_cache() -> dict:
    path = _get_index_cache_path()
    if not path.is_file():
        return {"v": INDEX_CACHE_VERSION, "sources": {}}
    try:
        data = _JSON_LOADS(path.read_bytes())
        if isinstance(data, dict) and data.get("v") == INDEX_CACHE_VERSION:
            return data
    except Exception:
        pass
    return {"v": INDEX_CACHE_VERSION, "sources": {}}


def _save_index_cache(cache: dict) -> None:
    path = _get_index_cache_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    content = _JSON_DUMPS(cache)
    if isinstance(content, bytes):
        tmp.write_bytes(content)
    else:
        tmp.write_text(content, encoding="utf-8")
    tmp.rename(path)


def _validate_cache_entry(entry: dict) -> bool:
    """Check if a cache entry is still valid (file exists, mtime+size unchanged)."""
    src = entry.get("source", "")
    if src.startswith("sqlite:"):
        db_path = Path(src[len("sqlite:"):])
        return db_path.is_file() and db_path.stat().st_mtime == entry.get("mtime", 0)
    p = Path(src)
    if not p.is_file():
        return False
    try:
        st = p.stat()
        return st.st_mtime == entry.get("mtime", 0) and st.st_size == entry.get("size", 0)
    except OSError:
        return False


def _file_cache_entry(path: Path, records: list[Record]) -> dict:
    """Build a cache entry from a file path and its parsed records."""
    try:
        st = path.stat()
        mtime, size = st.st_mtime, st.st_size
    except OSError:
        mtime, size = 0.0, 0
    return {
        "mtime": mtime,
        "size": size,
        "source": str(path),
        "records": [
            {
                "agent": r.agent,
                "session_id": r.session_id,
                "source": r.source,
                "updated": r.updated,
                "msgs": r.msgs,
                "cwd": r.cwd,
                "title": r.title,
                "prompts": r.prompts,
            }
            for r in records
        ],
    }


def _db_cache_entry(db: Path, records: list[Record]) -> dict:
    """Build a cache entry from a database path and its parsed records."""
    try:
        mtime = db.stat().st_mtime
    except OSError:
        mtime = 0.0
    return {
        "mtime": mtime,
        "size": 0,
        "source": f"sqlite:{db}",
        "records": [
            {
                "agent": r.agent,
                "session_id": r.session_id,
                "source": r.source,
                "updated": r.updated,
                "msgs": r.msgs,
                "cwd": r.cwd,
                "title": r.title,
                "prompts": r.prompts,
            }
            for r in records
        ],
    }


def _records_from_cache_entry(entry: dict) -> list[Record]:
    """Reconstruct Record objects from a cache entry."""
    records = []
    for d in entry.get("records", []):
        r = Record(
            agent=d.get("agent", ""),
            session_id=d.get("session_id", ""),
            source=d.get("source", ""),
            updated=d.get("updated", 0.0),
            msgs=d.get("msgs", 0),
            cwd=d.get("cwd", ""),
            title=d.get("title", ""),
            prompts=d.get("prompts", []),
        )
        records.append(r)
    return records


def _collect_jsonl_paths(root: Path, glob_pat: str = "*.jsonl") -> list[Path]:
    """Collect all JSONL paths from a root directory (non-recursive by default)."""
    if not root.is_dir():
        return []
    paths: list[Path] = []
    for proj in root.iterdir():
        if not proj.is_dir():
            continue
        paths.extend(proj.glob(glob_pat))
    return paths


@dataclass
class Record:
    agent: str
    session_id: str
    source: str
    updated: float = 0.0  # unix seconds
    msgs: int = 0
    cwd: str = ""
    title: str = ""
    prompts: list = field(default_factory=list)
    _prompt_set: set = field(default_factory=set)

    def add_search_text(self, text: str) -> None:
        # Clip before cleaning so the regex sub only walks ~140 chars
        # instead of potentially megabytes of assistant output.
        t = clean(text[: MAX_SEARCH_CHUNK * 2])[:MAX_SEARCH_CHUNK]
        if not t or t in self._prompt_set:
            return
        for prefix in _SEARCH_NOISE_PREFIXES:
            if t.startswith(prefix):
                return
        self._prompt_set.add(t)
        self.prompts.append(t)

    def add_prompt(self, text: str) -> None:
        if not is_real_prompt(text):
            return
        self.add_search_text(text)

    def as_row(self) -> str:
        iso = (
            datetime.fromtimestamp(self.updated, tz=timezone.utc)
            .strftime("%Y-%m-%d %H:%M")
            if self.updated
            else ""
        )
        title = display_text(clean(self.title)[:MAX_TITLE] or "(no title)", DISPLAY_TITLE)
        cwd = self.cwd or ""
        if cwd.startswith(HOME_STR):
            cwd = "~" + cwd[len(HOME_STR):]
        cwd = display_text(cwd, DISPLAY_CWD)
        blob = " · ".join(self.prompts)
        blob = _RE_WS_TAB_NL.sub(" ", blob)[:MAX_SEARCH_BLOB]
        return "\t".join(
            [
                self.agent,
                self.session_id,
                self.source,
                colorise(self.agent),
                f"{iso:<16}",
                f"{self.msgs:>4}",
                f"{cwd:<{DISPLAY_CWD}}",
                title,
                blob,
            ]
        )


def clean(s: str) -> str:
    if not s:
        return ""
    # Fast path: most prompt chunks are plain ASCII with no whitespace
    # runs. Skipping the regex for these rows cuts ~30% off the
    # indexer's re.sub time on a typical run. We require the string to
    # be ASCII so non-ASCII whitespace (U+00A0, U+202F, …) still flows
    # through the regex path where \s matches them.
    if s.isascii() and "\t" not in s and "\n" not in s and "\r" not in s and "  " not in s:
        return s.strip()
    return _RE_WS.sub(" ", s).strip()


def display_text(text: str, width: int) -> str:
    if len(text) <= width:
        return text
    if width <= 1:
        return text[:width]
    return text[: width - 1] + "…"


def iter_text_fragments(parts) -> Iterator[str]:
    if parts is None:
        return
    if isinstance(parts, str):
        yield parts
        return
    if not isinstance(parts, list):
        return
    for part in parts:
        if not isinstance(part, dict):
            continue
        kind = part.get("type")
        if kind in ("text", "input_text", "output_text"):
            text = part.get("text")
            if text is None:
                text = (part.get("data") or {}).get("text")
            if text:
                yield text


_JUNK_PREFIXES = (
    "<",  # xml-ish tags, <command-name>, <system-reminder>, <env...>
    "Caveat:",
    "[Request interrupted",
)
_RE_JUNK = re.compile(r'^\s*(?:<|Caveat:|\[Request interrupted)')


def is_real_prompt(text: str) -> bool:
    if not text:
        return False
    return not _RE_JUNK.match(text)


_iso_cache: dict[str, float] = {}

def iso_to_epoch(s: str) -> float:
    if not s:
        return 0.0
    v = _iso_cache.get(s)
    if v is not None:
        return v
    try:
        s2 = s.replace("Z", "+00:00")
        v = datetime.fromisoformat(s2).timestamp()
    except Exception:
        v = 0.0
    _iso_cache[s] = v
    return v


def _parallel_read(
    paths: list[Path],
    reader: Callable[[Path], Optional[Record]],
    label: str,
) -> Iterator[Record]:
    if not paths:
        return
    if len(paths) == 1 or INDEX_JOBS == 1:
        for path in paths:
            try:
                rec = reader(path)
                if rec:
                    yield rec
            except Exception as e:
                print(f"{label}: {path}: {e}", file=sys.stderr)
        return

    with ProcessPoolExecutor(max_workers=min(INDEX_JOBS, len(paths))) as pool:
        futures = {pool.submit(reader, path): path for path in paths}
        for future in as_completed(futures):
            path = futures[future]
            try:
                rec = future.result()
                if rec:
                    yield rec
            except Exception as e:
                print(f"{label}: {path}: {e}", file=sys.stderr)


# --------------------------------------------------------------------------
# Claude Code  -- ~/.claude/projects/<enc-cwd>/<uuid>.jsonl
# --------------------------------------------------------------------------
def walk_claude(cache: dict = None) -> Iterator[Record]:
    root = HOME / ".claude" / "projects"
    sources = cache.get("sources", {}) if cache else {}
    uncached_paths: list[Path] = []
    for path in _collect_jsonl_paths(root):
        entry = sources.get(str(path))
        if entry and _validate_cache_entry(entry):
            yield from _records_from_cache_entry(entry)
        else:
            uncached_paths.append(path)
    for rec in _parallel_read(uncached_paths, _read_claude_file, "claude"):
        sources[str(rec.source)] = _file_cache_entry(Path(rec.source), [rec])
        yield rec


def _read_claude_file(path: Path) -> Optional[Record]:
    sid = path.stem
    rec = Record(
        agent="claude",
        session_id=sid,
        source=str(path),
        updated=path.stat().st_mtime,
    )
    title = ""
    cwd = ""
    msgs = 0
    skip_search = False
    # Claude jsonl files are dominated by 'progress' / 'file-history-snapshot'
    # noise. Fast-skip any line whose type we don't care about before
    # hitting json.loads (the profiler's #1 hotspot). We also accept any
    # line carrying "cwd" so that sessions with no user/assistant turns
    # (e.g. a summary-only file) still surface their working directory.
    with path.open("r", encoding="utf-8", errors="replace", buffering=1 << 16) as f:
        for line in f:
            has_msg = '"type":"user"' in line or '"type":"assistant"' in line
            has_cwd = not cwd and '"cwd":' in line
            if not (has_msg or has_cwd):
                continue
            try:
                obj = _JSON_LOADS(line)
            except stdjson.JSONDecodeError:
                continue
            if not cwd:
                c = obj.get("cwd")
                if c:
                    cwd = c
            t = obj.get("type")
            if t not in ("user", "assistant"):
                continue
            msgs += 1
            if not skip_search and not obj.get("isMeta"):
                m = obj.get("message") or {}
                is_user = t == "user"
                for tx in iter_text_fragments(m.get("content")):
                    if is_user and is_real_prompt(tx):
                        if not title:
                            title = tx
                        rec.add_prompt(tx)
                    else:
                        rec.add_search_text(tx)
                if len(rec.prompts) >= MAX_PROMPTS:
                    skip_search = True
            ts = obj.get("timestamp")
            if ts:
                e = iso_to_epoch(ts)
                if e > rec.updated:
                    rec.updated = e
    rec.title = title
    rec.cwd = cwd
    rec.msgs = msgs
    return rec


# --------------------------------------------------------------------------
# Codex  -- ~/.codex/sessions/YYYY/MM/DD/rollout-<ts>-<uuid>.jsonl
# --------------------------------------------------------------------------
def walk_codex(cache: dict = None) -> Iterator[Record]:
    root = HOME / ".codex" / "sessions"
    sources = cache.get("sources", {}) if cache else {}
    uncached_paths: list[Path] = []
    for path in root.rglob("rollout-*.jsonl"):
        entry = sources.get(str(path))
        if entry and _validate_cache_entry(entry):
            yield from _records_from_cache_entry(entry)
        else:
            uncached_paths.append(path)
    for rec in _parallel_read(uncached_paths, _read_codex_file, "codex"):
        sources[str(rec.source)] = _file_cache_entry(Path(rec.source), [rec])
        yield rec


def _read_codex_file(path: Path) -> Optional[Record]:
    m = re.match(
        r"rollout-\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}-([0-9a-f-]+)\.jsonl$",
        path.name,
    )
    sid = m.group(1) if m else path.stem
    rec = Record(
        agent="codex",
        session_id=sid,
        source=str(path),
        updated=path.stat().st_mtime,
    )
    title = ""
    cwd = ""
    msgs = 0
    skip_search = False
    # Codex rollouts are noisy: on a ~900-line session the only lines we
    # read are ~83 'message' items plus the one 'session_meta'. The rest
    # (reasoning / function_call / event_msg) can be substring-filtered
    # without decoding JSON, which is ~5x faster and cuts this reader
    # from 0.23s to well under 0.05s on a 900-line file.
    with path.open("r", encoding="utf-8", errors="replace", buffering=1 << 16) as f:
        for line in f:
            is_message = '"type":"message"' in line
            is_meta = not is_message and '"type":"session_meta"' in line
            if not (is_message or is_meta):
                continue
            try:
                obj = _JSON_LOADS(line)
            except stdjson.JSONDecodeError:
                continue
            t = obj.get("type")
            payload = obj.get("payload") or {}
            if t == "session_meta":
                c = payload.get("cwd")
                if c:
                    cwd = c
                pid = payload.get("id")
                if pid:
                    rec.session_id = pid
                continue
            if t != "response_item" or payload.get("type") != "message":
                continue
            role = payload.get("role")
            if role not in ("user", "assistant"):
                continue
            msgs += 1
            if not skip_search:
                is_user = role == "user"
                for tx in iter_text_fragments(payload.get("content") or ()):
                    if is_user and is_real_prompt(tx):
                        if not title:
                            title = tx
                        rec.add_prompt(tx)
                    else:
                        rec.add_search_text(tx)
                if len(rec.prompts) >= MAX_PROMPTS:
                    skip_search = True
            ts = obj.get("timestamp")
            if ts:
                e = iso_to_epoch(ts)
                if e > rec.updated:
                    rec.updated = e
    rec.title = title
    rec.cwd = cwd
    rec.msgs = msgs
    return rec


# --------------------------------------------------------------------------
# opencode -- sqlite at ~/.opencode/opencode.db
# --------------------------------------------------------------------------
# opencode has two schema generations in the wild plus per-project databases:
#   * current:  ~/.local/share/opencode/opencode.db
#               tables: session(id, directory, title, time_updated, ...),
#                       message(id, session_id, time_created, data JSON),
#                       part(message_id, session_id, data JSON, ...)
#   * legacy:   ~/.opencode/opencode.db
#               tables: sessions, messages, files  (flat)
#   * per-project: <repo>/.opencode/opencode.db  (same legacy layout)
# --------------------------------------------------------------------------
def _opencode_db_paths() -> list[Path]:
    """Return the opencode databases that actually back live sessions.

    Opencode >=1.x consolidated all sessions into a single global database
    at ``~/.local/share/opencode/opencode.db``. The legacy
    ``~/.opencode/opencode.db`` and any per-project ``<repo>/.opencode/...``
    databases are leftovers from older versions whose sessions are not
    resumable by the current CLI, so we deliberately skip them to avoid
    surfacing ghost sessions that fail on resume.
    """
    seen: dict[Path, float] = {}
    candidate = HOME / ".local" / "share" / "opencode" / "opencode.db"
    if candidate.is_file():
        seen[candidate.resolve()] = candidate.stat().st_mtime
    # Also include the legacy db only if the current one doesn't exist at
    # all (fresh migration case). Skipping it otherwise prevents duplicate
    # or unresumable entries from cluttering the picker.
    if not seen:
        legacy = HOME / ".opencode" / "opencode.db"
        if legacy.is_file():
            seen[legacy.resolve()] = legacy.stat().st_mtime
    return sorted(seen, key=lambda p: -seen[p])


def _opencode_table_set(con: sqlite3.Connection) -> set[str]:
    cur = con.execute("SELECT name FROM sqlite_master WHERE type='table'")
    return {r[0] for r in cur.fetchall()}


def _opencode_walk_current(con: sqlite3.Connection, db: Path) -> Iterator[Record]:
    cur = con.execute(
        "SELECT id, directory, title, time_updated FROM session ORDER BY time_updated DESC"
    )
    sessions = cur.fetchall()
    if not sessions:
        return

    msg_count: dict[str, int] = {}
    for sid, c in con.execute(
        "SELECT session_id, COUNT(*) FROM message "
        "WHERE data LIKE '%' || char(34) || 'role' || char(34) || ':' || char(34) || '%' "
        "GROUP BY session_id"
    ):
        msg_count[sid] = c

    # Single JOIN query replaces 3 separate queries (msg_role lookup + parts scan).
    # json_extract on the JOIN is fast enough when the LIKE pre-filter skips
    # non-text rows on part.data first.
    search_text: dict[str, list[tuple[str, str]]] = {}
    for sid, role, text in con.execute(
        """
        SELECT p.session_id,
               json_extract(m.data, '$.role'),
               json_extract(p.data, '$.text')
        FROM part p
        JOIN message m ON m.id = p.message_id AND m.session_id = p.session_id
        WHERE p.data LIKE '{"type":"text"%'
          AND p.data NOT LIKE '%"synthetic":true%'
        ORDER BY p.time_created ASC
        """
    ):
        if not text or role not in ("user", "assistant"):
            continue
        search_text.setdefault(sid, []).append((role, text))

    for sid, directory, title, updated in sessions:
        rec = Record(
            agent="opencode",
            session_id=sid,
            source=f"sqlite:{db}",
            updated=(updated or 0) / 1000.0,
            msgs=msg_count.get(sid, 0),
            cwd=directory or "",
        )
        for role, text in search_text.get(sid, []):
            if role == "user" and is_real_prompt(text):
                if not rec.title:
                    rec.title = text
                rec.add_prompt(text)
            else:
                rec.add_search_text(text)
        if not rec.title:
            rec.title = title or ""
        yield rec


def _opencode_walk_legacy(con: sqlite3.Connection, db: Path) -> Iterator[Record]:
    """Legacy schema: sessions / messages / files."""
    cur = con.execute(
        "SELECT id, title, message_count, updated_at FROM sessions ORDER BY updated_at DESC"
    )
    rows = cur.fetchall()
    if not rows:
        return
    search_text: dict[str, list[tuple[str, str]]] = {}
    for sid, role, parts in con.execute(
        "SELECT session_id, role, parts FROM messages ORDER BY created_at ASC"
    ):
        try:
            p = stdjson.loads(parts)
        except Exception:
            continue
        if not isinstance(p, list):
            continue
        for part in p:
            if not isinstance(part, dict) or part.get("type") != "text":
                continue
            tx = part.get("text") or (
                (part.get("data") or {}).get("text")
            ) or ""
            if tx:
                search_text.setdefault(sid, []).append((role, tx))
    # cwd via files table (best-effort)
    cwd_map: dict[str, str] = {}
    try:
        for sid, pth in con.execute(
            "SELECT session_id, path FROM files GROUP BY session_id"
        ):
            if pth and pth.startswith("/"):
                cwd_map[sid] = os.path.dirname(pth)
    except sqlite3.Error:
        pass
    for sid, title, msgs, updated in rows:
        updated_s = (
            (updated or 0) / 1000.0
            if (updated or 0) > 1_000_000_000_000
            else (updated or 0)
        )
        rec = Record(
            agent="opencode",
            session_id=sid,
            source=f"sqlite:{db}",
            updated=updated_s,
            msgs=msgs or 0,
            cwd=cwd_map.get(sid, ""),
            title="",
        )
        for role, text in search_text.get(sid, []):
            if role == "user" and is_real_prompt(text):
                if not rec.title:
                    rec.title = text
                rec.add_prompt(text)
            else:
                rec.add_search_text(text)
        if not rec.title:
            rec.title = title or ""
        yield rec


def walk_opencode(cache: dict = None) -> Iterator[Record]:
    sources = cache.get("sources", {}) if cache else {}
    for db in _opencode_db_paths():
        src_key = f"sqlite:{db}"
        entry = sources.get(src_key)
        if entry and _validate_cache_entry(entry):
            yield from _records_from_cache_entry(entry)
            continue
        records: list[Record] = []
        seen_sids: set[str] = set()
        try:
            con = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
        except sqlite3.Error as e:
            print(f"opencode: cannot open {db}: {e}", file=sys.stderr)
            continue
        # Performance PRAGMAs are best-effort on read-only connections.
        # Some systems (e.g. with FUSE temp dirs) reject these silently.
        try:
            con.execute("PRAGMA journal_mode=WAL")
            con.execute("PRAGMA cache_size=-64000")
        except sqlite3.Error:
            pass
        try:
            tables = _opencode_table_set(con)
            if {"session", "message", "part"}.issubset(tables):
                walker = _opencode_walk_current(con, db)
            elif {"sessions", "messages"}.issubset(tables):
                walker = _opencode_walk_legacy(con, db)
            else:
                continue
            for rec in walker:
                if rec.session_id in seen_sids:
                    continue
                seen_sids.add(rec.session_id)
                records.append(rec)
                yield rec
        except sqlite3.Error as e:
            print(f"opencode: {db}: {e}", file=sys.stderr)
        finally:
            con.close()
        if records:
            sources[src_key] = _db_cache_entry(db, records)


# --------------------------------------------------------------------------
# Droid / Factory -- ~/.factory/sessions/<proj>/<uuid>.jsonl
# --------------------------------------------------------------------------
def walk_droid(cache: dict = None) -> Iterator[Record]:
    root = HOME / ".factory" / "sessions"
    sources = cache.get("sources", {}) if cache else {}
    uncached_paths: list[Path] = []
    for path in _collect_jsonl_paths(root):
        entry = sources.get(str(path))
        if entry and _validate_cache_entry(entry):
            yield from _records_from_cache_entry(entry)
        else:
            uncached_paths.append(path)
    for rec in _parallel_read(uncached_paths, _read_droid_file, "droid"):
        sources[str(rec.source)] = _file_cache_entry(Path(rec.source), [rec])
        yield rec


def _read_droid_file(path: Path) -> Optional[Record]:
    sid = path.stem
    rec = Record(
        agent="droid",
        session_id=sid,
        source=str(path),
        updated=path.stat().st_mtime,
    )
    title = ""
    cwd = ""
    msgs = 0
    skip_search = False
    with path.open("r", encoding="utf-8", errors="replace", buffering=1 << 16) as f:
        for line in f:
            is_message = '"type":"message"' in line
            is_start = not is_message and '"type":"session_start"' in line
            if not (is_message or is_start):
                continue
            try:
                obj = _JSON_LOADS(line)
            except stdjson.JSONDecodeError:
                continue
            t = obj.get("type")
            if t == "session_start":
                c = obj.get("cwd")
                if c:
                    cwd = c
                tt = obj.get("title")
                if tt and not title:
                    title = tt
                continue
            if t != "message":
                continue
            m = obj.get("message") or {}
            role = m.get("role")
            if role not in ("user", "assistant"):
                continue
            msgs += 1
            if not skip_search:
                is_user = role == "user"
                for tx in iter_text_fragments(m.get("content") or ()):
                    if is_user and is_real_prompt(tx):
                        if not title:
                            title = tx
                        rec.add_prompt(tx)
                    else:
                        rec.add_search_text(tx)
                if len(rec.prompts) >= MAX_PROMPTS:
                    skip_search = True
            ts = obj.get("timestamp")
            if ts:
                e = iso_to_epoch(ts)
                if e > rec.updated:
                    rec.updated = e
    rec.title = title
    rec.cwd = cwd
    rec.msgs = msgs
    return rec


# --------------------------------------------------------------------------
# pi  -- ~/.pi/agent/sessions/<enc-cwd>/<iso>_<uuid>.jsonl
# --------------------------------------------------------------------------
def walk_pi(cache: dict = None) -> Iterator[Record]:
    root = HOME / ".pi" / "agent" / "sessions"
    sources = cache.get("sources", {}) if cache else {}
    uncached_paths: list[Path] = []
    for path in _collect_jsonl_paths(root):
        entry = sources.get(str(path))
        if entry and _validate_cache_entry(entry):
            yield from _records_from_cache_entry(entry)
        else:
            uncached_paths.append(path)
    for rec in _parallel_read(uncached_paths, _read_pi_file, "pi"):
        sources[str(rec.source)] = _file_cache_entry(Path(rec.source), [rec])
        yield rec


def _read_pi_file(path: Path) -> Optional[Record]:
    m = re.match(
        r".*_([0-9a-f-]+)$",
        path.stem,
    )
    sid = m.group(1) if m else path.stem
    rec = Record(
        agent="pi",
        session_id=sid,
        source=str(path),
        updated=path.stat().st_mtime,
    )
    title = ""
    cwd = ""
    msgs = 0
    skip_search = False
    with path.open("r", encoding="utf-8", errors="replace", buffering=1 << 16) as f:
        for line in f:
            is_message = '"type":"message"' in line
            is_session = not is_message and '"type":"session"' in line
            if not (is_message or is_session):
                continue
            try:
                obj = _JSON_LOADS(line)
            except stdjson.JSONDecodeError:
                continue
            t = obj.get("type")
            if t == "session":
                c = obj.get("cwd")
                if c:
                    cwd = c
                oid = obj.get("id")
                if oid:
                    rec.session_id = oid
                continue
            if t != "message":
                continue
            m = obj.get("message") or {}
            role = m.get("role")
            if role not in ("user", "assistant"):
                continue
            msgs += 1
            if not skip_search:
                is_user = role == "user"
                for tx in iter_text_fragments(m.get("content")):
                    if is_user and is_real_prompt(tx):
                        if not title:
                            title = tx
                        rec.add_prompt(tx)
                    else:
                        rec.add_search_text(tx)
                if len(rec.prompts) >= MAX_PROMPTS:
                    skip_search = True
            ts = obj.get("timestamp")
            if ts:
                e = iso_to_epoch(ts)
                if e > rec.updated:
                    rec.updated = e
    rec.title = title
    rec.cwd = cwd
    rec.msgs = msgs
    return rec


# --------------------------------------------------------------------------
WALKERS = {
    "claude": walk_claude,
    "codex": walk_codex,
    "opencode": walk_opencode,
    "droid": walk_droid,
    "pi": walk_pi,
}


def _resolve_walker(agent: str):
    """Resolve a walker function for an agent name.

    Checks built-in WALKERS first, then discovers plugins from
    the ``stores/`` directory.
    """
    w = WALKERS.get(agent)
    if w is not None:
        return w
    # Try plugin discovery
    try:
        from stores import discover_stores
        stores_dir = Path(__file__).resolve().parent / "stores"
        plugins = discover_stores(stores_dir)
        return plugins.get(agent)
    except (ImportError, Exception):
        pass
    return None


def main() -> int:
    import signal
    signal.signal(signal.SIGPIPE, signal.SIG_DFL)
    cache = _load_index_cache()
    agents = sys.argv[1:] or list(WALKERS.keys())
    records: list[Record] = []
    if len(agents) > 1 and INDEX_JOBS > 1:
        from concurrent.futures import ThreadPoolExecutor
        with ThreadPoolExecutor(max_workers=len(agents)) as pool:
            future_map = {}
            for a in agents:
                walker = _resolve_walker(a)
                if not walker:
                    print(f"unknown agent: {a}", file=sys.stderr)
                    continue
                future_map[pool.submit(lambda w: list(w(cache)), walker)] = a
            for future in as_completed(future_map):
                a = future_map[future]
                try:
                    records.extend(future.result())
                except Exception as e:
                    print(f"{a}: {e}", file=sys.stderr)
    else:
        for a in agents:
            walker = _resolve_walker(a)
            if not walker:
                print(f"unknown agent: {a}", file=sys.stderr)
                continue
            try:
                records.extend(walker(cache))
            except Exception as e:
                print(f"{a}: {e}", file=sys.stderr)
    try:
        _save_index_cache(cache)
    except Exception:
        pass
    records.sort(key=lambda r: r.updated, reverse=True)
    # One big write is a touch faster and, more importantly, atomic: fzf
    # won't see partially-written rows if the caller redirects to a pipe.
    sys.stdout.write("\n".join(r.as_row() for r in records))
    if records:
        sys.stdout.write("\n")
    return 0


if __name__ == "__main__":
    sys.exit(main())
