db.ts

// SQLite FTS5-backed session index and query layer.
import { DatabaseSync } from "node:sqlite";
import { stat } from "node:fs/promises";
import { parseFile, type ParsedSession, type SessionFile, type Source } from "./parsers.ts";

export interface DiscoveryOptions {
  query: string;
  limit: number;
  sort?: "relevance" | "newest" | "oldest";
  roles?: ("user" | "assistant" | "tool")[];
}

export interface WindowMessage {
  id: number;
  role: string;
  content: string;
  ts: number;
  match?: boolean;
}

export interface DiscoveryHit {
  session_id: string;
  source: Source;
  title: string;
  cwd: string;
  when: number;
  snippet: string;
  bookend_start: WindowMessage[];
  messages: WindowMessage[];
  bookend_end: WindowMessage[];
  match_message_id: number;
  messages_before: number;
  messages_after: number;
}

const SNIPPET_LEN = 14; // fts5 snippet tokens
const MSG_RETURN_CHARS = 350; // per-message truncation when returning to the LLM
const ROLE_ALIASES: Record<string, "user" | "assistant" | "tool"> = {
  user: "user",
  assistant: "assistant",
  tool: "tool",
  toolresult: "tool",
};

export function parseRoles(filter: string | undefined): ("user" | "assistant" | "tool")[] | undefined {
  if (!filter) return undefined;
  const roles = filter
    .split(",")
    .map((r) => ROLE_ALIASES[r.trim().toLowerCase()])
    .filter((r): r is "user" | "assistant" | "tool" => r !== undefined);
  return roles.length > 0 ? [...new Set(roles)] : undefined;
}

function clip(content: string): string {
  return content.length > MSG_RETURN_CHARS ? content.slice(0, MSG_RETURN_CHARS) + "…" : content;
}

/** Count searchable (letter/number) characters; trigram needs >= 3. */
function searchableChars(q: string): number {
  let n = 0;
  for (const ch of q) {
    if (/[\p{L}\p{N}]/u.test(ch)) n++;
  }
  return n;
}

/** Escape LIKE wildcards so user input matches literally. */
function escapeLike(s: string): string {
  return s.replace(/[\\%_]/g, (c) => "\\" + c);
}

export class SessionSearchDb {
  private db: DatabaseSync;
  private indexing: Promise<number> | null = null;
  indexedBytes = 0;

