// src/kb/KbStore.ts
import { Database, type Statement } from "bun:sqlite";
import { createHash } from "node:crypto";
import { chmodSync, existsSync, copyFileSync } from "node:fs";
import { ensureSecureDir } from "../install/blocks/secureWrite.ts";
import { join } from "node:path";
import { homedir } from "node:os";
import type { KbSearchResult } from "./types.ts";
import { chunkContent } from "./KbChunker.ts";
import { jaccardDedup } from "./jaccardDedup.ts";
import { profileLog, pnow } from "../proxy/profile.ts";

const KB_DIR = join(homedir(), ".fewtok", "kb");
const TTL_DAYS = Number(process.env.FEWTOK_KB_TTL_DAYS ?? 14);

function projectHash(projectDir: string): string {
  return createHash("sha256").update(projectDir.slice(0, 64)).digest("hex").slice(0, 16);
}

function sha256hex(content: string): string {
  return createHash("sha256").update(content).digest("hex");
}

export class KbStore {
  private readonly db: Database;
  private insertCount = 0;
  private readonly PRUNE_EVERY = 1000;
  private stmtSelectExisting!: Statement<{ id: number }, [string, string, string]>;
  private insertAllTx!: (label: string, sessionId: string, contentHash: string, chunks: string[], byteLen: number) => void;

  constructor(projectDir: string) {
    ensureSecureDir(KB_DIR);
    const dbPath = join(KB_DIR, `${projectHash(projectDir)}.db`);
    this.db = new Database(dbPath);
    chmodSync(dbPath, 0o600);
    this.initSchema();
  }

  private initSchema(): void {
    // N=2 (Phase 3): both children open this per-project DB at handler boot; serialize
    // the WAL/DDL write lock so the race loser waits instead of throwing SQLITE_BUSY.
    this.db.run("PRAGMA busy_timeout=5000");
    // auto_vacuum must be set before any write (including WAL header commit)
    this.db.run("PRAGMA auto_vacuum=INCREMENTAL");
    // WAL mode: parallel readers + single writer, no write contention
    this.db.run("PRAGMA journal_mode=WAL");
    this.db.run("PRAGMA synchronous=OFF");

    this.db.run(`
      CREATE TABLE IF NOT EXISTS sources (
        id           INTEGER PRIMARY KEY AUTOINCREMENT,
        label        TEXT NOT NULL,
        session_id   TEXT NOT NULL,
        content_hash TEXT NOT NULL,
        indexed_at   TEXT NOT NULL DEFAULT (datetime('now')),
        UNIQUE(label, session_id, content_hash)
      )
    `);
    this.migrateSourcesSessionScopedUnique();
    this.migrateToContentAddressed();

    this.db.run(`
      CREATE TABLE IF NOT EXISTS blobs (
        content_hash TEXT PRIMARY KEY,
        byte_len     INTEGER NOT NULL,
        has_trigram  INTEGER NOT NULL DEFAULT 1,
        indexed_at   TEXT NOT NULL DEFAULT (datetime('now'))
      )
    `);

    this.db.run(`
      CREATE VIRTUAL TABLE IF NOT EXISTS chunks USING fts5(
        content_hash UNINDEXED,
        content,
        tokenize = 'porter unicode61'
      )
    `);

    this.db.run(`
      CREATE VIRTUAL TABLE IF NOT EXISTS chunks_trigram USING fts5(
        content_hash UNINDEXED,
        content,
        tokenize = 'trigram'
      )
    `);

    this.db.run(
      "CREATE INDEX IF NOT EXISTS idx_sources_indexed ON sources(indexed_at)",
    );
    this.db.run(
      "CREATE INDEX IF NOT EXISTS idx_sources_chash ON sources(content_hash)",
    );

    this.stmtSelectExisting = this.db.prepare(
      "SELECT id FROM sources WHERE label = ? AND session_id = ? AND content_hash = ?",
    );

    this.insertAllTx = this.db.transaction(
      (label: string, sessionId: string, contentHash: string, chunks: string[], byteLen: number) => {
        this.db.run(
          "INSERT OR IGNORE INTO sources(label, session_id, content_hash) VALUES (?, ?, ?)",
          [label, sessionId, contentHash],
        );
        const have = this.db
          .query<{ x: number }, [string]>("SELECT 1 AS x FROM blobs WHERE content_hash = ?")
          .get(contentHash);
        if (have) return; // blob already indexed across the whole DB
        const trigramMax = Number(process.env.FEWTOK_KB_TRIGRAM_MAX ?? 32 * 1024);
        const trigram = byteLen <= trigramMax;
        for (const c of chunks) {
          this.db.run("INSERT INTO chunks(content_hash, content) VALUES (?, ?)", [contentHash, c]);
          if (trigram) this.db.run("INSERT INTO chunks_trigram(content_hash, content) VALUES (?, ?)", [contentHash, c]);
        }
        this.db.run(
          "INSERT INTO blobs(content_hash, byte_len, has_trigram) VALUES (?, ?, ?)",
          [contentHash, byteLen, trigram ? 1 : 0],
        );
      },
    );
  }

