from __future__ import annotations

import base64
import json
from aiohttp import ClientSession

from ..typing import AsyncResult, Messages, ImageType
from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
from ..image import to_bytes, is_accepted_format
from ..errors import MissingAuthError

class GeminiPro(AsyncGeneratorProvider, ProviderModelMixin):
    url = "https://ai.google.dev"
    working = True
    supports_message_history = True
    default_model = "gemini-pro"
    models = ["gemini-pro", "gemini-pro-vision"]

    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: Messages,
        stream: bool = False,
        proxy: str = None,
        api_key: str = None,
        image: ImageType = None,
        **kwargs
    ) -> AsyncResult:
        model = "gemini-pro-vision" if not model and image else model
        model = cls.get_model(model)
        if not api_key:
            raise MissingAuthError('Missing "api_key" for auth')
        headers = {
            "Content-Type": "application/json",
        }
        async with ClientSession(headers=headers) as session:
            method = "streamGenerateContent" if stream else "generateContent"
            url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:{method}"
            contents = [
                {
                    "role": "model" if message["role"] == "assistant" else message["role"],
                    "parts": [{"text": message["content"]}]
                }
                for message in messages
            ]
            if image:
                image = to_bytes(image)
                contents[-1]["parts"].append({
                    "inline_data": {
                        "mime_type": is_accepted_format(image),
                        "data": base64.b64encode(image).decode()
                    }
                })
            data = {
                "contents": contents,
                "generationConfig": {
                    "stopSequences": kwargs.get("stop"),
                    "temperature": kwargs.get("temperature"),
                    "maxOutputTokens": kwargs.get("max_tokens"),
                    "topP": kwargs.get("top_p"),
                    "topK": kwargs.get("top_k"),
                }
            }
            async with session.post(url, params={"key": api_key}, json=data, proxy=proxy) as response:
                if not response.ok:
                    data = await response.json()
                    raise RuntimeError(data[0]["error"]["message"])
                if stream:
                    lines = []
                    async for chunk in response.content:
                        if chunk == b"[{\n":
                            lines = [b"{\n"]
                        elif chunk == b",\r\n" or chunk == b"]":
                            try:
                                data = b"".join(lines)
                                data = json.loads(data)
                                yield data["candidates"][0]["content"]["parts"][0]["text"]
                            except:
                                data = data.decode() if isinstance(data, bytes) else data
                                raise RuntimeError(f"Read text failed. data: {data}")
                            lines = []
                        else:
                            lines.append(chunk)
                else:
                    data = await response.json()
                    yield data["candidates"][0]["content"]["parts"][0]["text"]