import json
import random
from typing import Any, Optional

import requests

from webscout.litagent import LitAgent
from webscout.Provider.TTI.base import BaseImages, TTICompatibleProvider
from webscout.Provider.TTI.utils import ImageData, ImageResponse


class Images(BaseImages):
    """Handles image generation requests for the Miragic AI provider."""

    def __init__(self, client: "MiragicAI"):
        self._client = client

    def create(
        self,
        *,
        model: str = "flux",
        prompt: str,
        n: int = 1,
        size: str = "1024x1024",
        response_format: str = "url",
        user: Optional[str] = None,
        style: str = "none",
        aspect_ratio: str = "1:1",
        timeout: int = 60,
        image_format: str = "png",
        seed: Optional[int] = None,
        **kwargs,
    ) -> ImageResponse:
        """
        Creates images using the Miragic AI API.

        Args:
            model (str): The model to use. Options: "flux", "turbo", "gptimage". Defaults to "flux".
            prompt (str): The text description of the image to generate.
            n (int): Number of images to generate. Defaults to 1.
            size (str): Image dimensions in "WxH" format. Defaults to "1024x1024".
            response_format (str): Format of the response ("url" or "b64_json"). Defaults to "url".
            user (Optional[str]): Optional user identifier.
            style (str): Optional style parameter.
            aspect_ratio (str): Optional aspect ratio parameter.
            timeout (int): Request timeout in seconds. Defaults to 60.
            image_format (str): Output image format ("png" or "jpeg"). Defaults to "png".
            seed (Optional[int]): Random seed for reproducibility.
            **kwargs: Additional parameters:
                - enhance_prompt (bool): Use AI to enhance the prompt.
                - safe_filter (bool): Enable strict NSFW filtering.
                - image_url (str): Source image URL for image-to-image models.

        Returns:
            ImageResponse: Object containing the generated image data.

        Raises:
            RuntimeError: If image generation or retrieval fails.
        """
        width, height = 1024, 1024
        if size:
            try:
                parts = size.split('x')
                if len(parts) == 2:
                    width, height = int(parts[0]), int(parts[1])
            except ValueError:
                pass

        width = max(256, min(2048, width))
        height = max(256, min(2048, height))

        enhance = kwargs.get("enhance_prompt", False)
        safe = kwargs.get("safe_filter", False)
        image_url = kwargs.get("image_url")

        images_data = []

        for _ in range(n):
            current_seed = seed if seed is not None else random.randint(0, 2**32 - 1)

            payload = {
                "data": [
                    prompt,
                    model,
                    width,
                    height,
                    current_seed,
                    image_url,
                    enhance,
                    safe
                ]
            }

            try:
                post_url = f"{self._client.api_endpoint}/call/generate_image_via_api_secure"
                resp = self._client.session.post(post_url, json=payload, timeout=timeout)
                resp.raise_for_status()

                event_id = resp.json().get("event_id")
                if not event_id:
                    raise RuntimeError(f"Failed to obtain event_id: {resp.text}")

                stream_url = f"{self._client.api_endpoint}/call/generate_image_via_api_secure/{event_id}"
                image_url_result = None

                with self._client.session.get(stream_url, stream=True, timeout=timeout + 60) as stream_resp:
                    stream_resp.raise_for_status()
                    for line in stream_resp.iter_lines():
                        if not line:
                            continue
                        line_text = line.decode('utf-8')

                        if line_text.startswith('data: '):
                            data_str = line_text[6:]
                            try:
                                data_json = json.loads(data_str)
                                if isinstance(data_json, list) and len(data_json) > 0:
                                    result_item = data_json[0]
                                    if isinstance(result_item, dict):
                                        image_url_result = result_item.get("url")
                            except json.JSONDecodeError:
                                continue
            except Exception as e:
                raise RuntimeError(f"Image generation failed: {e}")

            if image_url_result:
                images_data.append(ImageData(url=image_url_result))
            else:
                 raise RuntimeError("Failed to retrieve image URL from the response stream.")

        return ImageResponse(data=images_data)

class MiragicAI(TTICompatibleProvider):
    """
    Miragic AI TTI Provider implementation.

    Reverse engineered from the Hugging Face Space:
    https://huggingface.co/spaces/Miragic-AI/Miragic-AI-Image-Generator
    """

    required_auth: bool = False
    working: bool = True
    AVAILABLE_MODELS = ["flux", "turbo", "gptimage"]

    def __init__(self, **kwargs: Any):
        """Initializes the MiragicAI provider with a persistent session."""
        self.api_endpoint = "https://miragic-ai-miragic-ai-image-generator.hf.space/gradio_api"
        self.session = requests.Session()
        self.session.headers.update({
            "User-Agent": LitAgent().random(),
            "Referer": "https://miragic-ai-miragic-ai-image-generator.hf.space/",
            "Origin": "https://miragic-ai-miragic-ai-image-generator.hf.space",
            "Content-Type": "application/json"
        })
        self.images = Images(self)

    @property
    def models(self):
        """Returns a list of available models for this provider."""
        class _ModelList:
            def list(inner_self):
                return type(self).AVAILABLE_MODELS
        return _ModelList()

if __name__ == "__main__":
    from rich import print

    try:
        client = MiragicAI()
        print(f"Available Models: {client.models.list()}")

        print("Generating sample image...")
        response = client.images.create(
            prompt="A serene landscape with a lake and mountains, oil painting style",
            model="flux",
            size="1024x1024"
        )
        print(response)

    except Exception as error:
        print(f"Error during execution: {error}")
