CipherTalk/electron/services/ai/aiDatabase.ts

import Database from 'better-sqlite3'
import { existsSync, mkdirSync } from 'fs'
import { join, dirname } from 'path'

/**
 * AI 专用数据库管理
 */
export class AIDatabase {
  private db: Database.Database | null = null
  private dbPath: string | null = null

  /**
   * 初始化数据库
   */
  init(cachePath: string, wxid: string): void {
    // AI 数据库放在缓存根目录下，不放在账号目录内
    this.dbPath = join(cachePath, 'ai_summary.db')

    // 确保目录存在
    const dir = dirname(this.dbPath)
    if (!existsSync(dir)) {
      mkdirSync(dir, { recursive: true })
    }

    // 打开数据库
    this.db = new Database(this.dbPath)

    // 创建表
    this.createTables()
  }

  /**
   * 创建表结构
   */
  private createTables(): void {
    if (!this.db) return

    // 摘要记录表
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS summaries (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id TEXT NOT NULL,
        time_range_start INTEGER NOT NULL,
        time_range_end INTEGER NOT NULL,
        time_range_days INTEGER NOT NULL,
        message_count INTEGER NOT NULL,
        summary_text TEXT NOT NULL,
        tokens_used INTEGER,
        cost REAL,
        provider TEXT NOT NULL,
        model TEXT NOT NULL,
        created_at INTEGER NOT NULL
      );
      
      CREATE INDEX IF NOT EXISTS idx_summaries_session ON summaries(session_id);
      CREATE INDEX IF NOT EXISTS idx_summaries_created ON summaries(created_at);
      CREATE INDEX IF NOT EXISTS idx_summaries_time_range ON summaries(time_range_start, time_range_end);
    `)

    // 使用统计表
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS usage_stats (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        date TEXT NOT NULL,
        provider TEXT NOT NULL,
        model TEXT,
        total_tokens INTEGER DEFAULT 0,
        total_cost REAL DEFAULT 0,
        request_count INTEGER DEFAULT 0,
        UNIQUE(date, provider, model)
      );
      
      CREATE INDEX IF NOT EXISTS idx_usage_date ON usage_stats(date);
    `)

    // 缓存表
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS summary_cache (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        cache_key TEXT UNIQUE NOT NULL,
        summary_id INTEGER NOT NULL,
        expires_at INTEGER NOT NULL,
        FOREIGN KEY (summary_id) REFERENCES summaries(id)
      );
      
