#!/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("~"))
INDEX_JOBS = max(
    1,
    int(os.environ.get("FZFAI_INDEX_JOBS") or min(32, (os.cpu_count() or 4) + 4)),
)

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

# 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"


@dataclass
class Record:
    agent: str
    session_id: str
    source: str
    updated: float = 0.0  # unix seconds
    msgs: int = 0
    cwd: str = ""
    title: str = ""
    # message snippets joined with " · " for fuzzy search (hidden column)
    prompts: list = field(default_factory=list)

    def add_search_text(self, text: str) -> None:
        t = clean(text)[:MAX_SEARCH_CHUNK]
        if any(t.startswith(prefix) for prefix in _SEARCH_NOISE_PREFIXES):
            return
        if t and t not in self.prompts:
            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 ""
        # short cwd: replace $HOME with ~ for display compactness
        if cwd.startswith(str(HOME)):
            cwd = "~" + cwd[len(str(HOME)) :]
        cwd = display_text(cwd, DISPLAY_CWD)
        # Build search blob: strip TABs/newlines, cap total size.
        blob = " · ".join(self.prompts)
        blob = blob.replace("\t", " ").replace("\n", " ")[:MAX_SEARCH_BLOB]
        # Layout: columns 1..3 are machine-only (hidden from view); columns
        # 4..8 are visible. Column 9 is a hidden search-only blob.
        return "\t".join(
            [
                self.agent,                     # 1  raw agent   (hidden)
                self.session_id,                # 2  session id  (hidden)
                self.source,                    # 3  source      (hidden)
                colorise(self.agent),           # 4  agent badge (visible)
                f"{iso:<16}",                   # 5  date
                f"{self.msgs:>4}",              # 6  msgs
                f"{cwd:<{DISPLAY_CWD}}",        # 7  cwd
                title,                          # 8  title
                blob,                           # 9  search blob (hidden)
            ]
        )


def clean(s: str) -> str:
    if not s:
        return ""
    s = re.sub(r"\s+", " ", s)
    return 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",
)


def is_real_prompt(text: str) -> bool:
    """Return True only for what looks like an organic user prompt."""
    if not text:
        return False
    t = text.lstrip()
    if not t:
        return False
    for p in _JUNK_PREFIXES:
        if t.startswith(p):
            return False
    return True


def iso_to_epoch(s: str) -> float:
    if not s:
        return 0.0
    try:
        # handle trailing Z
        s = s.replace("Z", "+00:00")
        return datetime.fromisoformat(s).timestamp()
    except Exception:
        return 0.0


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() -> Iterator[Record]:
    root = HOME / ".claude" / "projects"
    if not root.is_dir():
        return
    paths: list[Path] = []
    for proj in root.iterdir():
        if not proj.is_dir():
            continue
        paths.extend(proj.glob("*.jsonl"))
    yield from _parallel_read(paths, _read_claude_file, "claude")


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
    with path.open("r", encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue
            if obj.get("type") in ("user", "assistant"):
                msgs += 1
            if not cwd and obj.get("cwd"):
                cwd = obj["cwd"]
            if obj.get("type") in ("user", "assistant") and not obj.get("isMeta"):
                m = obj.get("message") or {}
                for tx in iter_text_fragments(m.get("content")):
                    if obj.get("type") == "user" and is_real_prompt(tx):
                        if not title:
                            title = tx
                        rec.add_prompt(tx)
                    else:
                        rec.add_search_text(tx)
            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() -> Iterator[Record]:
    root = HOME / ".codex" / "sessions"
    if not root.is_dir():
        return
    paths = list(root.rglob("rollout-*.jsonl"))
    yield from _parallel_read(paths, _read_codex_file, "codex")


def _read_codex_file(path: Path) -> Optional[Record]:
    # filename: rollout-2025-09-25T21-35-23-<uuid>.jsonl
    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
    with path.open("r", encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue
            t = obj.get("type")
            payload = obj.get("payload") or {}
            if t == "session_meta":
                if payload.get("cwd"):
                    cwd = payload["cwd"]
                if payload.get("id"):
                    rec.session_id = payload["id"]
            if t == "response_item" and payload.get("type") == "message":
                role = payload.get("role")
                if role in ("user", "assistant"):
                    msgs += 1
                if role in ("user", "assistant"):
                    for tx in iter_text_fragments(payload.get("content", []) or []):
                        if role == "user" and is_real_prompt(tx):
                            if not title:
                                title = tx
                            rec.add_prompt(tx)
                        else:
                            rec.add_search_text(tx)
            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]:
    """Current schema: session / message / part.

    We push as much work as possible down to SQLite using json_extract:
      * message counts per session via a GROUP BY
      * user-prompt text pulled directly as a column, joining message
        to filter on role = 'user' at the DB layer
    This keeps the Python pass O(sessions + user_text_parts), not O(parts).
    """
    cur = con.execute(
        "SELECT id, directory, title, time_updated FROM session ORDER BY time_updated DESC"
    )
    sessions = cur.fetchall()
    if not sessions:
        return

    # Message count per session (user + assistant only).
    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

    # Search text extracted in SQL. We keep the title derived from real user
    # prompts but also index assistant text so queries can match session
    # content without showing it in the main list rows.
    search_text: dict[str, list[tuple[str, str]]] = {}
    for sid, role, text in con.execute(
        """
        SELECT p.session_id,
               json_extract(m.data, '$.role') AS role,
               json_extract(p.data, '$.text') AS txt
        FROM part p
        JOIN message m ON p.message_id = m.id
        WHERE p.data LIKE '{"type":"text"%'
          AND (json_extract(p.data, '$.synthetic') IS NULL
               OR json_extract(p.data, '$.synthetic') = 0)
        ORDER BY p.time_created ASC
        """
    ):
        if text and role in ("user", "assistant"):
            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() -> Iterator[Record]:
    seen_sids: set[str] = set()
    for db in _opencode_db_paths():
        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
        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)
                yield rec
        except sqlite3.Error as e:
            print(f"opencode: {db}: {e}", file=sys.stderr)
        finally:
            con.close()


