#!/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 ThreadPoolExecutor, as_completed
import json
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

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]+")

# Pre-filter substrings used to skip JSON parsing on lines that are
# guaranteed to be irrelevant. JSON emitted by every agent puts
# "type":"<value>" as a compact, unquoted key so a plain substring
# test is both correct and ~5x faster than json.loads + dict lookup.
# Any line whose "type" isn't in the agent's whitelist is discarded
# before we even hit the JSON decoder.
try:
    import orjson
    _JSON_LOADS = orjson.loads
    _JSON_DUMPS = orjson.dumps
except ImportError:
    _JSON_LOADS = json.loads
    _JSON_DUMPS = json.dumps

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 ThreadPoolExecutor(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 json.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 json.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 '%\"role\":\"user\"%' OR data LIKE '%\"role\":\"assistant\"%' "
        "GROUP BY session_id"
    ):
        msg_count[sid] = c

    # Build message_id -> role once with a LIKE-based query. This is
    # dramatically faster than calling json_extract on every row because
    # LIKE short-circuits on the compact "role":"..." prefix and the
    # planner can use the primary-key index on message.
    msg_role: dict[str, str] = {}
    for mid, data in con.execute("SELECT id, data FROM message"):
        if '"role":"user"' in data:
            msg_role[mid] = "user"
        elif '"role":"assistant"' in data:
            msg_role[mid] = "assistant"

    # Narrow the parts query: skip synthetic via LIKE (json_extract was
    # >2x slower on large dbs). Also extract text with json_extract only
    # on text parts so the path stays hot on the index.
    search_text: dict[str, list[tuple[str, str]]] = {}
    for sid, mid, text in con.execute(
        """
        SELECT p.session_id,
               p.message_id,
               json_extract(p.data, '$.text') AS txt
        FROM part p
        WHERE p.data LIKE '{"type":"text"%'
          AND p.data NOT LIKE '%"synthetic":true%'
        ORDER BY p.time_created ASC
        """
    ):
        if not text:
            continue
        role = msg_role.get(mid)
        if role is None:
            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 = json.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)
            con.execute("PRAGMA journal_mode=WAL")
            con.execute("PRAGMA cache_size=-64000")
        except sqlite3.Error as e:
            print(f"opencode: cannot open {db}: {e}", file=sys.stderr)
            continue
        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 json.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 json.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 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:
        with ThreadPoolExecutor(max_workers=len(agents)) as pool:
            future_map = {}
            for a in agents:
                walker = WALKERS.get(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 = WALKERS.get(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())
