Nano-Banana-AI-Image-Editor/src/services/indexedDBService.ts

import { Generation, Edit } from '../types';

// 数据库配置
const DB_NAME = 'NanoBananaDB';
const DB_VERSION = 2; // 更新版本号
const GENERATIONS_STORE = 'generations';
const EDITS_STORE = 'edits';
const REFERENCE_IMAGES_STORE = 'referenceImages'; // 新增参考图像存储

// 重试配置
// const MAX_RETRIES = 3;
// const RETRY_DELAY = 1000;

// IndexedDB实例
let db: IDBDatabase | null = null;

/**
 * 初始化数据库
 */
export const initDB = (): Promise<void> => {
  return new Promise((resolve, reject) => {
    const request = indexedDB.open(DB_NAME, DB_VERSION);
    
    request.onerror = () => {
      console.error('数据库打开失败:', request.error);
      reject(request.error);
    };
    
    request.onsuccess = () => {
      db = request.result;
      resolve();
    };
    
    request.onupgradeneeded = (event) => {
      const db = (event.target as IDBOpenDBRequest).result;
      
      // 创建生成记录存储
      if (!db.objectStoreNames.contains(GENERATIONS_STORE)) {
        const genStore = db.createObjectStore(GENERATIONS_STORE, { keyPath: 'id' });
        genStore.createIndex('timestamp', 'timestamp', { unique: false });
      }
      
      // 创建编辑记录存储
      if (!db.objectStoreNames.contains(EDITS_STORE)) {
        const editStore = db.createObjectStore(EDITS_STORE, { keyPath: 'id' });
        editStore.createIndex('timestamp', 'timestamp', { unique: false });
        editStore.createIndex('parentGenerationId', 'parentGenerationId', { unique: false });
      }
      
      // 创建参考图像存储
      if (!db.objectStoreNames.contains(REFERENCE_IMAGES_STORE)) {
        const refImageStore = db.createObjectStore(REFERENCE_IMAGES_STORE, { keyPath: 'id' });
        refImageStore.createIndex('timestamp', 'timestamp', { unique: false });
      }
    };
  });
};

/**
 * 获取数据库实例
 */
const getDB = (): IDBDatabase => {
  if (!db) {
    throw new Error('数据库未初始化');
  }
  return db;
};

/**
 * 添加生成记录
 */
export const addGeneration = async (generation: Generation): Promise<void> => {
  // 创建轻量级生成记录，只存储必要的信息和上传后的URL
  const lightweightGeneration = {
    id: generation.id,
    prompt: generation.prompt,
    parameters: generation.parameters,
    modelVersion: generation.modelVersion,
    timestamp: generation.timestamp,
    uploadResults: generation.uploadResults,
    usageMetadata: generation.usageMetadata,
    // 只存储上传后的URL，不存储base64数据
    sourceAssets: generation.sourceAssets.map(asset => {
      const uploadedUrl = getUploadedAssetUrl(generation, asset.id, 'source');
      return {
        id: asset.id,
        type: asset.type,
        // 如果没有上传后的URL，则不存储URL以避免base64数据
        url: uploadedUrl || '',
        mime: asset.mime,
        width: asset.width,
        height: asset.height,
        checksum: asset.checksum
      };
    }),
    outputAssets: generation.outputAssets.map(asset => {
      const uploadedUrl = getUploadedAssetUrl(generation, asset.id, 'output');
      return {
        id: asset.id,
        type: asset.type,
        // 如果没有上传后的URL，则不存储URL以避免base64数据
        url: uploadedUrl || '',
        mime: asset.mime,
        width: asset.width,
        height: asset.height,
        checksum: asset.checksum
      };
    })
  };

  const db = getDB();
  const transaction = db.transaction([GENERATIONS_STORE], 'readwrite');
  const store = transaction.objectStore(GENERATIONS_STORE);
  
  return new Promise((resolve, reject) => {
    const request = store.add(lightweightGeneration);
    request.onsuccess = () => resolve();
    request.onerror = () => reject(request.error);
  });
};

