import os
import time

from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO

from main_rest.app.utils.date_utils import D
from const import TEMP_FOLDER


def gemini_image(prompt, folder=None):
    from custom_logger.main_log import main_logger as LOG
    
    image_path = ''
    client = genai.Client(api_key="AIzaSyCSw4pcPDYdAnjzBB7J9ZKXtRJJvunjWtA") # a2tec key

    for i in range(3):
        response = client.models.generate_content(
            model="gemini-2.5-flash-image",
            contents=prompt,
            config=types.GenerateContentConfig(
            response_modalities=['TEXT', 'IMAGE']
            )
        )
        if folder == None:
            folder = TEMP_FOLDER
        
        if not os.path.exists(folder):
            os.makedirs(folder)
        
        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                image = Image.open(BytesIO((part.inline_data.data)))
                image_path = os.path.join(folder,f"gemini_{D.date_file_name()}_query.png")
                image.save(image_path)
        
        LOG.info(f"image generate : {image_path}")
        time.sleep(2)
        
        if os.path.exists(image_path):
            return image_path
    
    return image_path