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

import { GoogleGenAI } from '@google/genai'

// 注意：在生产环境中，这应该通过后端代理处理
// 优先使用localStorage中的API密钥，如果没有则使用环境变量中的，最后使用默认值
const API_KEY = localStorage.getItem('VITE_GEMINI_API_KEY') || import.meta.env.VITE_GEMINI_API_KEY || 'demo-key'
const genAI = new GoogleGenAI({ apiKey: API_KEY })

export interface GenerationRequest {
  prompt: string
  referenceImages?: Blob[] // Blob数组
  temperature?: number
  seed?: number
  // 添加abortSignal参数
  abortSignal?: AbortSignal
}

export interface EditRequest {
  instruction: string
  originalImage: Blob // Blob
  referenceImages?: Blob[] // Blob数组
  maskImage?: Blob // Blob
  temperature?: number
  seed?: number
  // 添加abortSignal参数
  abortSignal?: AbortSignal
}

export interface UsageMetadata {
  totalTokenCount?: number
  promptTokenCount?: number
  candidatesTokenCount?: number
}

export interface SegmentationRequest {
  image: Blob // Blob
  query: string // "像素(x,y)处的对象" 或 "红色汽车"
  // 添加abortSignal参数
  abortSignal?: AbortSignal
}

export class GeminiService {
  // 缓存base64图像数据，确保它们不会被清除
  private base64ImagesCache: Map<string, string> = new Map()

  // 将Blob转换为base64的辅助函数
  private async blobToBase64(blob: Blob): Promise<string> {
    return new Promise((resolve, reject) => {
      const reader = new FileReader()
      reader.onload = () => {
        const result = reader.result as string
        const base64 = result.split(',')[1] // Remove data:image/png;base64, prefix
        resolve(base64)
      }
      reader.onerror = reject
      reader.readAsDataURL(blob)
    })
  }

  // 生成Blob的唯一标识符
  private async generateBlobId(blob: Blob): Promise<string> {
    // 使用Blob的部分内容生成唯一标识符
    const arrayBuffer = await blob.slice(0, 1024).arrayBuffer()
    const uint8Array = new Uint8Array(arrayBuffer)
    let hash = ''
    for (let i = 0; i < uint8Array.length; i++) {
      hash += uint8Array[i].toString(16).padStart(2, '0')
    }
    return `${blob.type}-${blob.size}-${hash.substring(0, 32)}`
  }

  // 清理过期的缓存项（可选）
  private cleanupExpiredCache(): void {
    // 在这个实现中，我们不自动清理缓存
    // 只有在显式调用clearBase64Cache时才清理
    console.log('缓存大小:', this.base64ImagesCache.size)
  }

