import os
import mimetypes
from google import genai
from google.genai import types
from .config import load_prompt_template


class imggenrator:
    def __init__(self):
        self.api_key = os.getenv("IMG_GEN_API_KEY")
        self.model = os.getenv("model")
        self.client = genai.Client(api_key=self.api_key)

    # Notez que le paramètre `save_fn` n'est plus nécessaire.
    def generate(self, image_path: str, user_input_text: str):
        with open(image_path, "rb") as f:
            image_data = f.read()
        image_mime_type = mimetypes.guess_type(image_path)[0]
        full_prompt = load_prompt_template(image_path, user_input_text)
        contents = [
            types.Content(
                role="user",
                parts=[
                    types.Part.from_bytes(data=image_data, mime_type=image_mime_type),
                    types.Part.from_text(text=full_prompt),
                ],
            ),
        ]
        config = types.GenerateContentConfig(response_modalities=["IMAGE", "TEXT"])
        generated_images_data = []

        for chunk in self.client.models.generate_content_stream(
            model=self.model,
            contents=contents,
            config=config,
        ):
            if chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts:
                for part in chunk.candidates[0].content.parts:
                    if hasattr(part, "inline_data") and part.inline_data:
                        generated_images_data.append(part.inline_data.data)
                    elif hasattr(part, "text") and part.text:
                        print(part.text)
        return generated_images_data