import logging
import json
import httpx
from typing import List, Dict, Any, Optional
from src.services.provider.base import BaseVideoService, ServiceResponse, TaskStatus, GenerationResult
from src.utils.image_processing import resolve_image_param

logger = logging.getLogger(__name__)

class VolcengineVideoService(BaseVideoService):
    BASE_URL = "https://ark.cn-beijing.volces.com/api/v3/contents/generations/tasks"

    def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key=api_key, **kwargs)
        
        if not self.api_key:
             logger.warning("API Key is not set. Volcengine video generation will fail.")

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

        try:
            return await self._submit_task_async(prompt, api_key, **kwargs)
        except Exception as e:
            logger.error(f"Volcengine Video submission failed: {e}")
            return ServiceResponse(
                status=TaskStatus.FAILED,
                error=str(e)
            )

    async def _submit_task_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"
        }
        
        # --- Unified Service Logic ---
        actual_model = self.model_name
        variants = getattr(self, "config", {}).get("variants")
        
        # Handle image inputs
        image_inputs = []
        if "image_inputs" in kwargs and isinstance(kwargs["image_inputs"], list):
            image_inputs.extend([resolve_image_param(img) for img in kwargs["image_inputs"]])
            
        if variants:
            if len(image_inputs) > 1 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)
        
        content = [
            {
                "type": "text",
                "text": prompt
            }
        ]
        
        # Add image content objects
        # 日志ic to determine 'role' based on API docs:
        # 1 image -> first_frame (default) or reference_image (if lite i2v)
        # 2 images -> first_frame + last_frame
        
        if image_inputs:
            # 简单 heuristic mapping
            if len(image_inputs) == 1:
                content.append({
                    "type": "image_url",
                    "image_url": {"url": image_inputs[0]},
                    "role": "first_frame" # 默认 to first frame
                })
            elif len(image_inputs) == 2:
                content.append({
                    "type": "image_url",
                    "image_url": {"url": image_inputs[0]},
                    "role": "first_frame"
                })
                content.append({
                    "type": "image_url",
                    "image_url": {"url": image_inputs[1]},
                    "role": "last_frame"
                })
        
        cmd_flags = []
        
        if kwargs.get("aspect_ratio"):
             cmd_flags.append(f"--rt {kwargs['aspect_ratio']}")
            
        if kwargs.get("resolution"):
            cmd_flags.append(f"--rs {kwargs['resolution']}")
            
        if "duration" in kwargs:
            cmd_flags.append(f"--dur {kwargs['duration']}")
            
        if "fps" in kwargs:
            cmd_flags.append(f"--fps {kwargs['fps']}")
            
        # Force watermark to false
        cmd_flags.append("--wm false")
            
        if cmd_flags:
            # 查找 the text content and append flags
            for item in content:
                if item["type"] == "text":
                    item["text"] = f"{item['text']} {' '.join(cmd_flags)}"
                    break
        
        payload = {
            "model": actual_model,
            "content": content
        }
        
        # Top-level parameters
        if "generate_audio" in kwargs:
            payload["generate_audio"] = kwargs["generate_audio"]

        try:
            logger.info(f"Submitting Volcengine Video task for model {actual_model}...")
            async with httpx.AsyncClient(timeout=60) as client:
                response = await client.post(
                    self.BASE_URL,
                    headers=headers,
                    data=json.dumps(payload)
                )

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

            response.raise_for_status()
            data = response.json()
            
            task_id = data.get("id")
            if not task_id:
                raise ValueError("No task_id returned from API")
                
            return ServiceResponse(
                status=TaskStatus.QUEUED,
                task_id=task_id
            )
            
        except Exception as e:
            raise e

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

        try:
            return await self._check_status_async(task_id, api_key)
        except Exception as e:
            logger.error(f"Volcengine Video status check failed: {e}")
            return ServiceResponse(
                status=TaskStatus.UNKNOWN,
                error=str(e)
            )
            
    async def _check_status_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"
        }
        
        url = f"{self.BASE_URL}/{task_id}"
        
        try:
            async with httpx.AsyncClient(timeout=60) as client:
                response = await client.get(
                    url,
                    headers=headers
                )
            response.raise_for_status()
            data = response.json()
            
            status_str = data.get("status")
            
            # 映射 API status to internal TaskStatus
            status_map = {
                "queued": TaskStatus.QUEUED,
                "running": TaskStatus.PROCESSING,
                "succeeded": TaskStatus.SUCCEEDED,
                "failed": TaskStatus.FAILED,
                "cancelled": TaskStatus.FAILED,
                "expired": TaskStatus.FAILED
            }
            
            status = status_map.get(status_str, TaskStatus.UNKNOWN)
            
            results = []
            if status == TaskStatus.SUCCEEDED:
                content = data.get("content", {})
                video_url = content.get("video_url")
                if video_url:
                    results.append(GenerationResult(url=video_url))
            
            error_msg = None
            if status == TaskStatus.FAILED:
                error_obj = data.get("error")
                if error_obj:
                    error_msg = f"{error_obj.get('code')}: {error_obj.get('message')}"
            
            return ServiceResponse(
                status=status,
                task_id=task_id,
                results=results,
                error=error_msg
            )
            
        except Exception as e:
            raise e
