"""
Qwen Image Generation Service

Handles Qwen image generation via DashScope API.
"""
import logging
from typing import Dict, List, Optional

import httpx

from src.services.provider.base import (
    BaseImageService, ServiceResponse, TaskStatus, GenerationResult, map_provider_status
)
from src.utils.image_processing import resolve_image_param

logger = logging.getLogger(__name__)


class QwenImageService(BaseImageService):
    """Qwen Image Generation Service"""

    # Class-level batch task mapping
    _batch_tasks: Dict[str, List[str]] = {}

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

        Args:
            model_name: The model identifier (e.g., 'qwen-image-plus', 'qwen-image')
            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_qwen_multimodal(self, prompt: str, model_name: str = None, api_key: str = None, **kwargs) -> ServiceResponse:
        """Handle Qwen Multimodal generation (Edit/Fusion) via HTTP API (Synchronous)."""
        target_model = model_name or self.model_name
        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}"
        }

        # Prepare content (images + text)
        images = kwargs.get("image_inputs") or []
        if isinstance(images, str):
            images = [images]

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

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

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

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

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

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

        try:
            logger.info(f"Sending Qwen Multimodal request to {url} with model {target_model}")
            async with httpx.AsyncClient(timeout=(30, 60)) as client:
                response = await client.post(url, headers=headers, json=payload)
            if response.status_code != 200:
                error_msg = f"Qwen Multimodal API Error: {response.status_code}"
                try:
                    err_json = response.json()
                    if "message" in err_json:
                        error_msg += f" - {err_json['message']}"
                    elif "code" in err_json:
                        error_msg += f" - {err_json['code']}"
                except (ValueError, KeyError):
                    error_msg += f" - {response.text}"
                logger.error(error_msg)
                raise RuntimeError(error_msg)
            result_json = response.json()
            output = result_json.get("output", {})
            usage = result_json.get("usage", {})

            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 item in output["choices"]:
                    msg = item.get("message", {})
                    content_data = msg.get("content")

                    url = None
                    text_content = None

                    if isinstance(content_data, list):
                        for part in content_data:
                            if isinstance(part, dict):
                                if "image" in part:
                                    url = part["image"]
                                elif "text" in part:
                                    text_content = part["text"]
                    elif isinstance(content_data, str):
                        text_content = content_data

                    results.append(GenerationResult(
                        url=url,
                        content=text_content,
                        usage=usage
                    ))

            return ServiceResponse(
                status=TaskStatus.SUCCEEDED,
                results=results,
                meta={"request_id": result_json.get("request_id")}
            )
        except Exception as e:
            logger.error(f"Qwen Multimodal Request Failed: {str(e)}")
            raise RuntimeError(f"Qwen Multimodal Service Error: {str(e)}")

    async def _generate_qwen_text2image_single(self, prompt: str, actual_model: str, api_key: str = None, **kwargs) -> str:
        """Generate single image async task, return task_id."""
        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", "1328*1328").replace("x", "*"),
            "n": 1,
            "prompt_extend": kwargs.get("prompt_extend", True),
            "watermark": kwargs.get("watermark", 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": actual_model,
            "input": input_data,
            "parameters": parameters
        }

        async with httpx.AsyncClient() as client:
            response = await client.post(url, headers=headers, json=payload)
        if response.status_code != 200:
            logger.error(f"Qwen-Image API Error: {response.status_code} - {response.text}")
            response.raise_for_status()
        result_json = response.json()
        return result_json.get("output", {}).get("task_id")

    async def generate(self, prompt: str, **kwargs) -> ServiceResponse:
        """Handle Qwen generation via HTTP API.
        Dispatches to appropriate method based on model name.
        """
        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)

        # Multimodal (Image Editing/Fusion) models
        if "edit" in actual_model or "multimodal" in actual_model:
            return await self._generate_qwen_multimodal(prompt, model_name=actual_model, api_key=api_key, **kwargs)

        # Default to Text-to-Image (Asynchronous)
        n = min(kwargs.get("n", 1), 4)

        try:
            task_ids = []
            for i in range(n):
                logger.info(f"Qwen-Image starting task {i+1}/{n}")
                call_kwargs = {**kwargs}
                if "seed" in kwargs:
                    call_kwargs["seed"] = kwargs["seed"] + i
                task_id = await self._generate_qwen_text2image_single(prompt, actual_model, api_key=api_key, **call_kwargs)
                if task_id:
                    task_ids.append(task_id)

            primary_task_id = task_ids[0] if task_ids else None
            if primary_task_id and len(task_ids) > 1:
                QwenImageService._batch_tasks[primary_task_id] = task_ids

            return ServiceResponse(
                status=TaskStatus.PENDING,
                task_id=primary_task_id,
                meta={
                    "batch_task_ids": task_ids,
                    "batch_size": n,
                    "is_batch": len(task_ids) > 1
                }
            )
        except Exception as e:
            logger.error(f"Qwen-Image Request Failed: {str(e)}")
            raise RuntimeError(f"Qwen-Image Service Error: {str(e)}")

    async def _check_single_task(self, task_id: str, api_key: str = None) -> ServiceResponse:
        """Check single task status."""
        url = f"https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}"
        effective_key = api_key or self.api_key
        headers = {"Authorization": f"Bearer {effective_key}"}

        async with httpx.AsyncClient() as client:
            response = await client.get(url, headers=headers)
        if response.status_code != 200:
            logger.error(f"Qwen 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)

        results = []
        if "results" in output:
            for item in output["results"]:
                results.append(GenerationResult(
                    url=item.get("url"),
                    orig_prompt=item.get("orig_prompt"),
                    actual_prompt=item.get("actual_prompt"),
                    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")
            }
        )

    async def check_status(self, task_id: str, user_id: Optional[str] = None) -> ServiceResponse:
        """Check status of Qwen-Image task.
        Supports batch task queries.
        """
        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"
            )

        try:
            batch_task_ids = QwenImageService._batch_tasks.get(task_id)

            if not batch_task_ids:
                return await self._check_single_task(task_id, api_key=api_key)

            # Batch task query
            logger.info(f"Checking batch task {task_id} with {len(batch_task_ids)} sub-tasks")
            all_results = []
            all_succeeded = True
            any_failed = False
            errors = []

            for sub_task_id in batch_task_ids:
                sub_response = await self._check_single_task(sub_task_id, api_key=api_key)

                if sub_response.status == TaskStatus.SUCCEEDED:
                    if sub_response.results:
                        all_results.extend(sub_response.results)
                elif sub_response.status == TaskStatus.FAILED:
                    any_failed = True
                    if sub_response.error:
                        errors.append(sub_response.error)
                else:
                    all_succeeded = False

            if all_succeeded:
                del QwenImageService._batch_tasks[task_id]
                return ServiceResponse(
                    status=TaskStatus.SUCCEEDED,
                    task_id=task_id,
                    results=all_results,
                    meta={"batch_size": len(batch_task_ids), "completed": True}
                )
            elif any_failed and not all_results:
                return ServiceResponse(
                    status=TaskStatus.FAILED,
                    task_id=task_id,
                    error=f"All batch tasks failed: {'; '.join(errors[:3])}",
                    meta={"batch_size": len(batch_task_ids), "errors": errors}
                )
            elif any_failed:
                return ServiceResponse(
                    status=TaskStatus.SUCCEEDED,
                    task_id=task_id,
                    results=all_results,
                    meta={"batch_size": len(batch_task_ids), "completed": True, "partial": True}
                )
            else:
                return ServiceResponse(
                    status=TaskStatus.PROCESSING,
                    task_id=task_id,
                    results=all_results if all_results else None,
                    meta={"batch_size": len(batch_task_ids), "completed": False}
                )

        except Exception as e:
            logger.error(f"Qwen Task Check Failed: {str(e)}")
            raise RuntimeError(f"Qwen Task Check Error: {str(e)}")
