"""
Wan Image Generation Service

Handles Wan2.6 and Wan2.5 image generation via DashScope API.
"""
import logging
from typing import Optional

import httpx

from src.services.provider.base import BaseImageService, ServiceResponse, TaskStatus
from src.utils.image_processing import resolve_image_param

logger = logging.getLogger(__name__)


class WanImageService(BaseImageService):
    """Wan Image Generation Service"""

    def __init__(self, model_name: str = "wan2.6-t2i", api_key: Optional[str] = None, **kwargs):
        """Initialize Image Generation Service (Wanx).

        Args:
            model_name: The model identifier (e.g., 'wan2.6-t2i')
            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_wan2_6(self, prompt: str, model_name: str = None, api_key: str = None, **kwargs) -> ServiceResponse:
        """Handle Wan2.6 generation via HTTP API (Asynchronous)."""
        target_model = model_name or self.model_name
        url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation"
        effective_key = api_key or self.api_key

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

        # Extract normalized image inputs
        images = kwargs.get("image_inputs") or []
        if isinstance(images, str):
            images = [images]

        # Construct content
        content = [{"text": prompt}]
        for img in images:
            resolved_img = resolve_image_param(img)
            content.append({"image": resolved_img})

        messages = [{"role": "user", "content": content}]

        # Construct parameters
        parameters = {
            "watermark": False,
            "prompt_extend": False,
            "n": kwargs.get("n", 1)
        }

        if "size" in kwargs:
            parameters["size"] = kwargs["size"].replace("x", "*")

        if "negative_prompt" in kwargs and kwargs["negative_prompt"]:
            parameters["negative_prompt"] = kwargs["negative_prompt"]

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

        payload = {
            "model": target_model,
            "input": {"messages": messages},
            "parameters": parameters
        }

        try:
            import json
            logger.info(f"Sending Wan2.6 request to {url} with model {target_model}")
            logger.info(f"Wan2.6 request payload: {json.dumps(payload, ensure_ascii=False, default=str)}")
            async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, read=60.0)) as client:
                response = await client.post(url, headers=headers, json=payload)
            if response.status_code != 200:
                logger.error(f"Wan2.6 API Error: {response.status_code} - {response.text}")
                logger.error(f"Wan2.6 request payload that caused error: {json.dumps(payload, ensure_ascii=False, default=str)}")
                response.raise_for_status()
            result_json = response.json()
            output = result_json.get("output", {})
            return ServiceResponse(
                status=TaskStatus.PENDING,
                task_id=output.get("task_id"),
                meta={"request_id": result_json.get("request_id")}
            )
        except Exception as e:
            logger.error(f"Wan2.6 Request Failed: {str(e)}")
            raise RuntimeError(f"Wan2.6 Service Error: {str(e)}")

    async def _generate_wan_t2i(self, prompt: str, model_name: str = None, api_key: str = None, **kwargs) -> ServiceResponse:
        """Handle Wan2.5 and below generation via HTTP API (Asynchronous)."""
        target_model = model_name or self.model_name
        url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"
        effective_key = api_key or self.api_key

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

        parameters = {
            "size": kwargs.get("size", "1024*1024"),
            "n": kwargs.get("n", 1),
            "watermark": False,
            "prompt_extend": False
        }

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

        input_data = {"prompt": prompt}
        if "negative_prompt" in kwargs and kwargs["negative_prompt"]:
            input_data["negative_prompt"] = kwargs["negative_prompt"]

        payload = {
            "model": target_model,
            "input": input_data,
            "parameters": parameters
        }

        try:
            logger.info(f"Sending Wan T2I request to {url} with model {target_model}")
            async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, read=60.0)) as client:
                response = await client.post(url, headers=headers, json=payload)
            if response.status_code != 200:
                logger.error(f"Wan T2I API Error: {response.status_code} - {response.text}")
                response.raise_for_status()
            result_json = response.json()
            output = result_json.get("output", {})
            return ServiceResponse(
                status=TaskStatus.PENDING,
                task_id=output.get("task_id"),
                meta={"request_id": result_json.get("request_id")}
            )
        except Exception as e:
            logger.error(f"Wan T2I Image Gen Request Failed: {str(e)}")
            raise RuntimeError(f"Wan T2I Image Gen Service Error: {str(e)}")

    async def _generate_wan_i2i(self, prompt: str, model_name: str = None, api_key: str = None, **kwargs) -> ServiceResponse:
        """Handle Wan2.5 I2I generation via HTTP API (Asynchronous)."""
        target_model = model_name or self.model_name
        url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"
        effective_key = api_key or self.api_key

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

        parameters = {
            "size": kwargs.get("size", "1280*1280"),
            "n": kwargs.get("n", 1),
            "prompt_extend": kwargs.get("prompt_extend", True),
            "watermark": kwargs.get("watermark", False)
        }

        input_data = {"prompt": prompt}

        # Handle images
        images = kwargs.get("image_inputs")
        if not images:
            raise ValueError("image_inputs is required for I2I model")

        if isinstance(images, str):
            images = [images]

        images = [resolve_image_param(img) for img in images]
        input_data["images"] = images

        if "negative_prompt" in kwargs and kwargs["negative_prompt"]:
            input_data["negative_prompt"] = kwargs["negative_prompt"]

        payload = {
            "model": target_model,
            "input": input_data,
            "parameters": parameters
        }

        try:
            logger.info(f"Sending Wan I2I request to {url} with model {target_model}")
            async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, read=60.0)) as client:
                response = await client.post(url, headers=headers, json=payload)
            if response.status_code != 200:
                logger.error(f"Wan I2I API Error: {response.status_code} - {response.text}")
                response.raise_for_status()
            result_json = response.json()
            output = result_json.get("output", {})
            return ServiceResponse(
                status=TaskStatus.PENDING,
                task_id=output.get("task_id"),
                meta={"request_id": result_json.get("request_id")}
            )
        except Exception as e:
            logger.error(f"Wan I2I Image Gen Request Failed: {str(e)}")
            raise RuntimeError(f"Wan I2I Image Gen Service Error: {str(e)}")

    async def check_status(self, task_id: str, user_id: Optional[str] = None) -> ServiceResponse:
        """Check status of Wan task."""
        from src.services.provider.base import GenerationResult, map_provider_status

        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"
            )

        url = f"https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}"
        headers = {"Authorization": f"Bearer {api_key}"}

        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, read=60.0)) as client:
                response = await client.get(url, headers=headers)
            if response.status_code != 200:
                logger.error(f"Wan Task Check Error: {response.status_code} - {response.text}")
                response.raise_for_status()
            result_json = response.json()
            output = result_json.get("output", {})
            usage = result_json.get("usage", {})

            status_str = output.get("task_status", "UNKNOWN")
            status = map_provider_status("dashscope", status_str)

            if status == TaskStatus.SUCCEEDED:
                logger.info(f"Wan Task Succeeded. Output: {output}")

            results = []
            if "results" in output:
                for item in output["results"]:
                    results.append(GenerationResult(
                        url=item.get("url"),
                        content=None,
                        orig_prompt=item.get("orig_prompt"),
                        actual_prompt=item.get("actual_prompt"),
                        usage=usage
                    ))
            elif "choices" in output:
                for choice in output["choices"]:
                    message = choice.get("message", {})
                    content_list = message.get("content", [])
                    for item in content_list:
                        if isinstance(item, dict) and item.get("type") == "image":
                            results.append(GenerationResult(
                                url=item.get("image"),
                                content=None,
                                usage=usage
                            ))

            return ServiceResponse(
                status=status,
                task_id=output.get("task_id", task_id),
                results=results if results else None,
                error=output.get("message") if status == TaskStatus.FAILED else None,
                meta={
                    "submit_time": output.get("submit_time"),
                    "end_time": output.get("end_time"),
                    "request_id": result_json.get("request_id")
                }
            )
        except Exception as e:
            logger.error(f"Wan Task Check Failed: {str(e)}")
            raise RuntimeError(f"Wan Task Check Error: {str(e)}")

    async def generate(self, prompt: str, **kwargs) -> ServiceResponse:
        """Generate image using Wanx.

        Supports text-to-image and image-to-image generation.
        """
        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)

        # Wan2.6 via unified HTTP API
        if actual_model.startswith("wan2.6"):
            return await self._generate_wan2_6(prompt, model_name=actual_model, api_key=api_key, **kwargs)

        # Support branch for wan2.5 I2I
        if "i2i" in actual_model and actual_model.startswith("wan2.5"):
            return await self._generate_wan_i2i(prompt, model_name=actual_model, api_key=api_key, **kwargs)

        # Default to standard T2I asynchronous endpoint
        return await self._generate_wan_t2i(prompt, model_name=actual_model, api_key=api_key, **kwargs)
