"""
Z-Image Generation Service

Handles Z-Image Turbo generation via DashScope API.
"""
import logging
import time
from typing import Optional

import httpx

from src.services.provider.base import BaseImageService, ServiceResponse, TaskStatus, GenerationResult

logger = logging.getLogger(__name__)


class ZImageService(BaseImageService):
    """Z-Image Generation Service"""

    def __init__(self, model_name: str = "z-image-turbo", api_key: Optional[str] = None, **kwargs):
        """Initialize Z-Image Generation Service.

        Supports:
        - z-image-turbo (Fast, lightweight model)

        Args:
            model_name: The model identifier.
            api_key: The API key for DashScope
            **kwargs: Additional arguments
        """
        super().__init__(model_name, api_key=api_key, **kwargs)

        if not self.api_key:
            logger.warning("API Key is not set. Image generation will fail.")

        # Validate model name using base class method
        import os
        config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models.json")
        self.validate_model_from_config(config_path, "image")

    async def _generate_single(self, prompt: str, actual_model: str, api_key: str = None, **kwargs) -> GenerationResult:
        """Generate single image (Z-Image API only supports one at a time)."""
        url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
        effective_key = api_key or self.api_key

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {effective_key}"
        }

        # Handle size parameter
        size = kwargs.get("size", "1024*1536")
        if size and isinstance(size, str):
            size = size.replace("x", "*")

        width = kwargs.get("width")
        height = kwargs.get("height")
        if width and height:
            size = f"{width}*{height}"

        prompt_extend = kwargs.get("prompt_extend", False)

        payload = {
            "model": actual_model,
            "input": {
                "messages": [
                    {
                        "role": "user",
                        "content": [{"text": prompt}]
                    }
                ]
            },
            "parameters": {
                "prompt_extend": prompt_extend,
                "negative_prompt": kwargs.get("negative_prompt", ""),
                "size": size
            }
        }

        if "seed" in kwargs:
            payload["parameters"]["seed"] = kwargs["seed"]

        logger.info(f"Sending Z-Image Request: {payload}")
        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(url, headers=headers, json=payload)
            if response.status_code != 200:
                logger.error(f"Z-Image API Error: {response.status_code} - {response.text}")
                response.raise_for_status()
            result_json = response.json()
        output = result_json.get("output", {})
        usage = result_json.get("usage", {})

        if "code" in result_json:
            raise RuntimeError(f"Z-Image API Error: {result_json.get('code')} - {result_json.get('message')}")

        choices = output.get("choices", [])
        if not choices:
            raise RuntimeError("No choices returned from Z-Image API")

        choice = choices[0]
        message = choice.get("message", {})
        content_list = message.get("content", [])

        image_url = None
        rewritten_prompt = None

        for item in content_list:
            if "image" in item:
                image_url = item["image"]
            if "text" in item:
                rewritten_prompt = item["text"]

        if not image_url:
            raise RuntimeError("No image URL found in Z-Image response")

        return GenerationResult(
            url=image_url,
            orig_prompt=prompt,
            actual_prompt=rewritten_prompt if prompt_extend else prompt,
            usage=usage
        )

    async def generate(self, prompt: str, **kwargs) -> ServiceResponse:
        """Generate image using Z-Image Turbo (Sync HTTP).
        Since Z-Image API only supports one at a time, call multiple times for multiple images.
        """
        user_id = kwargs.get('user_id')
        api_key = self.get_effective_api_key(user_id)
        if not api_key:
            logger.error(f"No API key available for {self._provider_id}")
            return ServiceResponse(
                status=TaskStatus.FAILED,
                error="API key not configured"
            )

        kwargs = self.normalize_image_kwargs(**kwargs)
        actual_model = self.model_name
        variants = getattr(self, "config", {}).get("variants")

        if variants:
            images = kwargs.get("image_inputs")
            if images:
                actual_model = variants.get("i2i", variants.get("image", actual_model))
            else:
                actual_model = variants.get("t2i", actual_model)

        # Max 4 images for Z-Image
        n = min(kwargs.get("n", 1), 4)

        try:
            results = []
            for i in range(n):
                logger.info(f"Z-Image generating {i+1}/{n}")
                call_kwargs = {**kwargs}
                if "seed" in kwargs:
                    call_kwargs["seed"] = kwargs["seed"] + i
                result = await self._generate_single(prompt, actual_model, api_key=api_key, **call_kwargs)
                results.append(result)

            return ServiceResponse(
                status=TaskStatus.SUCCEEDED,
                task_id=f"zimage-batch-{time.time()}",
                results=results,
                meta={"batch_size": n}
            )

        except Exception as e:
            logger.error(f"Z-Image Generation Failed: {str(e)}")
            raise RuntimeError(f"Z-Image Service Error: {str(e)}")

    async def check_status(self, task_id: str) -> ServiceResponse:
        """Z-Image is synchronous, so check_status is not strictly needed.
        Implemented for interface consistency.
        """
        return ServiceResponse(
            status=TaskStatus.UNKNOWN,
            task_id=task_id,
            error="Z-Image service is synchronous; status check is not applicable."
        )