  insert(label: string, sessionId: string, content: string): void;
  insert(label: string, content: string): void;
  insert(label: string, sessionIdOrContent: string, maybeContent?: string): void {
    const sessionId = maybeContent === undefined ? "__default__" : sessionIdOrContent;
    const content = maybeContent === undefined ? sessionIdOrContent : maybeContent;
    if (!content.trim()) return;
    const t0 = pnow();
    const contentHash = sha256hex(content);

    const existing = this.stmtSelectExisting.get(label, sessionId, contentHash);
    if (existing) return;

    const chunks = chunkContent(content);
    const byteLen = Buffer.byteLength(content, "utf8");

    this.insertAllTx(label, sessionId, contentHash, chunks, byteLen);

    this.insertCount++;
    if (this.insertCount % this.PRUNE_EVERY === 0) this.maintain();
    profileLog("kb.insert", pnow() - t0, `bytes=${content.length} chunks=${chunks.length}`);
  }

  private migrateSourcesSessionScopedUnique(): void {
    const row = this.db
      .query<{ sql: string | null }, []>(
        "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'sources'",
      )
      .get();
    const sql = row?.sql ?? "";
    if (!/UNIQUE\s*\(\s*label\s*,\s*content_hash\s*\)/i.test(sql)) return;

    const migrate = this.db.transaction(() => {
      this.db.run("ALTER TABLE sources RENAME TO sources_old");
      this.db.run(`
        CREATE TABLE sources (
          id           INTEGER PRIMARY KEY AUTOINCREMENT,
          label        TEXT NOT NULL,
          session_id   TEXT NOT NULL,
          content_hash TEXT NOT NULL,
          indexed_at   TEXT NOT NULL DEFAULT (datetime('now')),
          UNIQUE(label, session_id, content_hash)
        )
      `);
      this.db.run(`
        INSERT OR IGNORE INTO sources(id, label, session_id, content_hash, indexed_at)
        SELECT id, label, session_id, content_hash, indexed_at FROM sources_old
      `);
      this.db.run("DROP TABLE sources_old");
    });
    migrate();
  }

