photoai/src/main/indexDb.ts

import { app } from 'electron'
import initSqlJs, { type Database, type SqlJsStatic } from 'sql.js'
import { readFile, writeFile, mkdir } from 'node:fs/promises'
import { existsSync, readFileSync } from 'node:fs'
import { join, dirname } from 'node:path'
import type {
  AssetRecord,
  QualityScores,
  IndexedAsset,
  AssetQuery,
  QualityThresholds,
  ColorLabel
} from '@shared/types'
import { SUPPORTED_EXTENSIONS, SUPPORTED_VIDEO_EXTENSIONS } from '@shared/constants'
import { logger } from './logger'

/**
 * 라이브러리 인덱스 DB. WASM SQLite(sql.js) 사용 — 네이티브 빌드/ABI 재빌드 불필요.
 * sql.js는 인메모리 → 변경분을 주기적으로 파일(userData/index.db)로 export 하여 영속화.
 * (수천~수만 장 메타데이터 규모에 충분. 대규모/임베딩은 Phase 2+에서 별도 전략.)
 */
class IndexDb {
  private SQL: SqlJsStatic | null = null
  private db: Database | null = null
  private dbPath = ''
  private dirty = false

  async init(): Promise<void> {
    if (this.db) return

    // sql.js wasm 바이트를 직접 읽어 전달 (asar/패키징 환경에서도 안전).
    // wasmBinary는 ArrayBuffer를 기대하므로 Buffer → ArrayBuffer 변환.
    const wasmPath = join(app.getAppPath(), 'node_modules', 'sql.js', 'dist', 'sql-wasm.wasm')
    const wasmBuf = readFileSync(wasmPath)
    this.SQL = await initSqlJs({ wasmBinary: new Uint8Array(wasmBuf).buffer })

    this.dbPath = join(app.getPath('userData'), 'index.db')
    const existing = existsSync(this.dbPath) ? await readFile(this.dbPath) : undefined
    this.db = new this.SQL.Database(existing)
    this.migrate()
    await this.save()

    logger.info('인덱스 DB 준비', { path: this.dbPath, assets: this.count() })
  }

  private migrate(): void {
    this.db!.run(`
      CREATE TABLE IF NOT EXISTS asset (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        contentHash TEXT UNIQUE NOT NULL,
        path TEXT NOT NULL,
        ext TEXT,
        sizeBytes INTEGER,
        mtime INTEGER,
        width INTEGER,
        height INTEGER,
        exifYear TEXT,
        exifMonth TEXT,
        gpsLat REAL,
        gpsLon REAL,
        camera TEXT,
        indexedAt INTEGER
      );
      CREATE TABLE IF NOT EXISTS quality (
        assetId INTEGER PRIMARY KEY REFERENCES asset(id) ON DELETE CASCADE,
        focus REAL,
        exposure REAL,
        eyesOpen REAL,
        flag TEXT
      );
      CREATE TABLE IF NOT EXISTS usermeta (
        assetId INTEGER PRIMARY KEY REFERENCES asset(id) ON DELETE CASCADE,
        rating INTEGER DEFAULT 0,
        label TEXT
      );
      CREATE TABLE IF NOT EXISTS embedding (
        assetId INTEGER PRIMARY KEY REFERENCES asset(id) ON DELETE CASCADE,
        vec BLOB
      );
      CREATE INDEX IF NOT EXISTS idx_asset_hash ON asset(contentHash);
      CREATE INDEX IF NOT EXISTS idx_asset_path ON asset(path);
    `)

    // 기존 DB(컬럼 없음)에 대한 마이그레이션 — ALTER ADD COLUMN
    this.ensureColumn('asset', 'gpsLat', 'REAL')
    this.ensureColumn('asset', 'gpsLon', 'REAL')
    this.ensureColumn('asset', 'camera', 'TEXT')
    // metaVersion: 확장 메타(GPS/카메라) 적재 버전. 구버전 행(0/NULL)은 재색인 시 backfill
    this.ensureColumn('asset', 'metaVersion', 'INTEGER')
  }