  constructor(path: string) {
    this.db = new DatabaseSync(path);
    this.db.exec(
      "PRAGMA journal_mode=WAL; PRAGMA synchronous=NORMAL; PRAGMA busy_timeout=5000;",
    );
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS sessions (
        source TEXT NOT NULL,
        session_id TEXT NOT NULL,
        path TEXT NOT NULL,
        title TEXT NOT NULL DEFAULT '',
        cwd TEXT NOT NULL DEFAULT '',
        created_at INTEGER NOT NULL DEFAULT 0,
        updated_at INTEGER NOT NULL DEFAULT 0,
        msg_count INTEGER NOT NULL DEFAULT 0,
        mtime INTEGER NOT NULL DEFAULT 0,
        size INTEGER NOT NULL DEFAULT 0,
        PRIMARY KEY (source, session_id)
      );
      CREATE TABLE IF NOT EXISTS messages (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        source TEXT NOT NULL,
        session_id TEXT NOT NULL,
        seq INTEGER NOT NULL,
        msg_id TEXT NOT NULL DEFAULT '',
        role TEXT NOT NULL,
        content TEXT NOT NULL,
        ts INTEGER NOT NULL DEFAULT 0,
        UNIQUE (source, session_id, seq)
      );
      CREATE INDEX IF NOT EXISTS idx_messages_session ON messages (source, session_id, seq);
      CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(content, tokenize='trigram');
      CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
        INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
      END;
      CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
        INSERT INTO messages_fts(messages_fts, rowid, content) VALUES ('delete', old.id, old.content);
      END;
    `);
  }

  close() {
    try {
      this.db.close();
    } catch {
      /* ignore */
    }
  }

  // ------------------------------------------------------------ indexing

  /** Index one parsed session: replace its messages in a single transaction. */
  indexSession(session: ParsedSession, mtimeMs: number, size: number) {
    this.db.exec("BEGIN");
    try {
      // Same file may have changed its embedded session id; drop any stale
      // rows that previously mapped this path so path stays unique.
      const oldRows = this.db
        .prepare("SELECT source, session_id FROM sessions WHERE path = ?")
        .all(session.path) as { source: string; session_id: string }[];
      const delMsg = this.db.prepare("DELETE FROM messages WHERE source = ? AND session_id = ?");
      const delSes = this.db.prepare("DELETE FROM sessions WHERE source = ? AND session_id = ?");
      for (const o of oldRows) {
        if (o.source === session.source && o.session_id === session.sessionId) continue;
        delMsg.run(o.source, o.session_id);
        delSes.run(o.source, o.session_id);
      }
      delMsg.run(session.source, session.sessionId);
      const insMsg = this.db.prepare(
        "INSERT OR REPLACE INTO messages (source, session_id, seq, msg_id, role, content, ts) VALUES (?, ?, ?, ?, ?, ?, ?)",
      );
      for (const m of session.messages) {
        insMsg.run(session.source, session.sessionId, m.seq, m.msgId, m.role, m.content, m.ts);
      }
      this.db
        .prepare(
          `INSERT OR REPLACE INTO sessions
             (source, session_id, path, title, cwd, created_at, updated_at, msg_count, mtime, size)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
        )
        .run(
          session.source,
          session.sessionId,
          session.path,
          session.title,
          session.cwd,
          session.createdAt,
          session.updatedAt,
          session.messages.length,
          mtimeMs,
          size,
        );
      this.db.exec("COMMIT");
    } catch (e) {
      this.db.exec("ROLLBACK");
      throw e;
    }
  }

  private async needsReindex(f: SessionFile): Promise<boolean> {
    let st;
    try {
      st = await stat(f.path);
    } catch {
      return false; // file gone
    }
    // Track by canonical path (parsers may derive a different session id
    // from the file contents than the filename-based discovery id).
    const row = this.db
      .prepare("SELECT mtime, size FROM sessions WHERE path = ?")
      .get(f.path) as { mtime: number; size: number } | undefined;
    return row === undefined || row.mtime !== st.mtimeMs || row.size !== st.size;
  }

  /**
   * Incremental index: stat all known files, re-parse only changed ones.
   * Yields to the event loop between files and honors abort signals.
   * Returns number of files (re)indexed.
   */
  async incrementalIndex(files: SessionFile[], signal?: AbortSignal): Promise<number> {
    // Coalesce concurrent index runs (background warm-up + tool calls).
    if (this.indexing) return this.indexing;
    const run = this.runIndex(files, signal);
    this.indexing = run;
    try {
      return await run;
    } finally {
      this.indexing = null;
    }
  }

  private async runIndex(files: SessionFile[], signal?: AbortSignal): Promise<number> {
    let changed = 0;
    const pending: SessionFile[] = [];
    for (const f of files) {
      if (await this.needsReindex(f)) pending.push(f);
    }
    for (const f of pending) {
      if (signal?.aborted) break;
      let st;
      try {
        st = await stat(f.path);
      } catch {
        continue;
      }
      try {
        const session = await parseFile(f.source, f.path);
        // Skip empty results so a transiently unparseable file cannot wipe
        // a previously valid index.
        if (session.messages.length === 0) continue;
        this.indexSession(session, st.mtimeMs, st.size);
        this.indexedBytes += st.size;
        changed++;
      } catch {
        // skip unparseable files
      }
      // Yield so the TUI stays responsive during the first full build.
      await new Promise((r) => setImmediate(r));
    }
    this.pruneDeleted(files);
    return changed;
  }

  /** Drop sessions whose files no longer exist (moved/deleted/archived away). */
  private pruneDeleted(files: SessionFile[]) {
    const live = new Set(files.map((f) => f.path));
    const rows = this.db
      .prepare("SELECT source, session_id, path FROM sessions")
      .all() as { source: string; session_id: string; path: string }[];
    const delMsg = this.db.prepare("DELETE FROM messages WHERE source = ? AND session_id = ?");
    const delSes = this.db.prepare("DELETE FROM sessions WHERE source = ? AND session_id = ?");
    let n = 0;
    for (const r of rows) {
      if (live.has(r.path)) continue;
      delMsg.run(r.source, r.session_id);
      delSes.run(r.source, r.session_id);
      n++;
    }
    if (n > 0) this.indexedBytes = 0; // approximate; size is not tracked per row here
  }

  // ------------------------------------------------------------ queries

  /**
   * FTS5 match expression from a raw query. Returns null when the query is
   * too short for trigram matching (caller falls back to LIKE).
   */
  private buildMatch(query: string): string | null {
    const q = query.trim();
    // Trigram tokenizer needs at least 3 searchable characters per term.
    if (!q || searchableChars(q) < 3) return null;
    // Pass user FTS5 syntax through (AND default, OR/NOT/phrases supported).
    return q;
  }

  private roleClause(roles?: ("user" | "assistant" | "tool")[]): { sql: string; params: string[] } {
    if (!roles || roles.length === 0 || roles.length === 3) return { sql: "", params: [] };
    const placeholders = roles.map(() => "?").join(",");
    return { sql: ` AND m.role IN (${placeholders})`, params: roles };
  }

  private async bestHitRows(
    query: string,
    roles?: ("user" | "assistant" | "tool")[],
    capRows = 400,
  ): Promise<{ id: number; source: Source; session_id: string; seq: number; role: string; content: string; ts: number; snip: string }[]> {
    const match = this.buildMatch(query);
    const rc = this.roleClause(roles);

    if (match) {
      try {
        return this.db
          .prepare(
            `SELECT m.id, m.source, m.session_id, m.seq, m.role, m.content, m.ts,
                    snippet(messages_fts, 0, '«', '»', '…', ${SNIPPET_LEN}) AS snip
             FROM messages_fts
             JOIN messages m ON m.id = messages_fts.rowid
             WHERE messages_fts MATCH ?${rc.sql}
             ORDER BY rank
             LIMIT ?`,
          )
          .all(match, ...rc.params, capRows) as any[];
      } catch {
        // FTS5 syntax error: fall back to quoted-token AND search.
      }
      try {
        const tokens = query
          .split(/\s+/)
          .map((t) => t.replace(/[^\p{L}\p{N}_-]+/gu, ""))
          .filter((t) => t.length >= 3);
        if (tokens.length > 0) {
          const expr = tokens.map((t) => `"${t}"`).join(" AND ");
          return this.db
            .prepare(
              `SELECT m.id, m.source, m.session_id, m.seq, m.role, m.content, m.ts,
                      snippet(messages_fts, 0, '«', '»', '…', ${SNIPPET_LEN}) AS snip
               FROM messages_fts
               JOIN messages m ON m.id = messages_fts.rowid
               WHERE messages_fts MATCH ?${rc.sql}
               ORDER BY rank
               LIMIT ?`,
            )
            .all(expr, ...rc.params, capRows) as any[];
        }
      } catch {
        /* fall through to LIKE */
      }
    }

    // Short queries (<3 searchable chars) or FTS failure: LIKE scan over
    // user/assistant content with wildcards escaped for literal matching.
    const like = `%${escapeLike(query.trim())}%`;
    const rows = this.db
      .prepare(
        `SELECT m.id, m.source, m.session_id, m.seq, m.role, m.content, m.ts,
                CASE WHEN length(m.content) < 400 THEN m.content ELSE substr(m.content, 1, 400) END AS snip
         FROM messages AS m
         WHERE m.content LIKE ? ESCAPE '\\'${rc.sql}
         ORDER BY length(m.content)
         LIMIT ?`,
      )
      .all(like, ...rc.params, capRows) as any[];
    return rows;
  }

  private windowMessages(source: Source, sessionId: string, centerSeq: number, half: number) {
    const rows = this.db
      .prepare(
        `SELECT id, seq, role, content, ts FROM messages
         WHERE source = ? AND session_id = ? AND seq >= ? AND seq <= ?
         ORDER BY seq`,
      )
      .all(source, sessionId, Math.max(0, centerSeq - half), centerSeq + half) as any[];
    return rows.map((r) => ({ id: r.seq, role: r.role, content: clip(r.content), ts: r.ts }));
  }

  /** Discovery: full-text search, deduped per session, with bookends + window. */
  async discover(opts: DiscoveryOptions): Promise<DiscoveryHit[]> {
    const hits = await this.bestHitRows(opts.query, opts.roles);

    // Group hits by session, keep best-ranked hit per session.
    const bySession = new Map<string, typeof hits>();
    for (const h of hits) {
      const key = `${h.source}\u0000${h.session_id}`;
      if (!bySession.has(key)) bySession.set(key, []);
      bySession.get(key)!.push(h);
    }

    const perSession: { source: Source; session_id: string; hits: typeof hits }[] = [];
    for (const [key, hs] of bySession) {
      const [source, session_id] = key.split("\u0000");
      perSession.push({ source: source as Source, session_id, hits: hs });
    }

    // Order sessions: relevance by default, or by session recency.
    const getSession = (s: Source, id: string) =>
      this.db
        .prepare("SELECT * FROM sessions WHERE source = ? AND session_id = ?")
        .get(s, id) as Record<string, any> | undefined;

    if (opts.sort === "newest" || opts.sort === "oldest") {
      const dir = opts.sort === "newest" ? -1 : 1;
      perSession.sort((a, b) => {
        const sa = getSession(a.source, a.session_id)?.updated_at ?? 0;
        const sb = getSession(b.source, b.session_id)?.updated_at ?? 0;
        return (sa - sb) * dir;
      });
    }

    const results: DiscoveryHit[] = [];
    for (const ps of perSession.slice(0, opts.limit)) {
      const best = ps.hits[0];
      const meta = getSession(ps.source, ps.session_id);

      // Bookends: first/last 3 user+assistant messages of the session
      // (bounded queries, no full transcript scan).
      const bookendSql = `SELECT seq, role, content, ts FROM messages
        WHERE source = ? AND session_id = ? AND role IN ('user','assistant')`;
      const start = (this.db.prepare(bookendSql + " ORDER BY seq ASC LIMIT 3").all(ps.source, ps.session_id) as any[])
        .map((r) => ({ id: r.seq, role: r.role, content: clip(r.content), ts: r.ts }));
      const end = (this.db.prepare(bookendSql + " ORDER BY seq DESC LIMIT 3").all(ps.source, ps.session_id) as any[])
        .reverse()
        .map((r) => ({ id: r.seq, role: r.role, content: clip(r.content), ts: r.ts }));

      const win = this.windowMessages(ps.source, ps.session_id, best.seq, 5).map((m) =>
        m.id === best.seq
          ? { ...m, content: (best.snip ?? m.content).slice(0, MSG_RETURN_CHARS), match: true }
          : m,
      );

      const total = this.db
        .prepare("SELECT COUNT(*) AS c FROM messages WHERE source = ? AND session_id = ?")
        .get(ps.source, ps.session_id) as { c: number };

      results.push({
        session_id: ps.session_id,
        source: ps.source,
        title: meta?.title ?? "",
        cwd: meta?.cwd ?? "",
        when: meta?.updated_at ?? 0,
        snippet: (best.snip ?? best.content ?? "").slice(0, 600),
        bookend_start: start,
        messages: win,
        bookend_end: end,
        match_message_id: best.seq,
        messages_before: best.seq,
        messages_after: Math.max(0, total.c - best.seq - 1),
      });
    }
    return results;
  }

  /** Scroll: window of ±window messages around a message in one session. */
  scroll(
    source: Source | undefined,
    sessionId: string,
    aroundSeq: number,
    window: number,
  ): { messages: WindowMessage[]; before: number; after: number; ambiguous?: string[] } | null {
    // Session ids may collide across sources; resolve by source when given.
    const rows = this.db
      .prepare("SELECT source FROM sessions WHERE session_id = ?")
      .all(sessionId) as { source: Source }[];
    if (rows.length === 0) return null;
    if (rows.length > 1 && !source) {
      return { messages: [], before: 0, after: 0, ambiguous: rows.map((r) => r.source) };
    }
    const src = source ?? rows[0].source;
    const total = this.db
      .prepare("SELECT COUNT(*) AS c FROM messages WHERE source = ? AND session_id = ?")
      .get(src, sessionId) as { c: number };
    // ±window semantics: window=10 returns 10 messages before and after.
    const half = Math.max(1, Math.min(50, window));
    const messages = this.windowMessages(src, sessionId, aroundSeq, half).map((m) =>
      m.id === aroundSeq ? { ...m, match: true } : m,
    );
    return { messages, before: aroundSeq, after: Math.max(0, total.c - aroundSeq - 1) };
  }

  /** Browse: recent sessions chronologically. */
  browse(limit: number) {
    return this.db
      .prepare(
        `SELECT s.source, s.session_id, s.title, s.cwd, s.updated_at, s.msg_count,
                (SELECT content FROM messages m WHERE m.source = s.source AND m.session_id = s.session_id
                 AND m.role = 'user' ORDER BY m.seq LIMIT 1) AS preview
         FROM sessions s
         WHERE s.msg_count > 0
         ORDER BY s.updated_at DESC
         LIMIT ?`,
      )
      .all(limit) as any[];
  }

  stats() {
    return this.db
      .prepare(
        `SELECT (SELECT COUNT(*) FROM sessions) AS sessions,
                (SELECT COUNT(*) FROM messages) AS messages`,
      )
      .get() as { sessions: number; messages: number };
  }
}