server/db.ts

import { eq, desc, and, sql } from "drizzle-orm";
import { drizzle } from "drizzle-orm/better-sqlite3";
import Database from "better-sqlite3";
import path from "path";
import {
  InsertUser, users,
  sessions, InsertSession,
  messages, InsertMessage,
  settings, InsertSettings,
  toolPermissions, InsertToolPermission,
  costRecords, InsertCostRecord,
  mcpServers, InsertMcpServer,
} from "../drizzle/schema";

// ─── SQLite Database Setup ──────────────────────────────────────────────────

const DB_PATH = process.env.DATABASE_PATH || path.join(process.cwd(), "claw-web.db");

let _db: ReturnType<typeof drizzle> | null = null;
let _sqlite: InstanceType<typeof Database> | null = null;

export function getDb() {
  if (!_db) {
    _sqlite = new Database(DB_PATH);
    _sqlite.pragma("journal_mode = WAL");
    _sqlite.pragma("foreign_keys = ON");
    _db = drizzle(_sqlite);

    // Create tables if they don't exist
    _sqlite.exec(`
      CREATE TABLE IF NOT EXISTS users (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        openId TEXT NOT NULL UNIQUE,
        name TEXT,
        email TEXT,
        loginMethod TEXT,
        role TEXT NOT NULL DEFAULT 'user',
        createdAt TEXT NOT NULL DEFAULT (datetime('now')),
        updatedAt TEXT NOT NULL DEFAULT (datetime('now')),
        lastSignedIn TEXT NOT NULL DEFAULT (datetime('now'))
      );

      CREATE TABLE IF NOT EXISTS sessions (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        userId INTEGER NOT NULL,
        title TEXT NOT NULL DEFAULT 'New Session',
        model TEXT NOT NULL DEFAULT '',
        provider TEXT NOT NULL DEFAULT 'default',
        systemPrompt TEXT,
        isArchived INTEGER NOT NULL DEFAULT 0,
        createdAt TEXT NOT NULL DEFAULT (datetime('now')),
        updatedAt TEXT NOT NULL DEFAULT (datetime('now'))
      );

      CREATE TABLE IF NOT EXISTS messages (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        sessionId INTEGER NOT NULL,
        role TEXT NOT NULL,
        content TEXT,
        toolCalls TEXT,
        toolCallId TEXT,
        toolName TEXT,
        promptTokens INTEGER DEFAULT 0,
        completionTokens INTEGER DEFAULT 0,
        cost REAL DEFAULT 0,
        model TEXT,
        durationMs INTEGER DEFAULT 0,
        createdAt TEXT NOT NULL DEFAULT (datetime('now'))
      );

      CREATE TABLE IF NOT EXISTS settings (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        userId INTEGER NOT NULL UNIQUE,
        apiProvider TEXT NOT NULL DEFAULT 'default',
        apiKey TEXT,
        apiBaseUrl TEXT,
        model TEXT NOT NULL DEFAULT '',
        maxTokens INTEGER NOT NULL DEFAULT 16384,
        temperature REAL NOT NULL DEFAULT 0.7,
        topP REAL NOT NULL DEFAULT 1,
        systemPrompt TEXT,
        memoryContent TEXT,
        theme TEXT NOT NULL DEFAULT 'dark',
        expandToolCalls INTEGER NOT NULL DEFAULT 1,
        autoApproveSafeTools INTEGER NOT NULL DEFAULT 1,
        outputStyle TEXT NOT NULL DEFAULT 'markdown',
        createdAt TEXT NOT NULL DEFAULT (datetime('now')),
        updatedAt TEXT NOT NULL DEFAULT (datetime('now'))
      );

      CREATE TABLE IF NOT EXISTS toolPermissions (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        userId INTEGER NOT NULL,
        toolName TEXT NOT NULL,
        allowed INTEGER NOT NULL DEFAULT 1,
        requireConfirmation INTEGER NOT NULL DEFAULT 0,
        createdAt TEXT NOT NULL DEFAULT (datetime('now')),
        updatedAt TEXT NOT NULL DEFAULT (datetime('now'))
      );

      CREATE TABLE IF NOT EXISTS costRecords (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        userId INTEGER NOT NULL,
        sessionId INTEGER NOT NULL,
        model TEXT NOT NULL,
        promptTokens INTEGER NOT NULL DEFAULT 0,
        completionTokens INTEGER NOT NULL DEFAULT 0,
        totalCost REAL NOT NULL DEFAULT 0,
        createdAt TEXT NOT NULL DEFAULT (datetime('now'))
      );

      CREATE TABLE IF NOT EXISTS mcpServers (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        userId INTEGER NOT NULL,
        name TEXT NOT NULL,
        url TEXT NOT NULL,
        transport TEXT NOT NULL DEFAULT 'sse',
        enabled INTEGER NOT NULL DEFAULT 1,
        config TEXT,
        createdAt TEXT NOT NULL DEFAULT (datetime('now')),
        updatedAt TEXT NOT NULL DEFAULT (datetime('now'))
      );

      CREATE TABLE IF NOT EXISTS sessionState (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        sessionId INTEGER NOT NULL,
        stateKey TEXT NOT NULL,
        stateValue TEXT NOT NULL DEFAULT '{}',
        updatedAt TEXT NOT NULL DEFAULT (datetime('now')),
        UNIQUE(sessionId, stateKey)
      );
    `);

    // Create dev user if not exists
    const devUser = _sqlite.prepare("SELECT * FROM users WHERE openId = 'dev-user'").get();
    if (!devUser) {
      _sqlite.prepare(
        "INSERT INTO users (openId, name, email, role) VALUES ('dev-user', 'Developer', 'dev@claw.local', 'admin')"
      ).run();
      console.log("[DB] Created dev user");
    }
  }
  return _db;
}