/**
 * 添加编辑记录
 */
export const addEdit = async (edit: Edit): Promise<void> => {
  // 创建轻量级编辑记录，只存储必要的信息和上传后的URL
  const lightweightEdit = {
    id: edit.id,
    parentGenerationId: edit.parentGenerationId,
    maskAssetId: edit.maskAssetId,
    instruction: edit.instruction,
    timestamp: edit.timestamp,
    uploadResults: edit.uploadResults,
    parameters: edit.parameters,
    usageMetadata: edit.usageMetadata,
    // 只存储上传后的URL，不存储base64数据
    maskReferenceAsset: edit.maskReferenceAsset ? (() => {
      const uploadedUrl = getUploadedAssetUrl(edit, edit.maskReferenceAsset.id, 'mask');
      return {
        id: edit.maskReferenceAsset.id,
        type: edit.maskReferenceAsset.type,
        // 如果没有上传后的URL，则不存储URL以避免base64数据
        url: uploadedUrl || '',
        mime: edit.maskReferenceAsset.mime,
        width: edit.maskReferenceAsset.width,
        height: edit.maskReferenceAsset.height,
        checksum: edit.maskReferenceAsset.checksum
      };
    })() : undefined,
    outputAssets: edit.outputAssets.map(asset => {
      const uploadedUrl = getUploadedAssetUrl(edit, asset.id, 'output');
      return {
        id: asset.id,
        type: asset.type,
        // 如果没有上传后的URL，则不存储URL以避免base64数据
        url: uploadedUrl || '',
        mime: asset.mime,
        width: asset.width,
        height: asset.height,
        checksum: asset.checksum
      };
    })
  };

  const db = getDB();
  const transaction = db.transaction([EDITS_STORE], 'readwrite');
  const store = transaction.objectStore(EDITS_STORE);
  
  return new Promise((resolve, reject) => {
    const request = store.add(lightweightEdit);
    request.onsuccess = () => resolve();
    request.onerror = () => reject(request.error);
  });
};

/**
 * 从uploadResults中获取资产的上传后URL
 * 注意：这个函数需要根据资产在数组中的位置来匹配上传结果
 * - 输出资产的索引与uploadResults中的索引相对应
 * - 源资产（参考图像）的索引从outputAssets.length开始
 */
const getUploadedAssetUrl = (record: Generation | Edit, assetId: string, assetType: 'output' | 'source' | 'mask'): string | null => {
  if (!record.uploadResults || record.uploadResults.length === 0) {
    return null;
  }

  let assetIndex = -1;
  
  // 根据资产类型确定在uploadResults中的索引
  if (assetType === 'output') {
    // 输出资产的索引与在outputAssets数组中的索引相同
    assetIndex = record.outputAssets.findIndex(a => a.id === assetId);
  } else if (assetType === 'source') {
    // 源资产（参考图像）的索引从outputAssets.length开始
    assetIndex = record.sourceAssets?.findIndex(a => a.id === assetId);
    if (assetIndex >= 0) {
      assetIndex += record.outputAssets.length;
    }
  } else if (assetType === 'mask') {
    // 遮罩参考资产通常是第一个输出资产之后的第一个源资产
    assetIndex = record.outputAssets.length;
  }

  // 检查索引是否有效并且对应的上传结果是否存在且成功
  if (assetIndex >= 0 && assetIndex < record.uploadResults.length) {
    const uploadResult = record.uploadResults[assetIndex];
    if (uploadResult.success && uploadResult.url) {
      return uploadResult.url;
    }
  }

  return null;
};

/**
 * 获取所有生成记录（按时间倒序）
 */