  async generateImage(request: GenerationRequest): Promise<{ images: Blob[]; usageMetadata?: UsageMetadata }> {
    try {
      const contents: Array<{ text: string } | { inlineData: { mimeType: string; data: string } }> = [{ text: request.prompt }]

      // 如果提供了参考图像则添加
      if (request.referenceImages && request.referenceImages.length > 0) {
        // 将Blob转换为base64以发送到API
        const base64Images: string[] = []

        // 为每个参考图像生成或获取base64数据
        for (const blob of request.referenceImages) {
          // 生成Blob的唯一标识符
          const blobId = await this.generateBlobId(blob)

          let base64: string
          // 检查缓存中是否已有该图像的base64数据
          if (this.base64ImagesCache.has(blobId)) {
            // 从缓存中获取base64数据
            base64 = this.base64ImagesCache.get(blobId)!
            console.log('从缓存中获取参考图像base64数据')
          } else {
            // 转换Blob为base64并缓存结果
            base64 = await this.blobToBase64(blob)
            // 将base64数据存储到缓存中
            this.base64ImagesCache.set(blobId, base64)
            console.log('生成并缓存参考图像base64数据')
          }

          // 如果base64数据为空，重新生成
          if (!base64 || base64.length === 0) {
            console.warn('参考图像base64数据为空，重新生成')
            base64 = await this.blobToBase64(blob)
            // 更新缓存
            this.base64ImagesCache.set(blobId, base64)
          }

          base64Images.push(base64)
        }

        base64Images.forEach(image => {
          // 确保图像数据不为空
          if (image && image.length > 0) {
            contents.push({
              inlineData: {
                mimeType: 'image/png',
                data: image,
              },
            })
          } else {
            console.warn('跳过空的参考图像数据')
          }
        })
      }

      // 检查contents是否包含有效的图像数据或文本提示
      const hasImageData = contents.some(item => item.inlineData && item.inlineData.data && item.inlineData.data.length > 0)
      const hasTextPrompt = contents.some(item => item.text && item.text.length > 0)

      // 如果既没有图像数据也没有文本提示，抛出错误
      if (!hasImageData && !hasTextPrompt) {
        throw new Error('没有有效的图像数据或文本提示用于生成')
      }

      // 准备请求配置，包括abortSignal
      const generateContentParams: {
        model: string
        contents: Array<{ text: string } | { inlineData: { mimeType: string; data: string } }>
        config?: { httpOptions: { abortSignal: AbortSignal } }
      } = {
        model: 'gemini-2.5-flash-image-preview',
        contents,
      }

      // 如果提供了abortSignal，则添加到请求配置中
      if (request.abortSignal) {
        generateContentParams.config = {
          httpOptions: {
            abortSignal: request.abortSignal,
          },
        }
      }

      const response = await genAI.models.generateContent(generateContentParams)

      // 检查是否有被禁止的内容
      if (response.candidates && response.candidates.length > 0) {
        const candidate = response.candidates[0]
        if (candidate.finishReason === 'PROHIBITED_CONTENT') {
          throw new Error('内容被禁止：您的请求包含不允许的内容。请尝试其他提示。')
        }
        if (candidate.finishReason === 'IMAGE_SAFETY' || candidate.finishReason === 'IMAGE_OTHER') {
          throw new Error('图像安全检查失败：请求的图像内容可能不安全。请尝试调整提示词。')
        }
        // 检查finishReason为STOP但没有inlineData的情况
        if (candidate.finishReason === 'STOP') {
          // 检查是否有inlineData
          let hasInlineData = false
          if (candidate.content && candidate.content.parts) {
            for (const part of candidate.content.parts) {
              if (part.inlineData) {
                hasInlineData = true
                break
              }
            }
          }

          // 如果没有inlineData，则抛出错误
          if (!hasInlineData && candidate.content && candidate.content.parts && candidate.content.parts.length > 0) {
            throw new Error(candidate.content.parts[0].text || '生成失败：未返回图像数据')
          }
        }
      }

      const images: Blob[] = []

      // 检查响应是否存在以及是否有内容
      if (response.candidates && response.candidates.length > 0 && response.candidates[0].content && response.candidates[0].content.parts) {
        for (const part of response.candidates[0].content.parts) {
          if (part.inlineData) {
            // 将返回的base64数据转换为Blob
            const byteString = atob(part.inlineData.data)
            const mimeString = part.inlineData.mimeType || 'image/png'
            const ab = new ArrayBuffer(byteString.length)
            const ia = new Uint8Array(ab)
            for (let i = 0; i < byteString.length; i++) {
              ia[i] = byteString.charCodeAt(i)
            }
            const blob = new Blob([ab], { type: mimeString })
            images.push(blob)
          }
        }

        // 如果没有图像数据但有文本响应，抛出包含文本的错误
        if (images.length === 0) {
          let textResponse = ''
          for (const part of response.candidates[0].content.parts) {
            if (part.text) {
              textResponse += part.text
            }
          }
          if (textResponse) {
            throw new Error(`生成失败：${textResponse}`)
          }
        }
      }

      // 获取usageMetadata（如果存在）
      const usageMetadata = response.usageMetadata

      return { images, usageMetadata }
    } catch (error) {
      console.error('生成图像时出错:', error)
      // 检查是否是由于abortSignal导致的取消
      if (error instanceof Error && error.name === 'AbortError') {
        throw new Error('生成已取消')
      }
      if (error instanceof Error && error.message) {
        throw error
      }
      throw new Error(`生成图像失败: ${error instanceof Error ? error.message : '未知错误'}`)
    }
  }