// Initialize on import
getDb();

// ─── Users ───────────────────────────────────────────────────────────────────

export async function upsertUser(user: InsertUser): Promise<void> {
  if (!user.openId) throw new Error("User openId is required for upsert");
  const db = getDb();
  if (!db) return;

  const existing = await db.select().from(users).where(eq(users.openId, user.openId)).limit(1);
  if (existing.length > 0) {
    const updateSet: Record<string, unknown> = { updatedAt: new Date().toISOString() };
    if (user.name !== undefined) updateSet.name = user.name;
    if (user.email !== undefined) updateSet.email = user.email;
    if (user.loginMethod !== undefined) updateSet.loginMethod = user.loginMethod;
    if (user.role !== undefined) updateSet.role = user.role;
    if (user.lastSignedIn !== undefined) updateSet.lastSignedIn = user.lastSignedIn;
    await db.update(users).set(updateSet).where(eq(users.openId, user.openId));
  } else {
    await db.insert(users).values({
      openId: user.openId,
      name: user.name || null,
      email: user.email || null,
      loginMethod: user.loginMethod || null,
      role: user.role || "user",
    });
  }
}

export async function getUserByOpenId(openId: string) {
  const db = getDb();
  if (!db) return undefined;
  const result = await db.select().from(users).where(eq(users.openId, openId)).limit(1);
  return result.length > 0 ? result[0] : undefined;
}

// ─── Sessions ────────────────────────────────────────────────────────────────

export async function createSession(data: InsertSession) {
  const db = getDb();
  if (!db) throw new Error("DB not available");
  const result = await db.insert(sessions).values(data).returning();
  return result[0];
}

export async function getUserSessions(userId: number) {
  const db = getDb();
  if (!db) return [];
  return db.select().from(sessions)
    .where(and(eq(sessions.userId, userId), eq(sessions.isArchived, 0)))
    .orderBy(desc(sessions.updatedAt));
}

export async function getSession(sessionId: number, userId: number) {
  const db = getDb();
  if (!db) return undefined;
  const result = await db.select().from(sessions)
    .where(and(eq(sessions.id, sessionId), eq(sessions.userId, userId)))
    .limit(1);
  return result[0];
}