export const getAllGenerations = async (): Promise<Generation[]> => {
  const db = getDB();
  const transaction = db.transaction([GENERATIONS_STORE], 'readonly');
  const store = transaction.objectStore(GENERATIONS_STORE);
  const index = store.index('timestamp');
  
  return new Promise((resolve, reject) => {
    const request = index.getAll();
    request.onsuccess = () => {
      // 按时间倒序排列
      const generations = request.result.sort((a, b) => b.timestamp - a.timestamp);
      resolve(generations);
    };
    request.onerror = () => reject(request.error);
  });
};

/**
 * 获取所有编辑记录（按时间倒序）
 */
export const getAllEdits = async (): Promise<Edit[]> => {
  const db = getDB();
  const transaction = db.transaction([EDITS_STORE], 'readonly');
  const store = transaction.objectStore(EDITS_STORE);
  const index = store.index('timestamp');
  
  return new Promise((resolve, reject) => {
    const request = index.getAll();
    request.onsuccess = () => {
      // 按时间倒序排列
      const edits = request.result.sort((a, b) => b.timestamp - a.timestamp);
      resolve(edits);
    };
    request.onerror = () => reject(request.error);
  });
};

/**
 * 根据父生成ID获取编辑记录
 */
export const getEditsByParentGenerationId = async (parentGenerationId: string): Promise<Edit[]> => {
  const db = getDB();
  const transaction = db.transaction([EDITS_STORE], 'readonly');
  const store = transaction.objectStore(EDITS_STORE);
  const index = store.index('parentGenerationId');
  
  return new Promise((resolve, reject) => {
    const request = index.getAll(IDBKeyRange.only(parentGenerationId));
    request.onsuccess = () => {
      // 按时间倒序排列
      const edits = request.result.sort((a, b) => b.timestamp - a.timestamp);
      resolve(edits);
    };
    request.onerror = () => reject(request.error);
  });
};

/**
 * 删除最旧的记录以保持限制
 */
export const cleanupOldRecords = async (limit: number = 100): Promise<void> => {
  const db = getDB();
  
  // 清理生成记录
  const genTransaction = db.transaction([GENERATIONS_STORE], 'readwrite');
  const genStore = genTransaction.objectStore(GENERATIONS_STORE);
  
  // 清理编辑记录
  const editTransaction = db.transaction([EDITS_STORE], 'readwrite');
  const editStore = editTransaction.objectStore(EDITS_STORE);
  
  // 获取所有记录并按时间排序
  const allGenerations = await getAllGenerations();
  const allEdits = await getAllEdits();
  
  // 计算需要删除的记录数量
  if (allGenerations.length > limit) {
    const toDelete = allGenerations.slice(limit);
    for (const gen of toDelete) {
      genStore.delete(gen.id);
    }
  }
  
  if (allEdits.length > limit) {
    const toDelete = allEdits.slice(limit);
    for (const edit of toDelete) {
      editStore.delete(edit.id);
    }
  }
};

/**
 * 清理记录中的base64数据
 */
