import os
import io
import logging
from datetime import datetime
from typing import List, Union, Dict, Any
from PIL import Image

# Импортируем из нового SDK
from google import genai
from google.genai import types

# Для настройки логгера
logger = logging.getLogger(__name__)


class GoogleAdapter:
    def __init__(self, api_key: str):
        if not api_key:
            raise ValueError("API Key for Gemini is missing")
        self.client = genai.Client(api_key=api_key)
        # Укажите актуальную модель.
        # Если gemini-3-pro-image-preview недоступна, используйте gemini-2.0-flash-exp
        self.model_name = "gemini-3-pro-preview"

    def generate(
            self,
            prompt: str,
            image_bytes: bytes = None,
            generate_image: bool = False
    ) -> Dict[str, Any]:
        """
        Универсальный метод:
        - Если generate_image=True: просим модель вернуть картинку (Image Generation).
        - Если image_bytes переданы + generate_image=False: это Vision (описание фото).
        - Если image_bytes + generate_image=True: это Image-to-Image (редактирование).
        """
        if generate_image:
            self.model_name = "gemini-3-pro-image-preview"
        else :
            self.model_name = "gemini-3-pro-preview"
        contents = [prompt]

        # Если есть входное изображение (для Vision или для редактирования)
        if image_bytes:
            try:
                image = Image.open(io.BytesIO(image_bytes))
                contents.append(image)
            except Exception as e:
                logger.error(f"Error processing input image: {e}")
                return {"error": "Не удалось обработать входящее изображение."}

        # Настраиваем конфигурацию
        # Для генерации картинок добавляем 'IMAGE' в response_modalities
        modalities = ['TEXT']
        if generate_image:
            modalities.append('IMAGE')

        try:
            # Вызов API (синхронный метод в обертке, но aiogram вызывает его в треде,
            # либо используйте client.aio для асинхронности если поддерживается версией SDK)
            # В google-genai v0.3+ есть асинхронный клиент, но для простоты здесь стандартный вызов.
            # Чтобы не блокировать event loop, в main.py мы обернем это в to_thread при необходимости,
            # но пока используем стандартный вызов.

            response = self.client.models.generate_content(
                model=self.model_name,
                contents=contents,
                config=types.GenerateContentConfig(
                    response_modalities=modalities,
                    temperature=0.7 if not generate_image else 1.0,
                )
            )

            result = {"text": "", "images": []}

            # Парсим ответ (Text или Inline Data)
            if response.parts:
                for part in response.parts:
                    if part.text:
                        result["text"] += part.text

                    # Проверяем наличие сгенерированного изображения
                    if part.inline_data:
                        # ИСПРАВЛЕНИЕ: Берем "сырые" байты напрямую из ответа
                        # Это работает быстрее и не вызывает ошибку с PIL

                        # part.inline_data.data — это уже bytes
                        byte_arr = io.BytesIO(part.inline_data.data)
                        now = datetime.now()
                        # Имя файла для телеграма (формально)
                        byte_arr.name = f'{now.timestamp()}.png'

                        result["images"].append(byte_arr)

            return result

        except Exception as e:
            logger.error(f"Gemini API Error: {e}")
            return {"error": f"Ошибка API: {str(e)}"}