  async editImage(request: EditRequest): Promise<{ images: Blob[]; usageMetadata?: UsageMetadata }> {
    try {
      // 将原始图像Blob转换为base64以发送到API
      let originalImageBase64: string

      // 生成原始图像Blob的唯一标识符
      const originalBlobId = await this.generateBlobId(request.originalImage)

      // 检查缓存中是否已有该图像的base64数据
      if (this.base64ImagesCache.has(originalBlobId)) {
        // 从缓存中获取base64数据
        originalImageBase64 = this.base64ImagesCache.get(originalBlobId)!
        console.log('从缓存中获取原始图像base64数据')
      } else {
        // 转换Blob为base64并缓存结果
        originalImageBase64 = await this.blobToBase64(request.originalImage)
        // 将base64数据存储到缓存中
        this.base64ImagesCache.set(originalBlobId, originalImageBase64)
        console.log('生成并缓存原始图像base64数据')
      }

      // 如果base64数据为空，重新生成
      if (!originalImageBase64 || originalImageBase64.length === 0) {
        console.warn('原始图像base64数据为空，重新生成')
        originalImageBase64 = await this.blobToBase64(request.originalImage)
        // 更新缓存
        this.base64ImagesCache.set(originalBlobId, originalImageBase64)
      }

      const contents = [
        { text: this.buildEditPrompt(request) },
        {
          inlineData: {
            mimeType: 'image/png',
            data: originalImageBase64,
          },
        },
      ]

      // 如果提供了参考图像则添加
      if (request.referenceImages && request.referenceImages.length > 0) {
        // 将Blob转换为base64以发送到API
        const base64ReferenceImages: string[] = []

        // 为每个参考图像生成或获取base64数据
        for (const blob of request.referenceImages) {
          // 生成Blob的唯一标识符
          const blobId = await this.generateBlobId(blob)

          let base64: string
          // 检查缓存中是否已有该图像的base64数据
          if (this.base64ImagesCache.has(blobId)) {
            // 从缓存中获取base64数据
            base64 = this.base64ImagesCache.get(blobId)!
            console.log('从缓存中获取参考图像base64数据')
          } else {
            // 转换Blob为base64并缓存结果
            base64 = await this.blobToBase64(blob)
            // 将base64数据存储到缓存中
            this.base64ImagesCache.set(blobId, base64)
            console.log('生成并缓存参考图像base64数据')
          }

          // 如果base64数据为空，重新生成
          if (!base64 || base64.length === 0) {
            console.warn('参考图像base64数据为空，重新生成')
            base64 = await this.blobToBase64(blob)
            // 更新缓存
            this.base64ImagesCache.set(blobId, base64)
          }

          base64ReferenceImages.push(base64)
        }

        base64ReferenceImages.forEach(image => {
          // 确保图像数据不为空
          if (image && image.length > 0) {
            contents.push({
              inlineData: {
                mimeType: 'image/png',
                data: image,
              },
            })
          } else {
            console.warn('跳过空的参考图像数据')
          }
        })
      }

      if (request.maskImage) {
        // 将遮罩图像Blob转换为base64以发送到API
        let maskImageBase64: string

        // 生成遮罩图像Blob的唯一标识符
        const maskBlobId = await this.generateBlobId(request.maskImage)

        // 检查缓存中是否已有该图像的base64数据
        if (this.base64ImagesCache.has(maskBlobId)) {
          // 从缓存中获取base64数据
          maskImageBase64 = this.base64ImagesCache.get(maskBlobId)!
          console.log('从缓存中获取遮罩图像base64数据')
        } else {
          // 转换Blob为base64并缓存结果
          maskImageBase64 = await this.blobToBase64(request.maskImage)
          // 将base64数据存储到缓存中
          this.base64ImagesCache.set(maskBlobId, maskImageBase64)
          console.log('生成并缓存遮罩图像base64数据')
        }

        // 如果base64数据为空，重新生成
        if (!maskImageBase64 || maskImageBase64.length === 0) {
          console.warn('遮罩图像base64数据为空，重新生成')
          maskImageBase64 = await this.blobToBase64(request.maskImage)
          // 更新缓存
          this.base64ImagesCache.set(maskBlobId, maskImageBase64)
        }

        // 确保遮罩图像数据不为空
        if (maskImageBase64 && maskImageBase64.length > 0) {
          contents.push({
            inlineData: {
              mimeType: 'image/png',
              data: maskImageBase64,
            },
          })
        } else {
          console.warn('跳过空的遮罩图像数据')
        }
      }

      // 检查contents是否包含有效的图像数据或文本提示
      const hasImageData = contents.some(item => item.inlineData && item.inlineData.data && item.inlineData.data.length > 0)
      const hasTextPrompt = contents.some(item => item.text && item.text.length > 0)

      // 如果既没有图像数据也没有文本提示，抛出错误
      if (!hasImageData && !hasTextPrompt) {
        throw new Error('没有有效的图像数据或文本提示用于编辑')
      }

      // 准备请求配置，包括abortSignal
      const generateContentParams: {
        model: string
        contents: Array<{ text: string } | { inlineData: { mimeType: string; data: string } }>
        config?: { httpOptions: { abortSignal: AbortSignal } }
      } = {
        model: 'gemini-2.5-flash-image-preview',
        contents,
      }

      // 如果提供了abortSignal，则添加到请求配置中
      if (request.abortSignal) {
        generateContentParams.config = {
          httpOptions: {
            abortSignal: request.abortSignal,
          },
        }
      }

      const response = await genAI.models.generateContent(generateContentParams)

      // 检查是否有被禁止的内容
      if (response.candidates && response.candidates.length > 0) {
        const candidate = response.candidates[0]
        if (candidate.finishReason === 'PROHIBITED_CONTENT') {
          throw new Error('内容被禁止：您的请求包含不允许的内容。请尝试其他提示。')
        }
        if (candidate.finishReason === 'IMAGE_SAFETY' || candidate.finishReason === 'IMAGE_OTHER') {
          throw new Error('图像安全检查失败：请求的图像内容可能不安全。请尝试调整提示词。')
        }
        // 检查finishReason为STOP但没有inlineData的情况
        if (candidate.finishReason === 'STOP') {
          // 检查是否有inlineData
          let hasInlineData = false
          if (candidate.content && candidate.content.parts) {
            for (const part of candidate.content.parts) {
              if (part.inlineData) {
                hasInlineData = true
                break
              }
            }
          }

          // 如果没有inlineData，则抛出错误
          if (!hasInlineData && candidate.content && candidate.content.parts && candidate.content.parts.length > 0) {
            throw new Error(candidate.content.parts[0].text || '编辑失败：未返回图像数据')
          }
        }
      }

      const images: Blob[] = []

      // 检查响应是否存在以及是否有内容
      if (response.candidates && response.candidates.length > 0 && response.candidates[0].content && response.candidates[0].content.parts) {
        for (const part of response.candidates[0].content.parts) {
          if (part.inlineData) {
            // 将返回的base64数据转换为Blob
            const byteString = atob(part.inlineData.data)
            const mimeString = part.inlineData.mimeType || 'image/png'
            const ab = new ArrayBuffer(byteString.length)
            const ia = new Uint8Array(ab)
            for (let i = 0; i < byteString.length; i++) {
              ia[i] = byteString.charCodeAt(i)
            }
            const blob = new Blob([ab], { type: mimeString })
            images.push(blob)
          }
        }

        // 如果没有图像数据但有文本响应，抛出包含文本的错误
        if (images.length === 0) {
          let textResponse = ''
          for (const part of response.candidates[0].content.parts) {
            if (part.text) {
              textResponse += part.text
            }
          }
          if (textResponse) {
            throw new Error(`编辑失败：${textResponse}`)
          }
        }
      }

      // 获取usageMetadata（如果存在）
      const usageMetadata = response.usageMetadata

      return { images, usageMetadata }
    } catch (error) {
      console.error('编辑图像时出错:', error)
      // 检查是否是由于abortSignal导致的取消
      if (error instanceof Error && error.name === 'AbortError') {
        throw new Error('编辑已取消')
      }
      if (error instanceof Error && error.message) {
        throw error
      }
      throw new Error(`编辑图像失败: ${error instanceof Error ? error.message : '未知错误'}`)
    }
  }