  /** 테이블에 컬럼이 없으면 추가 (sql.js는 ADD COLUMN IF NOT EXISTS 미지원) */
  private ensureColumn(table: string, col: string, type: string): void {
    const res = this.db!.exec(`PRAGMA table_info(${table})`)
    const names = res.length ? res[0].values.map((r) => String(r[1])) : []
    if (!names.includes(col)) {
      this.db!.run(`ALTER TABLE ${table} ADD COLUMN ${col} ${type}`)
    }
  }

  /** 인메모리 DB를 디스크로 영속화 */
  async save(): Promise<void> {
    if (!this.db) return
    const data = this.db.export()
    await mkdir(dirname(this.dbPath), { recursive: true })
    await writeFile(this.dbPath, Buffer.from(data))
    this.dirty = false
  }

  /** 변경이 있을 때만 저장 (배치 색인 후 호출) */
  async saveIfDirty(): Promise<void> {
    if (this.dirty) await this.save()
  }

  count(): number {
    const res = this.db!.exec('SELECT COUNT(*) AS n FROM asset')
    return res.length ? Number(res[0].values[0][0]) : 0
  }

  /**
   * 같은 경로·mtime으로 이미 **확장 메타까지** 색인되었으면 true (빠른 스킵).
   * metaVersion이 없는 구버전 행은 false → 재색인 시 GPS/카메라 backfill.
   */
  isIndexedPath(path: string, mtime: number): boolean {
    const stmt = this.db!.prepare(
      'SELECT 1 FROM asset WHERE path = ? AND mtime = ? AND metaVersion >= 1 LIMIT 1'
    )
    try {
      stmt.bind([path, mtime])
      return stmt.step()
    } finally {
      stmt.free()
    }
  }

  /** 이미 색인되었고 mtime이 동일하면 재색인 불필요 */
  needsIndex(contentHash: string, mtime: number): boolean {
    const stmt = this.db!.prepare('SELECT mtime FROM asset WHERE contentHash = ?')
    try {
      stmt.bind([contentHash])
      if (!stmt.step()) return true // 미존재 → 색인 필요
      const row = stmt.getAsObject() as { mtime: number }
      return row.mtime !== mtime
    } finally {
      stmt.free()
    }
  }

  /**
   * 최근 색인순 자산 목록 + 품질 + 사용자메타. 품질 플래그는 임계값으로 **실시간 계산**
   * (임계값 변경 시 재분석 없이 즉시 반영). 별점/색라벨 필터 지원.
   */
  /** 자산+품질(실시간 플래그)+사용자메타를 결합하는 공통 SELECT */
  private innerSelect(th: QualityThresholds): string {
    const f = Number(th.focus)
    const x = Number(th.exposure)
    const e = Number(th.eyes)
    return `
      SELECT a.*,
        q.focus AS focus, q.exposure AS exposure, q.eyesOpen AS eyesOpen,
        COALESCE(um.rating, 0) AS rating, um.label AS label,
        CASE
          WHEN q.assetId IS NULL THEN NULL
          WHEN q.focus < ${f} THEN 'blurry'
          WHEN q.eyesOpen IS NOT NULL AND q.eyesOpen < ${e} THEN 'eyesClosed'
          WHEN q.exposure < ${x} THEN 'badExposure'
          ELSE 'candidate'
        END AS flag
      FROM asset a
      LEFT JOIN quality q ON q.assetId = a.id
      LEFT JOIN usermeta um ON um.assetId = a.id`
  }

