from __future__ import annotations

import json
from aiohttp import ClientSession

from ...typing import AsyncResult, Messages
from ..base_provider import AsyncGeneratorProvider, ProviderModelMixin
from ..helper import format_prompt


class NexraGeminiPro(AsyncGeneratorProvider, ProviderModelMixin):
    label = "Nexra Gemini PRO"
    api_endpoint = "https://nexra.aryahcr.cc/api/chat/complements"
    models = ['gemini-pro']

    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: Messages,
        proxy: str = None,
        **kwargs
    ) -> AsyncResult:
        headers = {
            "Content-Type": "application/json"
        }
        async with ClientSession(headers=headers) as session:
            data = {
                "messages": [
                    {'role': 'assistant', 'content': ''},
                    {'role': 'user', 'content': format_prompt(messages)}
                ],
                "markdown": False,
                "stream": True,
                "model": model
            }
            async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
                response.raise_for_status()
                full_response = ''
                async for line in response.content:
                    if line:
                        messages = line.decode('utf-8').split('\x1e')
                        for message_str in messages:
                            try:
                                message = json.loads(message_str)
                                if message.get('message'):
                                    full_response = message['message']
                                if message.get('finish'):
                                    yield full_response.strip()
                                    return
                            except json.JSONDecodeError:
                                pass