  async segmentImage(request: SegmentationRequest): Promise<{ masks: Array<{ label: string; box_2d: [number, number, number, number]; mask: string }> }> {
    try {
      // 将图像Blob转换为base64以发送到API
      let imageBase64: string

      // 生成图像Blob的唯一标识符
      const blobId = await this.generateBlobId(request.image)

      // 检查缓存中是否已有该图像的base64数据
      if (this.base64ImagesCache.has(blobId)) {
        // 从缓存中获取base64数据
        imageBase64 = this.base64ImagesCache.get(blobId)!
        console.log('从缓存中获取分割图像base64数据')
      } else {
        // 转换Blob为base64并缓存结果
        imageBase64 = await this.blobToBase64(request.image)
        // 将base64数据存储到缓存中
        this.base64ImagesCache.set(blobId, imageBase64)
        console.log('生成并缓存分割图像base64数据')
      }

      // 如果base64数据为空，重新生成
      if (!imageBase64 || imageBase64.length === 0) {
        console.warn('分割图像base64数据为空，重新生成')
        imageBase64 = await this.blobToBase64(request.image)
        // 更新缓存
        this.base64ImagesCache.set(blobId, imageBase64)
      }

      const prompt = [
        {
          text: `分析此图像并为以下对象创建分割遮罩: ${request.query}

返回具有此确切结构的JSON对象:
{
  "masks": [
    {
      "label": "分割对象的描述",
      "box_2d": [x, y, width, height],
      "mask": "base64编码的二进制遮罩图像"
    }
  ]
}

仅分割请求的特定对象或区域。遮罩应该是二进制PNG，其中白色像素(255)表示选定区域，黑色像素(0)表示背景。`,
        },
      ]

      // 确保图像数据不为空
      if (imageBase64 && imageBase64.length > 0) {
        prompt.push({
          inlineData: {
            mimeType: 'image/png',
            data: imageBase64,
          },
        })
      } else {
        console.warn('跳过空的分割图像数据')
      }

      // 检查prompt是否包含有效的图像数据或文本提示
      const hasImageData = prompt.some(item => item.inlineData && item.inlineData.data && item.inlineData.data.length > 0)
      const hasTextPrompt = prompt.some(item => item.text && item.text.length > 0)

      // 如果既没有图像数据也没有文本提示，抛出错误
      if (!hasImageData && !hasTextPrompt) {
        throw new Error('没有有效的图像数据或文本提示用于分割')
      }

      // 准备请求配置，包括abortSignal
      const generateContentParams: {
        model: string
        contents: Array<{ text: string } | { inlineData: { mimeType: string; data: string } }>
        config?: { httpOptions: { abortSignal: AbortSignal } }
      } = {
        model: 'gemini-2.5-flash-image-preview',
        contents: prompt,
      }

      // 如果提供了abortSignal，则添加到请求配置中
      if (request.abortSignal) {
        generateContentParams.config = {
          httpOptions: {
            abortSignal: request.abortSignal,
          },
        }
      }

      const response = await genAI.models.generateContent(generateContentParams)

      // 检查是否有被禁止的内容
      if (response.candidates && response.candidates.length > 0) {
        const candidate = response.candidates[0]
        if (candidate.finishReason === 'PROHIBITED_CONTENT') {
          throw new Error('内容被禁止：您的请求包含不允许的内容。请尝试其他提示。')
        }
        if (candidate.finishReason === 'IMAGE_SAFETY' || candidate.finishReason === 'IMAGE_OTHER') {
          throw new Error('图像安全检查失败：请求的图像内容可能不安全。请尝试调整提示词。')
        }
        // 检查finishReason为STOP但没有inlineData的情况
        if (candidate.finishReason === 'STOP') {
          // 检查是否有inlineData
          let hasInlineData = false
          if (candidate.content && candidate.content.parts) {
            for (const part of candidate.content.parts) {
              if (part.inlineData) {
                hasInlineData = true
                break
              }
            }
          }

          // 如果没有inlineData，则抛出错误
          if (!hasInlineData && candidate.content && candidate.content.parts && candidate.content.parts.length > 0) {
            throw new Error(candidate.content.parts[0].text || '分割失败：未返回结果数据')
          }
        }
      }

      const responseText = response.candidates[0].content.parts[0].text
      return JSON.parse(responseText)
    } catch (error) {
      console.error('分割图像时出错:', error)
      // 检查是否是由于abortSignal导致的取消
      if (error instanceof Error && error.name === 'AbortError') {
        throw new Error('分割已取消')
      }
      if (error instanceof Error && error.message) {
        throw error
      }
      throw new Error(`分割图像失败: ${error instanceof Error ? error.message : '未知错误'}`)
    }
  }

  private buildEditPrompt(request: EditRequest): string {
    const maskInstruction = request.maskImage ? '\n\n重要: 仅在遮罩图像显示白色像素(值255)的地方应用更改。完全不更改所有其他区域。精确遵守遮罩边界并在边缘保持无缝混合。' : ''

    return `根据以下指令编辑此图像: ${request.instruction}\n\n保持原始图像的光照、透视和整体构图。使更改看起来自然且无缝集成。${maskInstruction}\n\n保持图像质量并确保编辑看起来专业且逼真。`
  }

  // 公共方法：清除base64图像缓存
  public clearBase64Cache(): void {
    this.base64ImagesCache.clear()
    console.log('已清除base64图像缓存')
  }

  // 公共方法：获取缓存大小
  public getCacheSize(): number {
    return this.base64ImagesCache.size
  }
}

export const geminiService = new GeminiService()