  /** 쿼리 → (where절, 바인딩 파라미터). 고정 열거/숫자는 인라인, 사용자 값은 바인딩 */
  private buildWhere(query: AssetQuery): { where: string; params: (string | number)[] } {
    const conds: string[] = []
    const params: (string | number)[] = []
    if (query.filter === 'rejected') {
      conds.push("flag IN ('blurry', 'eyesClosed', 'badExposure')")
    } else if (query.filter !== 'all') {
      conds.push(`flag = '${query.filter}'`)
    }
    if (query.ratingMin > 0) conds.push(`rating >= ${Number(query.ratingMin)}`)
    if (query.kind === 'image') {
      conds.push(`ext IN (${SUPPORTED_EXTENSIONS.map((e) => `'${e}'`).join(',')})`)
    } else if (query.kind === 'video') {
      conds.push(`ext IN (${SUPPORTED_VIDEO_EXTENSIONS.map((e) => `'${e}'`).join(',')})`)
    }
    if (query.year) {
      conds.push('exifYear = ?')
      params.push(query.year)
    }
    if (query.camera) {
      conds.push('camera = ?')
      params.push(query.camera)
    }
    if (query.label) {
      conds.push('label = ?')
      params.push(query.label)
    }
    return { where: conds.length ? `WHERE ${conds.join(' AND ')}` : '', params }
  }

  listAssets(
    offset: number,
    limit: number,
    query: AssetQuery,
    th: QualityThresholds
  ): IndexedAsset[] {
    const inner = this.innerSelect(th)
    const { where, params } = this.buildWhere(query)
    const stmt = this.db!.prepare(
      `SELECT * FROM (${inner}) ${where} ORDER BY indexedAt DESC, id DESC LIMIT ? OFFSET ?`
    )
    const out: IndexedAsset[] = []
    try {
      stmt.bind([...params, limit, offset])
      while (stmt.step()) out.push(stmt.getAsObject() as unknown as IndexedAsset)
    } finally {
      stmt.free()
    }
    return out
  }

  /** 쿼리에 매칭되는 전체 자산 id (전체 선택/필터 전체 내보내기용) */
  listAssetIds(query: AssetQuery, th: QualityThresholds): number[] {
    const inner = this.innerSelect(th)
    const { where, params } = this.buildWhere(query)
    const stmt = this.db!.prepare(
      `SELECT id FROM (${inner}) ${where} ORDER BY indexedAt DESC, id DESC`
    )
    const out: number[] = []
    try {
      stmt.bind(params)
      while (stmt.step()) out.push(Number((stmt.getAsObject() as { id: number }).id))
    } finally {
      stmt.free()
    }
    return out
  }

  /** 컬렉션 패싯 집계 (연도/카메라/색라벨) */
  facets(): import('@shared/types').Facets {
    const q = (sql: string): import('@shared/types').FacetItem[] => {
      const res = this.db!.exec(sql)
      if (!res.length) return []
      return res[0].values.map((row) => ({ value: String(row[0]), count: Number(row[1]) }))
    }
    return {
      years: q(
        "SELECT exifYear, COUNT(*) FROM asset WHERE exifYear IS NOT NULL GROUP BY exifYear ORDER BY exifYear DESC"
      ),
      cameras: q(
        "SELECT camera, COUNT(*) FROM asset WHERE camera IS NOT NULL AND camera <> '' GROUP BY camera ORDER BY COUNT(*) DESC"
      ),
      labels: q(
        'SELECT label, COUNT(*) FROM usermeta WHERE label IS NOT NULL GROUP BY label ORDER BY COUNT(*) DESC'
      )
    }
  }

  setRating(assetId: number, rating: number): void {
    const r = Math.max(0, Math.min(5, Math.round(rating)))
    this.db!.run(
      `INSERT INTO usermeta (assetId, rating) VALUES (?, ?)
       ON CONFLICT(assetId) DO UPDATE SET rating = excluded.rating`,
      [assetId, r]
    )
    this.dirty = true
  }

  setLabel(assetId: number, label: ColorLabel): void {
    this.db!.run(
      `INSERT INTO usermeta (assetId, label) VALUES (?, ?)
       ON CONFLICT(assetId) DO UPDATE SET label = excluded.label`,
      [assetId, label]
    )
    this.dirty = true
  }

  // ---- 임베딩 / 검색 (Phase 2) ----

