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 BaseVideoService, ServiceResponse, TaskStatus, GenerationResult

logger = logging.getLogger(__name__)

class ModelScopeVideoService(BaseVideoService):
    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, 'video')
        
        if not self.api_key:
             logger.warning("API Key is not set. API calls will fail.")

    async def _generate_async(self, prompt: str, api_key: str = None, **kwargs) -> 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-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 "steps" in kwargs:
            payload["steps"] = kwargs["steps"]
        if "guidance" in kwargs:
            payload["guidance"] = kwargs["guidance"]
        if "loras" in kwargs:
            payload["loras"] = kwargs["loras"]
        elif "lora_model" in kwargs:
            # 如果提供了单个 LoRA 参数，转换为 loras 数组格式
            lora_entry = {
                "model": kwargs["lora_model"]
            }
            if "lora_strength" in kwargs:
                lora_entry["strength"] = kwargs["lora_strength"]
            payload["loras"] = [lora_entry]

        # 注意: ModelScope video generation might use 'text' or 'prompt' key depending on model?
        # 标准 API usually uses 'prompt'.

        # We need to guess the endpoint for video.
        # Common convention: /videos/generations
        # Or maybe it uses the same /images/generations endpoint but model type differs?
        # Unlikely.
        # Let's try /video/generations (singular) or /videos/generations (plural).
        # Most likely /videos/generations based on /images/generations.

        endpoint = f"{self.BASE_URL}/videos/generations"

        try:
            # 1. Submit Task
            logger.info(f"Submitting ModelScope Video task for model {self.model_name}...")
            async with httpx.AsyncClient(timeout=30) as client:
                response = await client.post(
                    endpoint,
                    headers=headers,
                    content=json.dumps(payload, ensure_ascii=False).encode('utf-8'),
                )

                # Fallback for endpoint if 404?
                if response.status_code == 404:
                    logger.warning(f"Endpoint {endpoint} not found. Trying /video/generations...")
                    endpoint = f"{self.BASE_URL}/video/generations"
                    response = await client.post(
                        endpoint,
                        headers=headers,
                        content=json.dumps(payload, ensure_ascii=False).encode('utf-8'),
                    )

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

            logger.info(f"ModelScope Video task submitted. Task ID: {task_id}")

            return await self._poll_task_async(task_id, api_key=api_key)

        except Exception as e:
            logger.error(f"ModelScope video 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:
        # ModelScope video task type is 'video_generation'
        # 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": "video_generation"
        }

        start_time = asyncio.get_event_loop().time()
        timeout = 600 # 10 minutes for video

        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":
                    # Video result parsing
                    # ModelScope video response usually has 'output_video' or 'output_videos'
                    # We'll try to collect all available video URLs
                    output_video = data.get("output_video")
                    output_videos = data.get("output_videos", [])

                    results = []

                    if output_video:
                         output_videos.append(output_video)

                    for vid_url in output_videos:
                        # 返回 remote URL directly as requested
                        results.append(GenerationResult(
                            url=vid_url,
                            content="Video 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
                    )

                await asyncio.sleep(5)

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


    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)
