"""
快手可灵AI客户端
支持AI视频生成、AI图片生成、AI对口型等功能
"""
import asyncio
import httpx
import json
import time
from typing import Optional, Dict, Any, List
from enum import Enum


class BizType(str, Enum):
    """业务类型"""
    AI_IMAGE = "aiImage"
    AI_VIDEO = "aiVideo"
    AI_LIP_SYNC = "aiLipSync"


class TaskStatus(str, Enum):
    """任务状态"""
    PROCESS = "process"
    FINISHED = "finished"
    FAILED = "failed"


class KlingClient:
    """快手可灵AI客户端"""
    
    def __init__(self, cookie: Optional[str] = None, timeout: int = 30):
        """
        初始化客户端
        
        Args:
            cookie: 认证Cookie
            timeout: 请求超时时间（秒）
        """
        self.cookie = cookie
        self.timeout = timeout
        self.base_url = "https://kling.kuaishou.com/api"
        
    def _get_headers(self) -> Dict[str, str]:
        """获取请求头"""
        headers = {
            "Content-Type": "application/json",
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
        }
        if self.cookie:
            headers["Cookie"] = self.cookie
        return headers
    
    async def create_image_task(
        self,
        prompt: str,
        negative_prompt: Optional[str] = None,
        cfg: str = "50",
        aspect_ratio: str = "16:9",
        image_count: int = 4,
        version: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        创建AI图片生成任务
        
        Args:
            prompt: 生成内容的提示词
            negative_prompt: 不希望呈现的内容
            cfg: 创意想象力与创意相关性比例
            aspect_ratio: 长宽比 (9:16, 16:9, 1:1)
            image_count: 生成图片数量(1-4)
            version: 模型版本
            
        Returns:
            包含task_id的响应
        """
        url = f"{self.base_url}/image/generate"
        payload = {
            "prompt": prompt,
            "cfg": cfg,
            "aspect_ratio": aspect_ratio,
            "image_count": image_count
        }
        
        if negative_prompt:
            payload["negative_prompt"] = negative_prompt
        if version:
            payload["version"] = version
            
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            response = await client.post(
                url,
                json=payload,
                headers=self._get_headers()
            )
            response.raise_for_status()
            return response.json()
    
    async def create_video_task(
        self,
        prompt: str,
        mode: str = "text2video",
        image_url: Optional[str] = None,
        aspect_ratio: str = "16:9",
        add_audio: bool = False,
        start_frame_image: Optional[str] = None,
        end_frame_image: Optional[str] = None,
        version: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        创建AI视频生成任务
        
        Args:
            prompt: 生成内容的提示词
            mode: 生成模式 (text2video, image2video)
            image_url: 参考图片地址
            aspect_ratio: 长宽比
            add_audio: 是否自动添加音频
            start_frame_image: 首帧图片URL
            end_frame_image: 尾帧图片URL
            version: 模型版本
            
        Returns:
            包含task_id的响应
        """
        url = f"{self.base_url}/video/generate"
        payload = {
            "prompt": prompt,
            "mode": mode,
            "aspect_ratio": aspect_ratio,
            "add_audio": add_audio
        }
        
        if image_url:
            payload["image_url"] = image_url
        if start_frame_image:
            payload["start_frame_image"] = start_frame_image
        if end_frame_image:
            payload["end_frame_image"] = end_frame_image
        if version:
            payload["version"] = version
            
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            response = await client.post(
                url,
                json=payload,
                headers=self._get_headers()
            )
            response.raise_for_status()
            return response.json()
    
    async def create_lipsync_task(
        self,
        video_id: Optional[str] = None,
        video_url: Optional[str] = None,
        mode: str = "text2video",
        text: Optional[str] = None,
        voice_id: Optional[str] = None,
        voice_language: str = "zh",
        voice_speed: float = 1.0,
        audio_type: Optional[str] = None,
        audio_file: Optional[str] = None,
        audio_url: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        创建AI对口型任务
        
        Args:
            video_id: 视频ID
            video_url: 视频URL
            mode: 生成模式 (text2video, audio2video)
            text: 对口型文本内容
            voice_id: 音色ID
            voice_language: 音色语种
            voice_speed: 语速
            audio_type: 音频类型 (file, url)
            audio_file: 音频文件路径
            audio_url: 音频URL
            
        Returns:
            包含task_id的响应
        """
        url = f"{self.base_url}/lipsync/generate"
        payload = {
            "mode": mode,
            "voice_language": voice_language,
            "voice_speed": voice_speed
        }
        
        if video_id:
            payload["video_id"] = video_id
        if video_url:
            payload["video_url"] = video_url
        if text:
            payload["text"] = text
        if voice_id:
            payload["voice_id"] = voice_id
        if audio_type:
            payload["audio_type"] = audio_type
        if audio_file:
            payload["audio_file"] = audio_file
        if audio_url:
            payload["audio_url"] = audio_url
            
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            response = await client.post(
                url,
                json=payload,
                headers=self._get_headers()
            )
            response.raise_for_status()
            return response.json()
    
    async def query_task_status(
        self,
        task_id: str,
        biz_type: BizType
    ) -> Dict[str, Any]:
        """
        查询任务状态
        
        Args:
            task_id: 任务ID
            biz_type: 业务类型
            
        Returns:
            任务状态信息
        """
        url = f"{self.base_url}/task/status"
        params = {
            "task_id": task_id,
            "biz_type": biz_type.value
        }
        
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            response = await client.get(
                url,
                params=params,
                headers=self._get_headers()
            )
            response.raise_for_status()
            return response.json()
    
    async def wait_for_task(
        self,
        task_id: str,
        biz_type: BizType,
        max_wait_time: int = 300,
        poll_interval: int = 5
    ) -> Dict[str, Any]:
        """
        等待任务完成
        
        Args:
            task_id: 任务ID
            biz_type: 业务类型
            max_wait_time: 最大等待时间（秒）
            poll_interval: 轮询间隔（秒）
            
        Returns:
            最终任务状态
        """
        start_time = time.time()
        
        while time.time() - start_time < max_wait_time:
            result = await self.query_task_status(task_id, biz_type)
            status = result.get("status")
            
            if status == TaskStatus.FINISHED:
                return result
            elif status == TaskStatus.FAILED:
                return result
            
            await asyncio.sleep(poll_interval)
        
        raise TimeoutError(f"Task {task_id} did not complete within {max_wait_time} seconds")