export async function updateSession(sessionId: number, userId: number, data: Partial<InsertSession>) {
  const db = getDb();
  if (!db) return;
  await db.update(sessions).set({ ...data, updatedAt: new Date().toISOString() })
    .where(and(eq(sessions.id, sessionId), eq(sessions.userId, userId)));
}

export async function deleteSession(sessionId: number, userId: number) {
  const db = getDb();
  if (!db) return;
  await db.delete(messages).where(eq(messages.sessionId, sessionId));
  await db.delete(sessions).where(and(eq(sessions.id, sessionId), eq(sessions.userId, userId)));
}

// ─── Messages ────────────────────────────────────────────────────────────────

export async function addMessage(data: InsertMessage) {
  const db = getDb();
  if (!db) throw new Error("DB not available");
  const result = await db.insert(messages).values(data).returning();
  return result[0];
}

export async function getSessionMessages(sessionId: number) {
  const db = getDb();
  if (!db) return [];
  return db.select().from(messages)
    .where(eq(messages.sessionId, sessionId))
    .orderBy(messages.id);
}

export async function deleteMessage(messageId: number) {
  const db = getDb();
  if (!db) return;
  await db.delete(messages).where(eq(messages.id, messageId));
}

export async function clearSessionMessages(sessionId: number) {
  const db = getDb();
  if (!db) return;
  await db.delete(messages).where(eq(messages.sessionId, sessionId));
}

// ─── Settings ────────────────────────────────────────────────────────────────

export async function getUserSettings(userId: number) {
  const db = getDb();
  if (!db) return null;
  const result = await db.select().from(settings).where(eq(settings.userId, userId)).limit(1);
  return result[0] || null;
}

export async function upsertSettings(userId: number, data: Partial<InsertSettings>) {
  const db = getDb();
  if (!db) return;
  const existing = await getUserSettings(userId);
  if (existing) {
    await db.update(settings).set({ ...data, updatedAt: new Date().toISOString() }).where(eq(settings.userId, userId));
  } else {
    await db.insert(settings).values({ userId, ...data });
  }
}

// ─── Tool Permissions ────────────────────────────────────────────────────────

export async function getToolPermissions(userId: number) {
  const db = getDb();
  if (!db) return [];
  return db.select().from(toolPermissions).where(eq(toolPermissions.userId, userId));
}

export async function upsertToolPermission(userId: number, toolName: string, allowed: number, requireConfirmation: number) {
  const db = getDb();
  if (!db) return;
  const existing = await db.select().from(toolPermissions)
    .where(and(eq(toolPermissions.userId, userId), eq(toolPermissions.toolName, toolName)))
    .limit(1);
  if (existing.length > 0) {
    await db.update(toolPermissions).set({ allowed, requireConfirmation, updatedAt: new Date().toISOString() })
      .where(and(eq(toolPermissions.userId, userId), eq(toolPermissions.toolName, toolName)));
  } else {
    await db.insert(toolPermissions).values({ userId, toolName, allowed, requireConfirmation });
  }
}

// ─── Cost Records ────────────────────────────────────────────────────────────

export async function addCostRecord(data: InsertCostRecord) {
  const db = getDb();
  if (!db) return;
  await db.insert(costRecords).values(data);
}

export async function getUserCostSummary(userId: number) {
  const db = getDb();
  if (!db) return { totalCost: 0, totalPromptTokens: 0, totalCompletionTokens: 0 };
  const result = await db.select({
    totalCost: sql<number>`COALESCE(SUM(totalCost), 0)`,
    totalPromptTokens: sql<number>`COALESCE(SUM(promptTokens), 0)`,
    totalCompletionTokens: sql<number>`COALESCE(SUM(completionTokens), 0)`,
  }).from(costRecords).where(eq(costRecords.userId, userId));
  return result[0] || { totalCost: 0, totalPromptTokens: 0, totalCompletionTokens: 0 };
}

