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

import { GoogleGenAI } from '@google/genai'

// 注意：在生产环境中，这应该通过后端代理处理
const API_KEY = import.meta.env.VITE_GEMINI_API_KEY || 'demo-key'
const genAI = new GoogleGenAI({ apiKey: API_KEY })

export interface GenerationRequest {
  prompt: string
  referenceImages?: string[] // base64数组
  temperature?: number
  seed?: number
}

export interface EditRequest {
  instruction: string
  originalImage: string // base64
  referenceImages?: string[] // base64数组
  maskImage?: string // base64
  temperature?: number
  seed?: number
}

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

export interface SegmentationRequest {
  image: string // base64
  query: string // "像素(x,y)处的对象" 或 "红色汽车"
}

export class GeminiService {
  async generateImage(request: GenerationRequest): Promise<{ images: string[]; usageMetadata?: any }> {
    try {
      const contents: any[] = [{ text: request.prompt }]

      // 如果提供了参考图像则添加
      if (request.referenceImages && request.referenceImages.length > 0) {
        request.referenceImages.forEach(image => {
          contents.push({
            inlineData: {
              mimeType: 'image/png',
              data: image,
            },
          })
        })
      }

      const response = await genAI.models.generateContent({
        model: 'gemini-2.5-flash-image-preview',
        contents,
      })

      // 检查是否有被禁止的内容
      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') {
          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: string[] = []

      // 检查响应是否存在以及是否有内容
      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) {
            images.push(part.inlineData.data)
          }
        }
      }

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

      return { images, usageMetadata }
    } catch (error) {
      console.error('生成图像时出错:', error)
      if (error instanceof Error && error.message) {
        throw error
      }
      throw new Error(`生成图像失败: ${error instanceof Error ? error.message : '未知错误'}`)
    }
  }

  async editImage(request: EditRequest): Promise<{ images: string[]; usageMetadata?: any }> {
    try {
      const contents = [
        { text: this.buildEditPrompt(request) },
        {
          inlineData: {
            mimeType: 'image/png',
            data: request.originalImage,
          },
        },
      ]

      // 如果提供了参考图像则添加
      if (request.referenceImages && request.referenceImages.length > 0) {
        request.referenceImages.forEach(image => {
          contents.push({
            inlineData: {
              mimeType: 'image/png',
              data: image,
            },
          })
        })
      }

      if (request.maskImage) {
        contents.push({
          inlineData: {
            mimeType: 'image/png',
            data: request.maskImage,
          },
        })
      }

      const response = await genAI.models.generateContent({
        model: 'gemini-2.5-flash-image-preview',
        contents,
      })

      // 检查是否有被禁止的内容
      if (response.candidates && response.candidates.length > 0) {
        const candidate = response.candidates[0]
        if (candidate.finishReason === 'PROHIBITED_CONTENT') {
          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: string[] = []

      // 检查响应是否存在以及是否有内容
      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) {
            images.push(part.inlineData.data)
          }
        }
      }

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

      return { images, usageMetadata }
    } catch (error) {
      console.error('编辑图像时出错:', error)
      if (error instanceof Error && error.message) {
        throw error
      }
      throw new Error(`编辑图像失败: ${error instanceof Error ? error.message : '未知错误'}`)
    }
  }

  async segmentImage(request: SegmentationRequest): Promise<any> {
    try {
      const prompt = [
        {
          text: `分析此图像并为以下对象创建分割遮罩: ${request.query}

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

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

      const response = await genAI.models.generateContent({
        model: 'gemini-2.5-flash-image-preview',
        contents: prompt,
      })

      // 检查是否有被禁止的内容
      if (response.candidates && response.candidates.length > 0) {
        const candidate = response.candidates[0]
        if (candidate.finishReason === 'PROHIBITED_CONTENT') {
          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)
      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}

保持原始图像的光照、透视和整体构图。使更改看起来自然且无缝集成。${maskInstruction}

保持图像质量并确保编辑看起来专业且逼真。`
  }
}

export const geminiService = new GeminiService()