export const cleanupBase64Data = async (): Promise<void> => {
  try {
    // 获取所有生成记录
    const generations = await getAllGenerations();
    
    // 获取所有编辑记录
    const edits = await getAllEdits();
    
    const db = getDB();
    
    // 更新生成记录
    for (const generation of generations) {
      // 检查是否有base64数据需要清理
      let needsUpdate = false;
      
      // 清理源资产中的base64数据
      const cleanedSourceAssets = generation.sourceAssets.map((asset) => {
        if (asset.url && asset.url.startsWith('data:')) {
          needsUpdate = true;
          return {
            ...asset,
            url: '' // 移除base64数据
          };
        }
        return asset;
      });
      
      // 清理输出资产中的base64数据
      const cleanedOutputAssets = generation.outputAssets.map((asset) => {
        if (asset.url && asset.url.startsWith('data:')) {
          needsUpdate = true;
          return {
            ...asset,
            url: '' // 移除base64数据
          };
        }
        return asset;
      });
      
      // 如果需要更新，则保存清理后的记录
      if (needsUpdate) {
        const cleanedGeneration = {
          ...generation,
          sourceAssets: cleanedSourceAssets,
          outputAssets: cleanedOutputAssets
        };
        
        const transaction = db.transaction([GENERATIONS_STORE], 'readwrite');
        const store = transaction.objectStore(GENERATIONS_STORE);
        await new Promise((resolve, reject) => {
          const request = store.put(cleanedGeneration);
          request.onsuccess = () => resolve(undefined);
          request.onerror = () => reject(request.error);
        });
      }
    }
    
    // 更新编辑记录
    for (const edit of edits) {
      // 检查是否有base64数据需要清理
      let needsUpdate = false;
      
      // 清理遮罩参考资产中的base64数据
      let cleanedMaskReferenceAsset = edit.maskReferenceAsset;
      if (edit.maskReferenceAsset && edit.maskReferenceAsset.url && edit.maskReferenceAsset.url.startsWith('data:')) {
        needsUpdate = true;
        cleanedMaskReferenceAsset = {
          ...edit.maskReferenceAsset,
          url: '' // 移除base64数据
        };
      }
      
      // 清理输出资产中的base64数据
      const cleanedOutputAssets = edit.outputAssets.map((asset) => {
        if (asset.url && asset.url.startsWith('data:')) {
          needsUpdate = true;
          return {
            ...asset,
            url: '' // 移除base64数据
          };
        }
        return asset;
      });
      
      // 如果需要更新，则保存清理后的记录
      if (needsUpdate) {
        const cleanedEdit = {
          ...edit,
          maskReferenceAsset: cleanedMaskReferenceAsset,
          outputAssets: cleanedOutputAssets
        };
        
        const transaction = db.transaction([EDITS_STORE], 'readwrite');
        const store = transaction.objectStore(EDITS_STORE);
        await new Promise((resolve, reject) => {
          const request = store.put(cleanedEdit);
          request.onsuccess = () => resolve(undefined);
          request.onerror = () => reject(request.error);
        });
      }
    }
    
    console.log('IndexedDB中的base64数据清理完成');
  } catch (error) {
    console.error('清理IndexedDB中的base64数据时出错:', error);
  }
};

/**
 * 删除指定的生成记录
 */
export const deleteGeneration = async (id: string): Promise<void> => {
  const db = getDB();
  const transaction = db.transaction([GENERATIONS_STORE], 'readwrite');
  const store = transaction.objectStore(GENERATIONS_STORE);
  
  return new Promise((resolve, reject) => {
    const request = store.delete(id);
    request.onsuccess = () => resolve();
    request.onerror = () => reject(request.error);
  });
};

/**
 * 删除指定的编辑记录
 */
export const deleteEdit = async (id: string): Promise<void> => {
  const db = getDB();
  const transaction = db.transaction([EDITS_STORE], 'readwrite');
  const store = transaction.objectStore(EDITS_STORE);
  
  return new Promise((resolve, reject) => {
    const request = store.delete(id);
    request.onsuccess = () => resolve();
    request.onerror = () => reject(request.error);
  });
};

/**
 * 批量删除生成记录
 */
export const deleteGenerations = async (ids: string[]): Promise<void> => {
  const db = getDB();
  const transaction = db.transaction([GENERATIONS_STORE], 'readwrite');
  const store = transaction.objectStore(GENERATIONS_STORE);
  
  const promises = ids.map(id => {
    return new Promise<void>((resolve, reject) => {
      const request = store.delete(id);
      request.onsuccess = () => resolve();
      request.onerror = () => reject(request.error);
    });
  });
  
  return Promise.all(promises).then(() => undefined);
};

/**
 * 批量删除编辑记录
 */
export const deleteEdits = async (ids: string[]): Promise<void> => {
  const db = getDB();
  const transaction = db.transaction([EDITS_STORE], 'readwrite');
  const store = transaction.objectStore(EDITS_STORE);
  
  const promises = ids.map(id => {
    return new Promise<void>((resolve, reject) => {
      const request = store.delete(id);
      request.onsuccess = () => resolve();
      request.onerror = () => reject(request.error);
    });
  });
  
  return Promise.all(promises).then(() => undefined);
};

