import { GoogleGenAI } from "@google/genai";
import { ModelType, BgType } from "../types";

export const removeBackground = async (
  base64Image: string,
  modelType: ModelType,
  bgType: BgType,
  bgValue?: string
): Promise<string> => {
  const ai = new GoogleGenAI({ apiKey: process.env.API_KEY });
  
  const cleanBase64 = base64Image.split(',')[1] || base64Image;
  const mimeType = base64Image.split(';')[0].split(':')[1] || 'image/png';

  let backgroundInstruction = "";
  if (bgType === 'transparent') {
    backgroundInstruction = "Isolate the main subject and place it on a transparent background (alpha channel). Ensure edges are smooth and free of artifacts.";
  } else if (bgType === 'color') {
    backgroundInstruction = `Isolate the main subject and place it on a solid flat background of color ${bgValue || '#FFFFFF'}. Ensure pixel-perfect edges.`;
  } else if (bgType === 'scenic') {
    backgroundInstruction = `Isolate the main subject and realistically composite it into a new background described as: ${bgValue || 'a professional studio background'}. Match lighting and shadows perfectly.`;
  }

  const prompt = `Task: High-precision background removal and subject isolation.
  Instruction: ${backgroundInstruction}
  Target Accuracy: >98.5%.
  Return only the updated image.`;

  try {
    const response = await ai.models.generateContent({
      model: modelType,
      contents: {
        parts: [
          {
            inlineData: {
              data: cleanBase64,
              mimeType: mimeType,
            },
          },
          {
            text: prompt,
          },
        ],
      },
      config: {
        imageConfig: {
          aspectRatio: "1:1",
        }
      }
    });

    if (!response.candidates?.[0]?.content?.parts) {
      throw new Error("Invalid response from model");
    }

    for (const part of response.candidates[0].content.parts) {
      if (part.inlineData) {
        return `data:image/png;base64,${part.inlineData.data}`;
      }
    }

    throw new Error("No image data returned from model. Check if the prompt was followed.");
  } catch (error) {
    console.error("Gemini API Error:", error);
    throw error;
  }
};