  private migrateToContentAddressed(): void {
    const hasBlobs = this.db
      .query<{ x: number }, []>("SELECT 1 AS x FROM sqlite_master WHERE type='table' AND name='blobs'")
      .get();
    if (hasBlobs) return;
    // detect legacy source_id-keyed chunks
    const chunksSql = this.db
      .query<{ sql: string | null }, []>("SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks'")
      .get()?.sql ?? "";
    const legacy = /source_id/i.test(chunksSql);
    // mandatory backup BEFORE any DDL (one-way door → recoverable)
    if (legacy) {
      const dbPath = this.db.filename;
      if (!existsSync(dbPath + ".pre-ca-backup")) copyFileSync(dbPath, dbPath + ".pre-ca-backup");
    }
    const migrate = this.db.transaction(() => {
      this.db.run("CREATE TABLE IF NOT EXISTS blobs (content_hash TEXT PRIMARY KEY, byte_len INTEGER NOT NULL, has_trigram INTEGER NOT NULL DEFAULT 1, indexed_at TEXT NOT NULL DEFAULT (datetime('now')))");
      this.db.run("CREATE INDEX IF NOT EXISTS idx_sources_chash ON sources(content_hash)");
      if (!legacy) return; // fresh DB: blobs created, nothing to backfill
      this.db.run("CREATE VIRTUAL TABLE chunks_v2 USING fts5(content_hash UNINDEXED, content, tokenize='porter unicode61')");
      this.db.run("CREATE VIRTUAL TABLE chunks_tri_v2 USING fts5(content_hash UNINDEXED, content, tokenize='trigram')");
      const trigramMax = Number(process.env.FEWTOK_KB_TRIGRAM_MAX ?? 32 * 1024);
      // distinct content_hash, representative = MIN(source id)
      const hashes = this.db
        .query<{ content_hash: string; rep: number; first_at: string }, []>(
          "SELECT content_hash, MIN(id) AS rep, MIN(indexed_at) AS first_at FROM sources GROUP BY content_hash",
        )
        .all();
      for (const h of hashes) {
        const rows = this.db
          .query<{ content: string }, [number]>("SELECT content FROM chunks WHERE source_id = ? ORDER BY rowid")
          .all(h.rep);
        if (rows.length === 0) continue;
        let byteLen = 0;
        for (const r of rows) byteLen += Buffer.byteLength(r.content, "utf8"); // canonical: SUM(byteLen(chunk))
        const trigram = byteLen <= trigramMax;
        for (const r of rows) {
          this.db.run("INSERT INTO chunks_v2(content_hash, content) VALUES (?, ?)", [h.content_hash, r.content]);
          if (trigram) this.db.run("INSERT INTO chunks_tri_v2(content_hash, content) VALUES (?, ?)", [h.content_hash, r.content]);
        }
        this.db.run("INSERT INTO blobs(content_hash, byte_len, has_trigram, indexed_at) VALUES (?, ?, ?, ?)", [h.content_hash, byteLen, trigram ? 1 : 0, h.first_at]);
      }
      this.db.run("DROP TABLE chunks");
      this.db.run("DROP TABLE chunks_trigram");
      this.db.run("ALTER TABLE chunks_v2 RENAME TO chunks");
      this.db.run("ALTER TABLE chunks_tri_v2 RENAME TO chunks_trigram");
    });
    migrate();
    if (legacy) this.db.run("VACUUM"); // one-shot reclaim; locks DB, needs ~2x file size transiently
  }

