import httpx
import logging
import asyncio
import json
import time
import jwt
import base64
from typing import Optional, Dict, Any, List, Union
from urllib.parse import urlparse

from src.services.provider.base import BaseVideoService, ServiceResponse, TaskStatus, GenerationResult
from src.config.settings import KLING_ACCESS_KEY, KLING_SECRET_KEY, KLING_API_BASE

logger = logging.getLogger(__name__)

class KlingVideoService(BaseVideoService):
    def __init__(self, model_name: str = "kling-v1", access_key: Optional[str] = None, secret_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key=access_key, **kwargs)
        self.access_key = self.api_key
        self.secret_key = secret_key

        if not self.access_key:
            self.access_key = KLING_ACCESS_KEY
            self.api_key = self.access_key
        if not self.secret_key:
            self.secret_key = KLING_SECRET_KEY

        if not self.access_key or not self.secret_key:
            logger.warning("KLING_ACCESS_KEY or KLING_SECRET_KEY is not set. Kling video generation will fail.")
        self.base_url = KLING_API_BASE.rstrip("/")

    async def _process_image_for_kling(self, image_input: str) -> str:
        """Process image for Kling API.
        Kling requires pure base64 string WITHOUT data URI scheme prefix.

        Input formats:
        - data:image/jpeg;base64,/9j/4AAQ... -> Extract base64 part
        - /9j/4AAQ... -> Use as is (already pure base64)
        - http://example.com/image.jpg -> Download and convert to base64

        Returns:
            Pure base64 string (no data URI scheme)
        """
        if not image_input:
            return image_input

        # 1. If it's a data URI, extract the base64 part
        if image_input.startswith("data:"):
            try:
                base64_data = image_input.split(",", 1)[1]
                logger.info(f"Extracted base64 from data URI (length: {len(base64_data)})")
                return base64_data
            except Exception as e:
                logger.error(f"Failed to extract base64 from data URI: {e}")
                return image_input

        # 2. If it's an HTTP(S) URL, download and convert to base64
        if image_input.startswith("http://") or image_input.startswith("https://"):
            try:
                logger.info(f"Downloading image from URL: {image_input[:100]}...")
                async with httpx.AsyncClient(timeout=30.0) as client:
                    response = await client.get(image_input)
                    response.raise_for_status()
                    image_bytes = response.content
                base64_data = base64.b64encode(image_bytes).decode('utf-8')
                logger.info(f"Converted URL image to base64 (length: {len(base64_data)})")
                return base64_data
            except Exception as e:
                logger.error(f"Failed to download and convert image from URL: {e}")
                return image_input

        # 3. Already pure base64 or unknown format, return as is
        logger.info(f"Using image input as is (length: {len(image_input)})")
        return image_input

    def _generate_jwt_token(self, access_key: str, secret_key: str) -> str:
        """ 生成 JWT token for Kling API authentication.
        """
        if not access_key or not secret_key:
            raise ValueError("KLING_ACCESS_KEY and KLING_SECRET_KEY are required")

        headers = {
            "alg": "HS256",
            "typ": "JWT"
        }
        payload = {
            "iss": access_key,
            "exp": int(time.time()) + 1800,  # Valid for 30 minutes
            "nbf": int(time.time()) - 5      # Valid from 5 seconds ago
        }

        # PyJWT encode returns a string in v2+
        token = jwt.encode(payload, secret_key, algorithm="HS256", headers=headers)
        return token

    def _get_headers(self, access_key: str, secret_key: str) -> Dict[str, str]:
        token = self._generate_jwt_token(access_key, secret_key)
        return {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}"
        }

    async def generate(self, prompt: str, **kwargs) -> ServiceResponse:
        """ 生成 video using Kling AI.
        Supports:
        - OmniVideo (kling-video-o1)
        - TextToVideo
        - ImageToVideo
        - MultiImageToVideo
        """
        user_id = kwargs.get('user_id')
        key_config = self.get_effective_api_key_with_config(user_id)
        if not key_config or not key_config.get('api_key'):
            logger.error(f"No API key available for {self._provider_id}")
            return ServiceResponse(
                status=TaskStatus.FAILED,
                error="API key not configured"
            )

        api_key = key_config['api_key']
        extra_config = key_config.get('extra_config', {})

        # Get secret_key from extra_config or fallback to system env
        secret_key = extra_config.get('secret_key') or self.secret_key

        # Pass keys through kwargs for internal methods
        kwargs['_access_key'] = api_key
        kwargs['_secret_key'] = secret_key

        # 1. OmniVideo Dispatch
        # Explicit model name or mode_type
        if self.model_name == "kling-video-o1" or kwargs.get("mode_type") == "omni":
            return await self._generate_omni(prompt, **kwargs)

        image_inputs = kwargs.get("image_inputs") or []
        if isinstance(image_inputs, str):
            image_inputs = [image_inputs]

        # 2. Multi-Image Dispatch
        # Check if image_list is provided and matches MultiImage format (list of dicts with 'image')
        image_list = kwargs.get("image_list")
        if image_list and isinstance(image_list, list) and len(image_list) > 0:
            first = image_list[0]
            # Omni uses "image_url", MultiImage uses "image"
            if "image" in first:
                return await self._generate_multi_image(prompt, **kwargs)
            elif "image_url" in first:
                # Likely intended for Omni, but let's see if model matches
                # user provided Omni format but didn't set model, maybe we should default to Omni?
                return await self._generate_omni(prompt, **kwargs)

        # 3. Image-to-Video Dispatch
        if image_inputs:
            return await self._generate_image2video(prompt, **kwargs)

        # 4. Default: Text-to-Video
        return await self._generate_text2video(prompt, **kwargs)

    def _get_kling_params(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
        """
        Resolve 'mode' and 'aspect_ratio' from kwargs.
        Mappings:
        - quality='1080p' -> mode='pro'
        - quality='720p' -> mode='std'
        - resolution='1920*1080' -> mode='pro', aspect_ratio='16:9'
        - resolution='1280*720' -> mode='std', aspect_ratio='16:9'
        """
        mode = kwargs.get("mode")
        aspect_ratio = kwargs.get("aspect_ratio")
        
        quality = kwargs.get("quality")
        resolution = kwargs.get("resolution")
        
        # 1. Resolve Mode
        if not mode:
            if quality == "Professional":
                mode = "pro"
            elif quality == "Standard":
                mode = "std"
            elif resolution:
                # 推断 from resolution
                # 专业: 1920*1080, 1080*1920
                if "1920" in resolution:
                    mode = "pro"
                # 标准: 1280*720, 720*1280
                elif "1280" in resolution or "720" in resolution:
                    mode = "std"
            
            if not mode:
                mode = "std" # 默认
        
        # 2. Resolve Aspect Ratio
        if not aspect_ratio:
            if resolution:
                try:
                    w, h = map(int, resolution.split("*"))
                    ratio = w / h
                    if 1.7 <= ratio <= 1.8:
                        aspect_ratio = "16:9"
                    elif 0.5 <= ratio <= 0.6:
                        aspect_ratio = "9:16"
                    elif 0.9 <= ratio <= 1.1:
                        aspect_ratio = "1:1"
                except (ValueError, TypeError, ZeroDivisionError):
                    pass
            
            if not aspect_ratio:
                aspect_ratio = "16:9" # 默认
                
        return {"mode": mode, "aspect_ratio": aspect_ratio}

    async def _generate_omni(self, prompt: str, **kwargs) -> ServiceResponse:
        url = f"{self.base_url}/videos/omni-video"
        params = self._get_kling_params(kwargs)
        
        # 进程 image_list for Kling format
        image_list = kwargs.get("image_list", [])
        
        # If no image_list was passed, build from normalized image_inputs
        if not image_list:
            image_inputs = kwargs.get("image_inputs") or []
            if isinstance(image_inputs, str):
                image_inputs = [image_inputs]
            if image_inputs:
                image_list = [{"image_url": image} for image in image_inputs]
        
        if image_list:
            processed_image_list = []
            for item in image_list:
                if isinstance(item, dict):
                    img_url = item.get("image_url") or item.get("image")
                    if img_url:
                        processed_item = item.copy()
                        processed_item["image_url"] = await self._process_image_for_kling(img_url)
                        # Remove "image" key if present to avoid confusion
                        processed_item.pop("image", None)
                        processed_image_list.append(processed_item)
                        logger.info(f"Processed Omni image_url (base64 length: {len(processed_item['image_url'])})")
                else:
                    processed_image_list.append(item)
            image_list = processed_image_list
        
        payload = {
            "model_name": self.model_name if self.model_name else "kling-video-o1",
            "prompt": prompt,
            "image_list": image_list,
            "element_list": kwargs.get("element_list", []),
            "video_list": kwargs.get("video_list", []),
            "mode": params["mode"],
            "aspect_ratio": params["aspect_ratio"],
            "duration": str(kwargs.get("duration", "5")),
            "callback_url": kwargs.get("callback_url"),
            "external_task_id": kwargs.get("external_task_id")
        }
        access_key = kwargs.get('_access_key')
        secret_key = kwargs.get('_secret_key')
        return await self._send_request(url, payload, "omni", access_key, secret_key)

    async def _generate_text2video(self, prompt: str, **kwargs) -> ServiceResponse:
        url = f"{self.base_url}/videos/text2video"
        params = self._get_kling_params(kwargs)
        payload = {
            "model_name": self.model_name,
            "prompt": prompt,
            "negative_prompt": kwargs.get("negative_prompt", ""),
            "cfg_scale": float(kwargs.get("cfg_scale", 0.5)),
            "mode": params["mode"],
            "aspect_ratio": params["aspect_ratio"],
            "duration": str(kwargs.get("duration", "5")),
            "camera_control": kwargs.get("camera_control"),
            "callback_url": kwargs.get("callback_url"),
            "external_task_id": kwargs.get("external_task_id")
        }
        access_key = kwargs.get('_access_key')
        secret_key = kwargs.get('_secret_key')
        return await self._send_request(url, payload, "text2video", access_key, secret_key)

    async def _generate_image2video(self, prompt: str, **kwargs) -> ServiceResponse:
        url = f"{self.base_url}/videos/image2video"
        image_inputs = kwargs.get("image_inputs") or []
        if isinstance(image_inputs, str):
            image_inputs = [image_inputs]
        image = image_inputs[0] if len(image_inputs) > 0 else None
        image_tail = image_inputs[1] if len(image_inputs) > 1 else None
        
        # 进程 images for Kling format (pure base64 without data URI scheme)
        if image:
            image = await self._process_image_for_kling(image)
            logger.info(f"Processed first frame image for Kling (base64 length: {len(image) if image else 0})")
        
        if image_tail:
            image_tail = await self._process_image_for_kling(image_tail)
            logger.info(f"Processed tail frame image for Kling (base64 length: {len(image_tail) if image_tail else 0})")
        
        params = self._get_kling_params(kwargs)
        payload = {
            "model_name": self.model_name,
            "image": image,
            "image_tail": image_tail,
            "prompt": prompt, # 可选 in I2V but good to have
            "negative_prompt": kwargs.get("negative_prompt", ""),
            "cfg_scale": float(kwargs.get("cfg_scale", 0.5)),
            "mode": params["mode"],
            "duration": str(kwargs.get("duration", "5")),
            "camera_control": kwargs.get("camera_control"),
            "callback_url": kwargs.get("callback_url"),
            "external_task_id": kwargs.get("external_task_id")
        }
        access_key = kwargs.get('_access_key')
        secret_key = kwargs.get('_secret_key')
        return await self._send_request(url, payload, "image2video", access_key, secret_key)

    async def _generate_multi_image(self, prompt: str, **kwargs) -> ServiceResponse:
        url = f"{self.base_url}/videos/multi-image2video"
        params = self._get_kling_params(kwargs)
        
        # 进程 image_list for Kling format
        image_list = kwargs.get("image_list", [])
        if image_list:
            processed_image_list = []
            for item in image_list:
                if isinstance(item, dict) and "image" in item:
                    # Multi-Image format: {"image": "...", ...}
                    processed_item = item.copy()
                    processed_item["image"] = await self._process_image_for_kling(item["image"])
                    processed_image_list.append(processed_item)
                    logger.info(f"Processed Multi-Image image (base64 length: {len(processed_item['image'])})")
                else:
                    processed_image_list.append(item)
            image_list = processed_image_list
        
        payload = {
            "model_name": self.model_name if self.model_name else "kling-v1-6",
            "image_list": image_list,
            "prompt": prompt,
            "negative_prompt": kwargs.get("negative_prompt", ""),
            "mode": params["mode"],
            "duration": str(kwargs.get("duration", "5")),
            "aspect_ratio": params["aspect_ratio"],
            "callback_url": kwargs.get("callback_url"),
            "external_task_id": kwargs.get("external_task_id")
        }
        access_key = kwargs.get('_access_key')
        secret_key = kwargs.get('_secret_key')
        return await self._send_request(url, payload, "multi-image2video", access_key, secret_key)

    async def _send_request(self, url: str, payload: Dict, task_type: str, access_key: Optional[str] = None, secret_key: Optional[str] = None) -> ServiceResponse:
        payload = {k: v for k, v in payload.items() if v is not None}

        try:
            headers = self._get_headers(access_key or self.access_key, secret_key or self.secret_key)

            async with httpx.AsyncClient(timeout=60.0) as client:
                resp = await client.post(url, headers=headers, json=payload)
                if resp.status_code >= 400:
                    try:
                        err_data = resp.json()
                        err_msg = err_data.get("message") or resp.text
                    except ValueError:
                        err_msg = resp.text
                    raise ValueError(f"Kling API Error ({resp.status_code}): {err_msg}")
                data = resp.json()

            if data.get("code") != 0:
                raise ValueError(f"Kling API Error: {data.get('message')}")

            task_data = data.get("data", {})
            task_id = task_data.get("task_id")

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

            return ServiceResponse(
                task_id=str(task_id),
                status=TaskStatus.QUEUED,
                meta={
                    "provider": "kling",
                    "task_type": task_type,
                    "raw_response": data
                }
            )
        except Exception as e:
            logger.error(f"Kling Request Failed: {e}")
            return ServiceResponse(
                task_id="failed",
                status=TaskStatus.FAILED,
                error=str(e)
            )

    async def check_status(self, task_id: str, **kwargs) -> ServiceResponse:
        """
        Check status. Since we don't persist the task type,
        we try endpoints in order of likelihood or user preference.
        """
        access_key = kwargs.get('access_key') or self.access_key
        secret_key = kwargs.get('secret_key') or self.secret_key

        endpoints = [
            ("text2video", f"{self.base_url}/videos/text2video/{task_id}"),
            ("image2video", f"{self.base_url}/videos/image2video/{task_id}"),
            ("omni", f"{self.base_url}/videos/omni-video/{task_id}"),
            ("multi-image", f"{self.base_url}/videos/multi-image2video/{task_id}")
        ]

        last_error = None
        found_data = None

        async with httpx.AsyncClient(timeout=30.0) as client:
            for name, url in endpoints:
                try:
                    headers = self._get_headers(access_key, secret_key)
                    resp = await client.get(url, headers=headers)

                    if resp.status_code == 200:
                        data = resp.json()
                        if data.get("code") == 0:
                            found_data = data
                            break
                    elif resp.status_code == 404:
                        continue

                except Exception as e:
                    last_error = e
                    continue

        if not found_data:
            return ServiceResponse(
                task_id=task_id,
                status=TaskStatus.UNKNOWN,
                error=f"Task not found in any endpoint. Last error: {last_error}"
            )

        return self._parse_status_response(found_data, task_id)

    def _parse_status_response(self, data: Dict, task_id: str) -> ServiceResponse:
        task_data = data.get("data", {})
        status_str = task_data.get("task_status", "").lower()
        
        status_map = {
            "submitted": TaskStatus.QUEUED,
            "processing": TaskStatus.PROCESSING,
            "succeed": TaskStatus.SUCCEEDED,
            "failed": TaskStatus.FAILED
        }
        
        status = status_map.get(status_str, TaskStatus.UNKNOWN)
        
        results = []
        if status == TaskStatus.SUCCEEDED:
             task_result = task_data.get("task_result", {})
             videos = task_result.get("videos", [])
             if videos:
                 for v in videos:
                     results.append(GenerationResult(
                         url=v.get("url"),
                         content_type="video/mp4",
                         meta=v
                     ))
        
        error_msg = None
        if status == TaskStatus.FAILED:
            error_msg = task_data.get("task_status_msg")
            
        return ServiceResponse(
            task_id=task_id,
            status=status,
            results=results,
            output={"video_url": results[0].url} if results else {},
            error=error_msg,
            meta=data
        )