  /** 임베딩 미보유 이미지(영상 제외) 목록 — 검색 색인 생성용 */
  listAssetsNeedingEmbedding(imageExts: string[]): { id: number; path: string }[] {
    const placeholders = imageExts.map(() => '?').join(',')
    const stmt = this.db!.prepare(
      `SELECT a.id AS id, a.path AS path
       FROM asset a LEFT JOIN embedding e ON e.assetId = a.id
       WHERE e.assetId IS NULL AND a.ext IN (${placeholders})
       ORDER BY a.id`
    )
    const out: { id: number; path: string }[] = []
    try {
      stmt.bind(imageExts)
      while (stmt.step()) out.push(stmt.getAsObject() as unknown as { id: number; path: string })
    } finally {
      stmt.free()
    }
    return out
  }

  setEmbedding(assetId: number, vec: number[]): void {
    const bytes = new Uint8Array(new Float32Array(vec).buffer)
    this.db!.run(
      `INSERT INTO embedding (assetId, vec) VALUES (?, ?)
       ON CONFLICT(assetId) DO UPDATE SET vec = excluded.vec`,
      [assetId, bytes]
    )
    this.dirty = true
  }

  /** GPS 좌표가 있는 자산(지도 마커용) */
  assetsWithGps(): { id: number; contentHash: string; path: string; gpsLat: number; gpsLon: number }[] {
    const stmt = this.db!.prepare(
      'SELECT id, contentHash, path, gpsLat, gpsLon FROM asset WHERE gpsLat IS NOT NULL AND gpsLon IS NOT NULL'
    )
    const out: { id: number; contentHash: string; path: string; gpsLat: number; gpsLon: number }[] =
      []
    try {
      while (stmt.step()) {
        const r = stmt.getAsObject() as unknown as {
          id: number
          contentHash: string
          path: string
          gpsLat: number
          gpsLon: number
        }
        out.push(r)
      }
    } finally {
      stmt.free()
    }
    return out
  }

  /** 특정 자산의 임베딩 (연관 탐색용) */
  getEmbedding(assetId: number): Float32Array | null {
    const stmt = this.db!.prepare('SELECT vec FROM embedding WHERE assetId = ?')
    try {
      stmt.bind([assetId])
      if (!stmt.step()) return null
      const u8 = (stmt.getAsObject() as { vec: Uint8Array }).vec
      return new Float32Array(u8.buffer, u8.byteOffset, Math.floor(u8.byteLength / 4))
    } finally {
      stmt.free()
    }
  }

  /** mtime이 ±window 이내인 자산 id (시간 연관용) */
  assetsNearTime(mtime: number, windowMs: number, excludeId: number, limit = 200): number[] {
    const stmt = this.db!.prepare(
      'SELECT id FROM asset WHERE mtime BETWEEN ? AND ? AND id <> ? LIMIT ?'
    )
    const out: number[] = []
    try {
      stmt.bind([mtime - windowMs, mtime + windowMs, excludeId, limit])
      while (stmt.step()) out.push(Number((stmt.getAsObject() as { id: number }).id))
    } finally {
      stmt.free()
    }
    return out
  }

  embeddingCount(): number {
    const res = this.db!.exec('SELECT COUNT(*) AS n FROM embedding')
    return res.length ? Number(res[0].values[0][0]) : 0
  }

  /** 전체 임베딩 로드 (브루트포스 코사인 검색용) */
  getAllEmbeddings(): { assetId: number; vec: Float32Array }[] {
    const stmt = this.db!.prepare('SELECT assetId, vec FROM embedding')
    const out: { assetId: number; vec: Float32Array }[] = []
    try {
      while (stmt.step()) {
        const row = stmt.getAsObject() as { assetId: number; vec: Uint8Array }
        const u8 = row.vec
        const vec = new Float32Array(u8.buffer, u8.byteOffset, Math.floor(u8.byteLength / 4))
        out.push({ assetId: Number(row.assetId), vec })
      }
    } finally {
      stmt.free()
    }
    return out
  }