  search(q: string, opts: { sessionId: string; limit?: number }): KbSearchResult[];
  search(q: string, limit?: number): KbSearchResult[];
  search(q: string, opts: { sessionId: string; limit?: number } | number = { sessionId: "__default__" }): KbSearchResult[] {
    const sessionId = typeof opts === "number" ? "__default__" : opts.sessionId ?? "__default__";
    const limit = typeof opts === "number" ? opts : opts.limit ?? 5;
    if (!q.trim()) return [];

    // Tier 1: porter FTS5 — stemmed search; blob join avoids label fanout
    let porterRows: Array<{ label: string; content: string; indexed_at: string; rank: number }> = [];
    try {
      porterRows = this.db
        .query<
          { label: string; content: string; indexed_at: string; rank: number },
          [string, string, string, number]
        >(`
          SELECT
            (SELECT s.label FROM sources s WHERE s.content_hash = c.content_hash AND s.session_id = ? LIMIT 1) AS label,
            c.content, b.indexed_at, c.rank
          FROM chunks c
          JOIN blobs b ON b.content_hash = c.content_hash
          WHERE chunks MATCH ?
            AND EXISTS (SELECT 1 FROM sources s WHERE s.content_hash = c.content_hash AND s.session_id = ?)
          ORDER BY c.rank
          LIMIT ?
        `)
        .all(sessionId, q, sessionId, limit * 2);
    } catch {
      // FTS5 MATCH syntax error (e.g. bare operator, colon query) — skip tier 1
    }

    const results: KbSearchResult[] = porterRows.map(r => {
      const hoursAgo = (Date.now() - new Date(r.indexed_at).getTime()) / 3_600_000;
      const recencyBoost = Math.max(0, 1.0 - hoursAgo / 24.0);
      return {
        toolCallId: r.label,
        sourceLabel: r.label,
        content: r.content,
        score: Math.abs(r.rank) * (1.0 + recencyBoost),
        tier: "porter" as const,
      };
    });

    // Tier 2: trigram fallback when porter found < 2 results
    if (results.length < 2) {
      let trigramRows: Array<{ label: string; content: string; rank: number }> = [];
      try {
        trigramRows = this.db
          .query<
            { label: string; content: string; rank: number },
            [string, string, string, number]
          >(`
            SELECT
              (SELECT s.label FROM sources s WHERE s.content_hash = c.content_hash AND s.session_id = ? LIMIT 1) AS label,
              c.content, c.rank
            FROM chunks_trigram c
            JOIN blobs b ON b.content_hash = c.content_hash
            WHERE chunks_trigram MATCH ?
              AND EXISTS (SELECT 1 FROM sources s WHERE s.content_hash = c.content_hash AND s.session_id = ?)
            ORDER BY c.rank
            LIMIT ?
          `)
          .all(sessionId, q, sessionId, limit);
      } catch {
        // FTS5 MATCH syntax error — skip trigram tier
      }
      for (const r of trigramRows) {
        results.push({
          toolCallId: r.label,
          sourceLabel: r.label,
          content: r.content,
          score: Math.abs(r.rank),
          tier: "trigram" as const,
        });
      }
    }

    results.sort((a, b) => b.score - a.score);
    return jaccardDedup(results, 0.8).slice(0, limit);
  }

  searchByLabel(label: string, q: string, opts: { sessionId: string; limit?: number }): KbSearchResult[] {
    const sessionId = opts.sessionId;
    const limit = opts.limit ?? 5;
    if (!label.trim()) return [];

    let rows: Array<{ label: string; content: string; indexed_at: string; rank: number }> = [];
    if (q.trim()) {
      try {
        rows = this.db
          .query<
            { label: string; content: string; indexed_at: string; rank: number },
            [string, string, string, string, number]
          >(`
            SELECT ? AS label, c.content, b.indexed_at, c.rank
            FROM chunks c JOIN blobs b ON b.content_hash = c.content_hash
            WHERE chunks MATCH ?
              AND EXISTS (SELECT 1 FROM sources s WHERE s.content_hash = c.content_hash AND s.session_id = ? AND s.label = ?)
            ORDER BY c.rank LIMIT ?
          `)
          .all(label, q, sessionId, label, limit * 2);
      } catch {
        rows = [];
      }
    }

    const ranked = jaccardDedup(rows.map(r => ({
      toolCallId: r.label,
      sourceLabel: r.label,
      content: r.content,
      score: Math.abs(r.rank),
      tier: "porter" as const,
    })), 0.8).slice(0, limit);
    if (ranked.length > 0) return ranked;

    const fallbackRows = this.db
      .query<
        { label: string; content: string },
        [string, string, string, number]
      >(`
        SELECT ? AS label, c.content
        FROM chunks c JOIN blobs b ON b.content_hash = c.content_hash
        WHERE EXISTS (SELECT 1 FROM sources s WHERE s.content_hash = c.content_hash AND s.session_id = ? AND s.label = ?)
        ORDER BY b.indexed_at DESC, c.rowid LIMIT ?
      `)
      .all(label, sessionId, label, limit);

    return jaccardDedup(
      fallbackRows.map(r => ({
        toolCallId: r.label,
        sourceLabel: r.label,
        content: r.content,
        score: 0,
        tier: "porter" as const,
      })),
      0.8,
    ).slice(0, limit);
  }

