from __future__ import annotations

import json
import uuid

from ...typing import AsyncResult, Messages
from ...providers.response import ImageResponse, ImagePreview, JsonConversation, Reasoning
from ...requests import StreamSession
from ...errors import ResponseError
from ..base_provider import AsyncGeneratorProvider, ProviderModelMixin
from ..helper import format_image_prompt
from .Janus_Pro_7B import get_zerogpu_token
from .raise_for_status import raise_for_status

class BlackForestLabsFlux1Dev(AsyncGeneratorProvider, ProviderModelMixin):
    url = "https://black-forest-labs-flux-1-dev.hf.space"
    space = "black-forest-labs/FLUX.1-dev"
    referer = f"{url}/?__theme=light"

    working = True

    default_model = 'black-forest-labs-flux-1-dev'
    default_image_model = default_model
    model_aliases = {"flux-dev": default_model, "flux": default_model}
    image_models = [default_image_model, *model_aliases.keys()]
    models = image_models

    @classmethod
    def run(cls, method: str, session: StreamSession, conversation: JsonConversation, data: list = None):
        headers = {
            "accept": "application/json",
            "content-type": "application/json",
            "x-zerogpu-token": conversation.zerogpu_token,
            "x-zerogpu-uuid": conversation.zerogpu_uuid,
            "referer": cls.referer,
        }
        if method == "post":
            return session.post(f"{cls.url}/gradio_api/queue/join?__theme=light", **{
                "headers": {k: v for k, v in headers.items() if v is not None},
                "json": {"data": data,"event_data":None,"fn_index":2,"trigger_id":4,"session_hash":conversation.session_hash}

            })
        return session.get(f"{cls.url}/gradio_api/queue/data?session_hash={conversation.session_hash}", **{
            "headers": {
                "accept": "text/event-stream",
                "content-type": "application/json",
                "referer": cls.referer,
            }
        })

    @classmethod
    async def create_async_generator(
        cls, 
        model: str, 
        messages: Messages,
        prompt: str = None,
        proxy: str = None,
        width: int = 1024,
        height: int = 1024,
        guidance_scale: float = 3.5,
        num_inference_steps: int = 28,
        seed: int = 0,
        randomize_seed: bool = True,
        cookies: dict = None,
        api_key: str = None,
        zerogpu_uuid: str = "[object Object]",
        **kwargs
    ) -> AsyncResult:
        model = cls.get_model(model)
        async with StreamSession(impersonate="chrome", proxy=proxy) as session:
            prompt = format_image_prompt(messages, prompt)
            data = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps]
            conversation = JsonConversation(zerogpu_token=api_key, zerogpu_uuid=zerogpu_uuid, session_hash=uuid.uuid4().hex)
            if conversation.zerogpu_token is None:
                conversation.zerogpu_uuid, conversation.zerogpu_token = await get_zerogpu_token(cls.space, session, conversation, cookies)
            async with cls.run(f"post", session, conversation, data) as response:
                await raise_for_status(response)
                assert (await response.json()).get("event_id")
                async with cls.run("get", session, conversation) as event_response:
                    await raise_for_status(event_response)
                    async for chunk in event_response.iter_lines():
                        if chunk.startswith(b"data: "):
                            try:
                                json_data = json.loads(chunk[6:])
                                if json_data is None:
                                    continue
                                if json_data.get('msg') == 'log':
                                    yield Reasoning(status=json_data["log"])

                                if json_data.get('msg') == 'progress':
                                    if 'progress_data' in json_data:
                                        if json_data['progress_data']:
                                            progress = json_data['progress_data'][0]
                                            yield Reasoning(status=f"{progress['desc']} {progress['index']}/{progress['length']}")
                                        else:
                                            yield Reasoning(status=f"Generating")

                                elif json_data.get('msg') == 'process_generating':
                                    for item in json_data['output']['data'][0]:
                                        if isinstance(item, dict) and "url" in item:
                                            yield ImagePreview(item["url"], prompt)
                                        elif isinstance(item, list) and len(item) > 2 and "url" in item[1]:
                                            yield ImagePreview(item[2], prompt)

                                elif json_data.get('msg') == 'process_completed':
                                    if 'output' in json_data and 'error' in json_data['output']:
                                        json_data['output']['error'] = json_data['output']['error'].split(" <a ")[0]
                                        raise ResponseError(json_data['output']['error'])
                                    if 'output' in json_data and 'data' in json_data['output']:
                                        yield Reasoning(status="Finished")
                                        if len(json_data['output']['data']) > 0:
                                            yield ImageResponse(json_data['output']['data'][0]["url"], prompt)
                                    break
                            except (json.JSONDecodeError, KeyError, TypeError) as e:
                                raise RuntimeError(f"Failed to parse message: {chunk.decode(errors='replace')}", e)