  /** id 목록을 입력 순서대로 IndexedAsset으로 조회 (검색 결과 정렬 유지) */
  assetsByIds(ids: number[], th: QualityThresholds): IndexedAsset[] {
    if (ids.length === 0) return []
    const placeholders = ids.map(() => '?').join(',')
    const stmt = this.db!.prepare(
      `SELECT * FROM (${this.innerSelect(th)}) WHERE id IN (${placeholders})`
    )
    const byId = new Map<number, IndexedAsset>()
    try {
      stmt.bind(ids)
      while (stmt.step()) {
        const a = stmt.getAsObject() as unknown as IndexedAsset
        if (a.id != null) byId.set(a.id, a)
      }
    } finally {
      stmt.free()
    }
    return ids.map((id) => byId.get(id)).filter((a): a is IndexedAsset => !!a)
  }

  getById(id: number): AssetRecord | null {
    const stmt = this.db!.prepare('SELECT * FROM asset WHERE id = ?')
    try {
      stmt.bind([id])
      if (!stmt.step()) return null
      return stmt.getAsObject() as unknown as AssetRecord
    } finally {
      stmt.free()
    }
  }

  /** 자산 + 연관 메타(품질/사용자메타/임베딩) 전부 삭제 */
  deleteAsset(id: number): void {
    this.db!.run('DELETE FROM embedding WHERE assetId = ?', [id])
    this.db!.run('DELETE FROM quality WHERE assetId = ?', [id])
    this.db!.run('DELETE FROM usermeta WHERE assetId = ?', [id])
    this.db!.run('DELETE FROM asset WHERE id = ?', [id])
    this.dirty = true
  }

  getByHash(contentHash: string): AssetRecord | null {
    const stmt = this.db!.prepare('SELECT * FROM asset WHERE contentHash = ?')
    try {
      stmt.bind([contentHash])
      if (!stmt.step()) return null
      return stmt.getAsObject() as unknown as AssetRecord
    } finally {
      stmt.free()
    }
  }

  /** 자산 upsert (contentHash 기준). 반환: asset id */
  upsertAsset(r: AssetRecord): number {
    this.db!.run(
      `INSERT INTO asset
         (contentHash, path, ext, sizeBytes, mtime, width, height, exifYear, exifMonth,
          gpsLat, gpsLon, camera, metaVersion, indexedAt)
       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?)
       ON CONFLICT(contentHash) DO UPDATE SET
         path=excluded.path, ext=excluded.ext, sizeBytes=excluded.sizeBytes,
         mtime=excluded.mtime, width=excluded.width, height=excluded.height,
         exifYear=excluded.exifYear, exifMonth=excluded.exifMonth,
         gpsLat=excluded.gpsLat, gpsLon=excluded.gpsLon, camera=excluded.camera,
         metaVersion=1, indexedAt=excluded.indexedAt`,
      [
        r.contentHash,
        r.path,
        r.ext,
        r.sizeBytes,
        r.mtime,
        r.width,
        r.height,
        r.exifYear,
        r.exifMonth,
        r.gpsLat,
        r.gpsLon,
        r.camera,
        r.indexedAt
      ]
    )
    this.dirty = true
    const res = this.db!.exec('SELECT id FROM asset WHERE contentHash = ?', [r.contentHash])
    return res.length ? Number(res[0].values[0][0]) : -1
  }

  setQuality(assetId: number, q: QualityScores): void {
    this.db!.run(
      `INSERT INTO quality (assetId, focus, exposure, eyesOpen, flag)
       VALUES (?, ?, ?, ?, ?)
       ON CONFLICT(assetId) DO UPDATE SET
         focus=excluded.focus, exposure=excluded.exposure,
         eyesOpen=excluded.eyesOpen, flag=excluded.flag`,
      [assetId, q.focus, q.exposure, q.eyesOpen, q.flag]
    )
    this.dirty = true
  }

  close(): void {
    this.db?.close()
    this.db = null
  }
}

export const indexDb = new IndexDb()