export async function getSessionCostSummary(sessionId: number) {
  const db = getDb();
  if (!db) return { totalCost: 0, totalPromptTokens: 0, totalCompletionTokens: 0 };
  const result = await db.select({
    totalCost: sql<number>`COALESCE(SUM(totalCost), 0)`,
    totalPromptTokens: sql<number>`COALESCE(SUM(promptTokens), 0)`,
    totalCompletionTokens: sql<number>`COALESCE(SUM(completionTokens), 0)`,
  }).from(costRecords).where(eq(costRecords.sessionId, sessionId));
  return result[0] || { totalCost: 0, totalPromptTokens: 0, totalCompletionTokens: 0 };
}

// ─── MCP Servers ─────────────────────────────────────────────────────────────

export async function getMcpServers(userId: number) {
  const db = getDb();
  if (!db) return [];
  return db.select().from(mcpServers).where(eq(mcpServers.userId, userId));
}

export async function addMcpServer(data: InsertMcpServer) {
  const db = getDb();
  if (!db) throw new Error("DB not available");
  const result = await db.insert(mcpServers).values(data).returning();
  return result[0];
}

export async function updateMcpServer(id: number, userId: number, data: Partial<InsertMcpServer>) {
  const db = getDb();
  if (!db) return;
  await db.update(mcpServers).set({ ...data, updatedAt: new Date().toISOString() })
    .where(and(eq(mcpServers.id, id), eq(mcpServers.userId, userId)));
}

export async function deleteMcpServer(id: number, userId: number) {
  const db = getDb();
  if (!db) return;
  await db.delete(mcpServers).where(and(eq(mcpServers.id, id), eq(mcpServers.userId, userId)));
}

// ─── Session State Persistence ──────────────────────────────────────────────
// Persists in-memory state (todoLists, planModes, etc.) to SQLite so it
// survives server restarts. Uses raw SQLite for simplicity since sessionState
// table is not in the Drizzle schema.

export function saveSessionState(sessionId: number, key: string, value: unknown): void {
  try {
    if (!_sqlite) getDb();
    if (!_sqlite) return;
    const json = JSON.stringify(value);
    _sqlite.prepare(
      `INSERT INTO sessionState (sessionId, stateKey, stateValue, updatedAt)
       VALUES (?, ?, ?, datetime('now'))
       ON CONFLICT(sessionId, stateKey) DO UPDATE SET stateValue = excluded.stateValue, updatedAt = datetime('now')`
    ).run(sessionId, key, json);
  } catch (err: any) {
    console.error(`[db] Failed to save session state ${key}:`, err.message);
  }
}

export function loadSessionState<T>(sessionId: number, key: string): T | null {
  try {
    if (!_sqlite) getDb();
    if (!_sqlite) return null;
    const row = _sqlite.prepare(
      `SELECT stateValue FROM sessionState WHERE sessionId = ? AND stateKey = ?`
    ).get(sessionId, key) as { stateValue: string } | undefined;
    if (!row) return null;
    return JSON.parse(row.stateValue) as T;
  } catch (err: any) {
    console.error(`[db] Failed to load session state ${key}:`, err.message);
    return null;
  }
}

export function deleteSessionState(sessionId: number, key?: string): void {
  try {
    if (!_sqlite) getDb();
    if (!_sqlite) return;
    if (key) {
      _sqlite.prepare(`DELETE FROM sessionState WHERE sessionId = ? AND stateKey = ?`).run(sessionId, key);
    } else {
      _sqlite.prepare(`DELETE FROM sessionState WHERE sessionId = ?`).run(sessionId);
    }
  } catch (err: any) {
    console.error(`[db] Failed to delete session state:`, err.message);
  }
}

// ─── Graceful shutdown ────────────────────────────────────────────────────────

export function closeDb() {
  if (_sqlite) {
    try {
      _sqlite.close();
      console.log("[db] SQLite connection closed");
    } catch (err) {
      console.error("[db] Error closing SQLite:", err);
    }
    _sqlite = null;
    _db = null;
  }
}