import httpx
import os
from http import HTTPStatus
from typing import Optional, Dict, Any, List
from src.config.settings import UPLOAD_DIR
from src.services.provider.base import BaseVideoService, ServiceResponse, TaskStatus, GenerationResult, map_provider_status
from src.utils.image_processing import resolve_image_param
from src.services.storage_service import storage_manager
import logging

logger = logging.getLogger(__name__)

class WanVideoService(BaseVideoService):
    def __init__(self, model_name: str = "wan2.2-t2v-plus", api_key: Optional[str] = None, **kwargs):
        """ 初始化 Video Generation Service (Wanx).
        
        Args:
            model_name: The model identifier (e.g., 'wan2.2-t2v-plus')
            api_key: The API key for DashScope
            **kwargs: Additional arguments
        """
        super().__init__(model_name, api_key=api_key, **kwargs)
            
        if not api_key:
            logger.warning("API Key is not set. Video generation will fail.")
        
        # 验证 model name using base class method
        config_path = os.path.join(os.path.dirname(__file__), "models.json")
        self.validate_model_from_config(config_path, "video")

    async def generate(self,
                       prompt: str,
                       **kwargs) -> ServiceResponse:
        """ 生成 video via HTTP API (Asynchronous).
        Supports:
        - Text-to-Video (T2V): Provide prompt
        - Image-to-Video (I2V): Provide image_inputs[0]
        - Keyframe-to-Video (KF2V): Provide image_inputs[0] (first), image_inputs[1] (last)
        - Reference-to-Video (R2V): Provide video_inputs (videos) + image_inputs (images)
        """
        # Get effective API key (user key > system key)
        user_id = kwargs.get('user_id')
        api_key = self.get_effective_api_key(user_id)

        if not api_key:
            logger.error("No API key available for DashScope video generation")
            return ServiceResponse(
                status=TaskStatus.FAILED,
                error="API key not configured"
            )

        # Determine endpoint and input format based on model or params
        
        # 1. Normalize inputs
        raw_image_inputs = kwargs.get("image_inputs") or []
        if isinstance(raw_image_inputs, str):
            image_inputs = [raw_image_inputs]
        else:
            image_inputs = raw_image_inputs

        raw_video_inputs = kwargs.get("video_inputs") or []
        if isinstance(raw_video_inputs, str):
            video_inputs = [raw_video_inputs]
        else:
            video_inputs = raw_video_inputs

        raw_audio_inputs = kwargs.get("audio_inputs") or []
        if isinstance(raw_audio_inputs, str):
            audio_inputs = [raw_audio_inputs]
        else:
            audio_inputs = raw_audio_inputs
            
        enable_audio = kwargs.get("enable_audio")

        # Resolve all image URLs (local -> base64, OSS -> signed)
        image_inputs = [resolve_image_param(img) for img in image_inputs]
        
        # Collect reference video URLs (sign OSS URLs)
        video_inputs = [storage_manager.sign_url(v) for v in video_inputs if v]

        # --- Unified Service Logic ---
        # If the service is initialized with a base name (e.g., "wan2.6"),
        # determine the actual sub-model based on input parameters.
        actual_model = self.model_name
        
        variants = getattr(self, "config", {}).get("variants")
        
        if variants:
            has_last_frame = len(image_inputs) > 1
            
            if video_inputs:
                actual_model = variants.get("r2v", actual_model)
            elif has_last_frame and "kf2v" in variants:
                actual_model = variants.get("kf2v", actual_model)
            elif image_inputs:
                actual_model = variants.get("i2v", actual_model)
            else:
                actual_model = variants.get("t2v", actual_model)
        
        is_r2v = "r2v" in actual_model
        is_kf2v = "kf2v" in actual_model
        is_wan22 = "wan2.2" in actual_model or "wan2.5" in actual_model
        
        # 默认 endpoint (T2V and general I2V like wan2.6)
        url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation/video-synthesis"
        
        # Specific Keyframe-to-Video models use a different endpoint
        if is_kf2v:
             url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2video/video-synthesis"

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

        # Construct Input
        input_data = {"prompt": prompt}
        
        # Optional input fields
        if kwargs.get("negative_prompt"):
            input_data["negative_prompt"] = kwargs["negative_prompt"]
            
        # Audio handling
        audio_url = kwargs.get("audio_url")
        if audio_inputs and len(audio_inputs) > 0:
            audio_url = audio_inputs[0]

        if audio_url:
            input_data["audio_url"] = storage_manager.sign_url(audio_url)

        # --- R2V: Use unified reference_urls (images + videos mixed) ---
        if is_r2v:
            # wan2.6-r2v / wan2.6-r2v-flash use reference_urls (new API)
            # reference_urls supports mixed images + videos, up to 5 total (images 0~5, videos 0~3)
            reference_urls = []
            # Add reference videos first (they define character order: character1, character2, ...)
            reference_urls.extend(video_inputs)
            # Add reference images
            reference_urls.extend(image_inputs)
            
            if reference_urls:
                input_data["reference_urls"] = reference_urls
        
        # --- KF2V: first_frame_url + last_frame_url ---
        elif is_kf2v:
            if len(image_inputs) > 0:
                input_data["first_frame_url"] = image_inputs[0]
            if len(image_inputs) > 1:
                input_data["last_frame_url"] = image_inputs[1]
            
        # --- Standard T2V/I2V ---
        else:
            if video_inputs:
                input_data["reference_video_urls"] = video_inputs
            
            if len(image_inputs) > 0:
                input_data["img_url"] = image_inputs[0]

        # Construct Parameters
        parameters = {}
        
        # Size / Resolution handling
        # Support new dimension field
        dimension = kwargs.get("dimension")
        if dimension:
            parameters["size"] = dimension
        elif "size" in kwargs and kwargs["size"]:
             parameters["size"] = kwargs["size"]
        
        if "resolution" in kwargs and kwargs["resolution"]:
             parameters["resolution"] = kwargs["resolution"]

        # Models that require explicit size (W*H) format
        if (is_wan22 or is_r2v) and "size" not in parameters:
             res = parameters.get("resolution")
             if res == "720P":
                 parameters["size"] = "1280*720"
             elif res == "1080P":
                 parameters["size"] = "1920*1080"
             else:
                 parameters["size"] = "1920*1080"  # 默认 16:9
             
        # Known parameters to pass through
        known_params = [
            "duration", "shot_type", 
            "template"
        ]
        
        for key in known_params:
            if key in kwargs and kwargs[key] is not None:
                parameters[key] = kwargs[key]
        
        # Audio parameter: R2V supports audio generation (default true)
        if enable_audio is not None:
            parameters["audio"] = enable_audio
        
        # Watermark: default false, not exposed to frontend
        parameters["watermark"] = False
        
        # Prompt Extend: default false
        parameters["prompt_extend"] = False

        
        # Default duration if not provided
        if "duration" not in parameters:
            parameters["duration"] = 5

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

        try:
            logger.info(f"Sending Video Gen request to {url} with model {actual_model} (base: {self.model_name})")
            async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, connect=60.0)) as client:
                response = await client.post(url, headers=headers, json=payload)
                if response.status_code != 200:
                    logger.error(f"Video Gen API Error: {response.status_code} - {response.text}")
                    logger.error(f"Payload: {payload}")
                    response.raise_for_status()
                result_json = response.json()
            
            output = result_json.get("output", {})
            task_id = output.get("task_id")
            
            return ServiceResponse(
                status=TaskStatus.PENDING,
                task_id=task_id,
                meta={"request_id": result_json.get("request_id")}
            )

        except Exception as e:
            logger.error(f"Video Gen Call Failed: {str(e)}")
            raise RuntimeError(f"Video Gen Service Error: {str(e)}")

    async def check_status(self, task_id: str, user_id: Optional[str] = None) -> ServiceResponse:
        """
        Check the status of an asynchronous video generation task via HTTP API.
        """
        # Get effective API key (user key > system key)
        api_key = self.get_effective_api_key(user_id)

        if not api_key:
            logger.error("No API key available for DashScope video check status")
            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, connect=60.0)) as client:
                response = await client.get(url, headers=headers)
                if response.status_code != 200:
                    logger.error(f"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)
            
            error_msg = None
            if status == TaskStatus.FAILED:
                code = output.get("code", "UnknownCode")
                message = output.get("message", "Unknown Error")
                error_msg = f"{code}: {message}"
            
            results = []
            video_url = output.get("video_url")
            if video_url:
                results.append(GenerationResult(
                    url=video_url,
                    usage=usage
                ))
            
            return ServiceResponse(
                status=status,
                task_id=task_id,
                results=results,
                error=error_msg,
                meta={"request_id": result_json.get("request_id")}
            )

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