  pruneOld(ttlDays = 14): void {
    const interval = `-${ttlDays} days`;
    this.db.run(
      "DELETE FROM sources WHERE indexed_at < datetime('now', ?)",
      [interval],
    );
  }

  /** Bytes currently used by the DB file (pages * page_size). */
  private dbByteSize(): number {
    const pc = this.db.query<{ page_count: number }, []>("PRAGMA page_count").get()?.page_count ?? 0;
    const ps = this.db.query<{ page_size: number }, []>("PRAGMA page_size").get()?.page_size ?? 0;
    return pc * ps;
  }

  /** Delete chunk/blob rows whose content_hash no longer exists in sources. */
  private gcOrphanBlobs(): void {
    this.db.run("DELETE FROM chunks WHERE content_hash NOT IN (SELECT content_hash FROM sources)");
    this.db.run("DELETE FROM chunks_trigram WHERE content_hash NOT IN (SELECT content_hash FROM sources)");
    this.db.run("DELETE FROM blobs WHERE content_hash NOT IN (SELECT content_hash FROM sources)");
  }

  /**
   * TTL prune + hard size ceiling. Sources are the GC roots: always delete the
   * oldest sources first (by indexed_at), then their now-orphaned chunks.
   */
  maintain(): void {
    this.pruneOld(TTL_DAYS);
    this.gcOrphanBlobs();

    const autoVacuum =
      this.db.query<{ auto_vacuum: number }, []>("PRAGMA auto_vacuum").get()?.auto_vacuum ?? 0;
    if (autoVacuum !== 2) return; // size reclaim requires INCREMENTAL; skip loop on legacy DBs

    const maxBytes = Number(process.env.FEWTOK_KB_MAX_BYTES ?? 256 * 1024 * 1024);
    let size = this.dbByteSize();
    let guard = 0;
    while (size > maxBytes && guard++ < 1000) {
      const total = this.db.query<{ n: number }, []>("SELECT count(*) AS n FROM sources").get()?.n ?? 0;
      if (total === 0) break;
      const batch = Math.max(50, Math.floor(total * 0.05));
      if (total <= batch) break; // FLOOR GUARD: never drain to empty chasing an unreachable cap (over-cap beats deleted)
      this.db.run(
        `DELETE FROM sources WHERE id IN (SELECT id FROM sources ORDER BY indexed_at ASC, id ASC LIMIT ?)`,
        [batch],
      );
      this.gcOrphanBlobs();
      this.db.run("PRAGMA incremental_vacuum");
      const newSize = this.dbByteSize();
      if (newSize >= size) break; // DATA-LOSS GUARD: vacuum not reclaiming → STOP before emptying DB (over-cap beats deleted)
      size = newSize;
    }
  }

  close(): void {
    this.db.close();
  }

  /** Test/diagnostic: number of rows currently in the trigram FTS table. */
  debugTrigramRowCount(): number {
    const row = this.db.query<{ n: number }, []>("SELECT count(*) AS n FROM chunks_trigram").get();
    return row?.n ?? 0;
  }

  /** Test/diagnostic helpers. */
  debugDbByteSize(): number { return this.dbByteSize(); }
  debugSourceCount(): number { return this.db.query<{ n: number }, []>("SELECT count(*) AS n FROM sources").get()?.n ?? 0; }
  debugBlobCount(): number { return this.db.query<{ n: number }, []>("SELECT count(*) AS n FROM blobs").get()?.n ?? 0; }
  debugOrphanChunkCount(): number {
    const row = this.db
      .query<{ n: number }, []>("SELECT count(*) AS n FROM chunks WHERE content_hash NOT IN (SELECT content_hash FROM sources)")
      .get();
    return row?.n ?? 0;
  }
}