      CREATE INDEX IF NOT EXISTS idx_cache_key ON summary_cache(cache_key);
      CREATE INDEX IF NOT EXISTS idx_cache_expires ON summary_cache(expires_at);
    `)
    
    // 添加新列（如果不存在）
    try {
      this.db.exec("ALTER TABLE summaries ADD COLUMN prompt_text TEXT")
    } catch (e) {
      // 忽略错误，列已存在
    }
    
    try {
      this.db.exec("ALTER TABLE summaries ADD COLUMN custom_name TEXT")
    } catch (e) {
      // 忽略错误，列已存在
    }
  }

  /**
   * 获取数据库实例
   */
  getDb(): Database.Database {
    if (!this.db) {
      throw new Error('数据库未初始化')
    }
    return this.db
  }

  /**
   * 保存摘要
   */
  saveSummary(summary: {
    sessionId: string
    timeRangeStart: number
    timeRangeEnd: number
    timeRangeDays: number
    messageCount: number
    summaryText: string
    tokensUsed: number
    cost: number
    provider: string
    model: string
    promptText?: string
  }): number {
    const db = this.getDb()

    console.log('[AIDatabase] 保存摘要:', {
      sessionId: summary.sessionId,
      timeRangeDays: summary.timeRangeDays,
      messageCount: summary.messageCount,
      provider: summary.provider,
      model: summary.model
    })

    const result = db.prepare(`
      INSERT INTO summaries (
        session_id, time_range_start, time_range_end, time_range_days,
        message_count, summary_text, tokens_used, cost,
        provider, model, created_at, prompt_text
      ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    `).run(
      summary.sessionId,
      summary.timeRangeStart,
      summary.timeRangeEnd,
      summary.timeRangeDays,
      summary.messageCount,
      summary.summaryText,
      summary.tokensUsed,
      summary.cost,
      summary.provider,
      summary.model,
      Date.now(),
      summary.promptText || ''
    )

    console.log('[AIDatabase] 摘要已保存，ID:', result.lastInsertRowid)

    return result.lastInsertRowid as number
  }

  /**
   * 保存缓存
   */
  saveCache(cacheKey: string, summaryId: number, expiresAt: number): void {
    const db = this.getDb()

    db.prepare(`
      INSERT OR REPLACE INTO summary_cache (cache_key, summary_id, expires_at)
      VALUES (?, ?, ?)
    `).run(cacheKey, summaryId, expiresAt)
  }

  /**
   * 获取缓存的摘要
   */
  getCachedSummary(cacheKey: string): any | null {
    const db = this.getDb()
    const now = Date.now()

    const row: any = db.prepare(`
      SELECT s.* FROM summaries s
      JOIN summary_cache c ON s.id = c.summary_id
      WHERE c.cache_key = ? AND c.expires_at > ?
    `).get(cacheKey, now)

    if (!row) return null

    // 转换为 camelCase 格式
    return {
      id: row.id,
      sessionId: row.session_id,
      timeRangeStart: row.time_range_start,
      timeRangeEnd: row.time_range_end,
      timeRangeDays: row.time_range_days,
      messageCount: row.message_count,
      summaryText: row.summary_text,
      tokensUsed: row.tokens_used,
      cost: row.cost,
      provider: row.provider,
      model: row.model,
      createdAt: row.created_at,
      promptText: row.prompt_text
    }
  }

  /**
   * 更新使用统计
   */
  updateUsageStats(provider: string, model: string, tokens: number, cost: number): void {
    const db = this.getDb()
    const date = new Date().toISOString().split('T')[0] // YYYY-MM-DD

    db.prepare(`
      INSERT INTO usage_stats (date, provider, model, total_tokens, total_cost, request_count)
      VALUES (?, ?, ?, ?, ?, 1)
      ON CONFLICT(date, provider, model) DO UPDATE SET
        total_tokens = total_tokens + excluded.total_tokens,
        total_cost = total_cost + excluded.total_cost,
        request_count = request_count + 1
    `).run(date, provider, model, tokens, cost)
  }

  /**
   * 获取使用统计
   */
  getUsageStats(startDate?: string, endDate?: string): any[] {
    const db = this.getDb()

    let query = 'SELECT * FROM usage_stats'
    const params: any[] = []

    if (startDate && endDate) {
      query += ' WHERE date >= ? AND date <= ?'
      params.push(startDate, endDate)
    } else if (startDate) {
      query += ' WHERE date >= ?'
      params.push(startDate)
    } else if (endDate) {
      query += ' WHERE date <= ?'
      params.push(endDate)
    }

    query += ' ORDER BY date DESC'

    return db.prepare(query).all(...params)
  }

  /**
   * 获取会话的摘要历史
   */
  getSummaryHistory(sessionId: string, limit: number = 10): any[] {
    const db = this.getDb()

    console.log('[AIDatabase] 查询历史记录:', { sessionId, limit })

    const rows = db.prepare(`
      SELECT * FROM summaries
      WHERE session_id = ?
      ORDER BY created_at DESC
      LIMIT ?
    `).all(sessionId, limit)

    console.log('[AIDatabase] 查询到', rows.length, '条历史记录')

    // 转换为 camelCase 格式
    return rows.map((row: any) => ({
      id: row.id,
      sessionId: row.session_id,
      timeRangeStart: row.time_range_start,
      timeRangeEnd: row.time_range_end,
      timeRangeDays: row.time_range_days,
      messageCount: row.message_count,
      summaryText: row.summary_text,
      tokensUsed: row.tokens_used,
      cost: row.cost,
      provider: row.provider,
      model: row.model,
      createdAt: row.created_at,
      promptText: row.prompt_text,
      customName: row.custom_name
    }))
  }

  /**
   * 删除摘要
   */
  deleteSummary(id: number): boolean {
    const db = this.getDb()

    try {
      // 先删除相关的缓存
      db.prepare('DELETE FROM summary_cache WHERE summary_id = ?').run(id)
      
      // 删除摘要记录
      const result = db.prepare('DELETE FROM summaries WHERE id = ?').run(id)
      
      return result.changes > 0
    } catch (e) {
      console.error('[AIDatabase] 删除摘要失败:', e)
      return false
    }
  }

  /**
   * 重命名摘要
   */
  renameSummary(id: number, customName: string): boolean {
    const db = this.getDb()

    try {
      const result = db.prepare('UPDATE summaries SET custom_name = ? WHERE id = ?').run(customName, id)
      return result.changes > 0
    } catch (e) {
      console.error('[AIDatabase] 重命名摘要失败:', e)
      return false
    }
  }

  /**
   * 清理过期缓存
   */
  cleanExpiredCache(): void {
    const db = this.getDb()
    const now = Date.now()

    db.prepare('DELETE FROM summary_cache WHERE expires_at <= ?').run(now)
  }

  /**
   * 关闭数据库
   */
  close(): void {
    if (this.db) {
      this.db.close()
      this.db = null
    }
  }
}

// 导出单例
export const aiDatabase = new AIDatabase()