/**
 * 添加参考图像
 */
export const addReferenceImage = async (image: { id: string; data: string; timestamp: number }): Promise<void> => {
  const db = getDB();
  const transaction = db.transaction([REFERENCE_IMAGES_STORE], 'readwrite');
  const store = transaction.objectStore(REFERENCE_IMAGES_STORE);
  
  return new Promise((resolve, reject) => {
    const request = store.add(image);
    request.onsuccess = () => resolve();
    request.onerror = () => reject(request.error);
  });
};

/**
 * 获取所有参考图像
 */
export const getAllReferenceImages = async (): Promise<Array<{ id: string; data: string; timestamp: number }>> => {
  const db = getDB();
  const transaction = db.transaction([REFERENCE_IMAGES_STORE], 'readonly');
  const store = transaction.objectStore(REFERENCE_IMAGES_STORE);
  const index = store.index('timestamp');
  
  return new Promise((resolve, reject) => {
    const request = index.getAll();
    request.onsuccess = () => {
      // 按时间倒序排列
      const images = request.result.sort((a, b) => b.timestamp - a.timestamp);
      resolve(images);
    };
    request.onerror = () => reject(request.error);
  });
};

/**
 * 根据ID获取参考图像
 */
export const getReferenceImageById = async (id: string): Promise<{ id: string; data: string; timestamp: number } | undefined> => {
  const db = getDB();
  const transaction = db.transaction([REFERENCE_IMAGES_STORE], 'readonly');
  const store = transaction.objectStore(REFERENCE_IMAGES_STORE);
  
  return new Promise((resolve, reject) => {
    const request = store.get(id);
    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error);
  });
};

/**
 * 删除参考图像
 */
export const deleteReferenceImage = async (id: string): Promise<void> => {
  const db = getDB();
  const transaction = db.transaction([REFERENCE_IMAGES_STORE], 'readwrite');
  const store = transaction.objectStore(REFERENCE_IMAGES_STORE);
  
  return new Promise((resolve, reject) => {
    const request = store.delete(id);
    request.onsuccess = () => resolve();
    request.onerror = () => reject(request.error);
  });
};

/**
 * 清空所有参考图像
 */
export const clearAllReferenceImages = async (): Promise<void> => {
  const db = getDB();
  const transaction = db.transaction([REFERENCE_IMAGES_STORE], 'readwrite');
  const store = transaction.objectStore(REFERENCE_IMAGES_STORE);
  
  return new Promise((resolve, reject) => {
    const request = store.clear();
    request.onsuccess = () => resolve();
    request.onerror = () => reject(request.error);
  });
};

/**
 * 清空所有记录
 */
export const clearAllRecords = async (): Promise<void> => {
  const db = getDB();
  
  const genTransaction = db.transaction([GENERATIONS_STORE], 'readwrite');
  const genStore = genTransaction.objectStore(GENERATIONS_STORE);
  
  const editTransaction = db.transaction([EDITS_STORE], 'readwrite');
  const editStore = editTransaction.objectStore(EDITS_STORE);
  
  const refImageTransaction = db.transaction([REFERENCE_IMAGES_STORE], 'readwrite');
  const refImageStore = refImageTransaction.objectStore(REFERENCE_IMAGES_STORE);
  
  return Promise.all([
    new Promise<void>((resolve, reject) => {
      const request = genStore.clear();
      request.onsuccess = () => resolve();
      request.onerror = () => reject(request.error);
    }),
    new Promise<void>((resolve, reject) => {
      const request = editStore.clear();
      request.onsuccess = () => resolve();
      request.onerror = () => reject(request.error);
    }),
    new Promise<void>((resolve, reject) => {
      const request = refImageStore.clear();
      request.onsuccess = () => resolve();
      request.onerror = () => reject(request.error);
    })
  ]).then(() => undefined);
};
export const closeDB = (): void => {
  if (db) {
    db.close();
    db = null;
  }
};