from __future__ import annotations

from aiohttp import ClientSession
import json

from ...typing import AsyncResult, Messages
from ..base_provider import AsyncGeneratorProvider, ProviderModelMixin
from ...image import ImageResponse


class NexraSDLora(AsyncGeneratorProvider, ProviderModelMixin):
    label = "Nexra Stable Diffusion Lora"
    url = "https://nexra.aryahcr.cc/documentation/stable-diffusion/en"
    api_endpoint = "https://nexra.aryahcr.cc/api/image/complements"
    working = True
    
    default_model = 'sdxl-lora'
    models = [default_model]

    @classmethod
    def get_model(cls, model: str) -> str:
        return cls.default_model

    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: Messages,
        proxy: str = None,
        response: str = "url", # base64 or url
        guidance: str = 0.3, # Min: 0, Max: 5
        steps: str = 2, # Min: 2, Max: 10
        **kwargs
    ) -> AsyncResult:
        model = cls.get_model(model)
        
        headers = {
            "Content-Type": "application/json"
        }
        async with ClientSession(headers=headers) as session:
            prompt = messages[0]['content']
            data = {
                "prompt": prompt,
                "model": model,
                "response": response,
                "data": {
                    "guidance": guidance,
                    "steps": steps
                }
            }
            async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
                text_data = await response.text()
                
                if response.status == 200:
                    try:
                        json_start = text_data.find('{')
                        json_data = text_data[json_start:]
                        
                        data = json.loads(json_data)
                        if 'images' in data and len(data['images']) > 0:
                            image_url = data['images'][-1]
                            yield ImageResponse(image_url, prompt)
                        else:
                            yield ImageResponse("No images found in the response.", prompt)
                    except json.JSONDecodeError:
                        yield ImageResponse("Failed to parse JSON. Response might not be in JSON format.", prompt)
                else:
                    yield ImageResponse(f"Request failed with status: {response.status}", prompt)