import asyncio
import logging
import os
import json
import httpx
import uuid
import shutil
from typing import Dict, Any, Optional, List
from src.services.provider.base import BaseImageService, ServiceResponse, TaskStatus, GenerationResult
from src.utils.image_processing import resolve_image_param

logger = logging.getLogger(__name__)

class ModelScopeImageService(BaseImageService):
    BASE_URL = "https://api-inference.modelscope.cn/v1"

    def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key=api_key, **kwargs)

        # 验证 model
        config_path = os.path.join(os.path.dirname(__file__), 'models.json')
        self.validate_model_from_config(config_path, 'image')

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

    async def _generate_single_async(self, prompt: str, api_key: str = None, **kwargs) -> str:
        """生成单张图片，返回 task_id"""
        # Use provided api_key or fall back to instance api_key
        effective_key = api_key or self.api_key

        headers = {
            "Authorization": f"Bearer {effective_key}",
            "Content-Type": "application/json",
            "X-ModelScope-Async-Mode": "true"
        }

        payload = {
            "model": self.model_name,
            "prompt": prompt,
        }

        # Add optional parameters
        if "negative_prompt" in kwargs:
            payload["negative_prompt"] = kwargs["negative_prompt"]
        if "size" in kwargs:
            payload["size"] = kwargs["size"]
        if "seed" in kwargs:
            payload["seed"] = kwargs["seed"]
        if "steps" in kwargs:
            payload["steps"] = kwargs["steps"]
        if "guidance" in kwargs:
            payload["guidance"] = kwargs["guidance"]

        # Handle reference images
        ref_images = kwargs.get("image_inputs")
        if ref_images:
            if isinstance(ref_images, list) and len(ref_images) > 0:
                payload["image_url"] = resolve_image_param(ref_images[0])
            elif isinstance(ref_images, str):
                payload["image_url"] = resolve_image_param(ref_images)

        if "loras" in kwargs:
            payload["loras"] = kwargs["loras"]
        elif "lora_model" in kwargs:
            lora_entry = {"model": kwargs["lora_model"]}
            if "lora_strength" in kwargs:
                lora_entry["strength"] = kwargs["lora_strength"]
            payload["loras"] = [lora_entry]

        # 提交任务
        logger.info(f"Submitting ModelScope Image task for model {self.model_name}...")
        async with httpx.AsyncClient(timeout=30) as client:
            response = await client.post(
                f"{self.BASE_URL}/images/generations",
                headers=headers,
                content=json.dumps(payload, ensure_ascii=False).encode('utf-8'),
            )

        if response.status_code != 200:
            logger.error(f"ModelScope API Error: {response.text}")
        response.raise_for_status()

        task_data = response.json()
        task_id = task_data.get("task_id")

        if not task_id:
            raise ValueError(f"No task_id in response: {task_data}")

        return task_id

    async def _generate_async(self, prompt: str, api_key: str = None, **kwargs) -> ServiceResponse:
        """生成图片，支持多张（通过多次调用）"""
        # ModelScope 限制最多 4 张
        n = min(kwargs.get("n", 1), 4)

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

            # 由于是同步等待模式，需要轮询所有任务
            all_results = []
            for i, task_id in enumerate(task_ids):
                logger.info(f"Polling ModelScope task {i+1}/{len(task_ids)}: {task_id}")
                response = await self._poll_task_async(task_id, api_key=api_key)
                if response.status == TaskStatus.SUCCEEDED and response.results:
                    all_results.extend(response.results)
                else:
                    logger.warning(f"ModelScope task {task_id} failed or returned no results")

            return ServiceResponse(
                status=TaskStatus.SUCCEEDED if all_results else TaskStatus.FAILED,
                task_id=task_ids[0] if task_ids else None,
                results=all_results,
                meta={"batch_size": n, "completed": True}
            )

        except Exception as e:
            logger.error(f"ModelScope image generation failed: {e}")
            return ServiceResponse(
                status=TaskStatus.FAILED,
                error=str(e)
            )

    async def _poll_task_async(self, task_id: str, api_key: str = None) -> ServiceResponse:
        # Use provided api_key or fall back to instance api_key
        effective_key = api_key or self.api_key

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

        start_time = asyncio.get_event_loop().time()
        timeout = 300 # 5 minutes timeout

        while True:
            if asyncio.get_event_loop().time() - start_time > timeout:
                return ServiceResponse(status=TaskStatus.FAILED, error="Timeout waiting for task completion")

            try:
                async with httpx.AsyncClient(timeout=30) as client:
                    response = await client.get(
                        f"{self.BASE_URL}/tasks/{task_id}",
                        headers=headers,
                    )
                response.raise_for_status()
                data = response.json()

                status = data.get("task_status")

                if status == "SUCCEED":
                    output_images = data.get("output_images", [])
                    results = []

                    for img_url in output_images:
                        # 返回 remote URL directly as requested
                        results.append(GenerationResult(
                            url=img_url,
                            content="Image generated by ModelScope",
                            orig_prompt=data.get("output", {}).get("prompt", "")
                        ))

                    return ServiceResponse(
                        status=TaskStatus.SUCCEEDED,
                        results=results,
                        task_id=task_id
                    )

                elif status == "FAILED":
                    return ServiceResponse(
                        status=TaskStatus.FAILED,
                        error=data.get("message", "Unknown error"),
                        task_id=task_id
                    )

                # 运行中, PENDING
                await asyncio.sleep(2)

            except Exception as e:
                logger.warning(f"Error polling task {task_id}: {e}")
                await asyncio.sleep(2)


    async def generate(self, prompt: str, **kwargs) -> ServiceResponse:
        # Get effective API key (user key takes priority)
        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"
            )

        return await self._generate_async(prompt, api_key=api_key, **kwargs)

    async def check_status(self, task_id: str, user_id: Optional[str] = None) -> ServiceResponse:
        # Get effective API key (user key takes priority)
        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"
            )

        return await self._poll_task_async(task_id, api_key)