# --------------------------------------------------------------------------
# Droid / Factory -- ~/.factory/sessions/<proj>/<uuid>.jsonl
# --------------------------------------------------------------------------
def walk_droid() -> Iterator[Record]:
    root = HOME / ".factory" / "sessions"
    if not root.is_dir():
        return
    paths: list[Path] = []
    for proj in root.iterdir():
        if not proj.is_dir():
            continue
        paths.extend(proj.glob("*.jsonl"))
    yield from _parallel_read(paths, _read_droid_file, "droid")


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
    with path.open("r", encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue
            t = obj.get("type")
            if t == "session_start":
                if obj.get("cwd"):
                    cwd = obj["cwd"]
                if obj.get("title") and not title:
                    title = obj["title"]
            if t == "message":
                m = obj.get("message") or {}
                role = m.get("role")
                if role in ("user", "assistant"):
                    msgs += 1
                if role in ("user", "assistant"):
                    for tx in iter_text_fragments(m.get("content") or []):
                        if role == "user" and is_real_prompt(tx):
                            if not title:
                                title = tx
                            rec.add_prompt(tx)
                        else:
                            rec.add_search_text(tx)
            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() -> Iterator[Record]:
    root = HOME / ".pi" / "agent" / "sessions"
    if not root.is_dir():
        return
    paths: list[Path] = []
    for proj in root.iterdir():
        if not proj.is_dir():
            continue
        paths.extend(proj.glob("*.jsonl"))
    yield from _parallel_read(paths, _read_pi_file, "pi")


def _read_pi_file(path: Path) -> Optional[Record]:
    # filename: 2026-04-02T12-02-19-290Z_<uuid>.jsonl
    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
    with path.open("r", encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue
            t = obj.get("type")
            if t == "session":
                if obj.get("cwd"):
                    cwd = obj["cwd"]
                if obj.get("id"):
                    rec.session_id = obj["id"]
            if t == "message":
                m = obj.get("message") or {}
                role = m.get("role")
                if role in ("user", "assistant"):
                    msgs += 1
                if role in ("user", "assistant"):
                    for tx in iter_text_fragments(m.get("content")):
                        if role == "user" and is_real_prompt(tx):
                            if not title:
                                title = tx
                            rec.add_prompt(tx)
                        else:
                            rec.add_search_text(tx)
            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:
    agents = sys.argv[1:] or list(WALKERS.keys())
    records: list[Record] = []
    for a in agents:
        walker = WALKERS.get(a)
        if not walker:
            print(f"unknown agent: {a}", file=sys.stderr)
            continue
        try:
            records.extend(walker())
        except Exception as e:
            print(f"{a}: {e}", file=sys.stderr)
    records.sort(key=lambda r: r.updated, reverse=True)
    out = sys.stdout
    for r in records:
        out.write(r.as_row())
        out.write("\n")
    return 0


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