#!/usr/bin/env python3
"""
Video generation provider abstraction for Stage 10.

Supports 6 providers:
- Kling Direct (v2.6): Direct API with JWT auth
- Kling O1 (OmniVideo): Direct API with reference image support
- Pollo Kling (v2.6): Via Pollo.ai proxy
- Pollo Kling O1: Via Pollo.ai proxy with reference image support
- Pollo Veo3: Google Veo 3 via Pollo.ai
- LaoZhang Veo3: Google Veo 3 via LaoZhang API

All providers support start + end frame mode for frame interpolation.
Kling O1 (direct and Pollo) additionally supports reference images for character consistency.
"""

import os
import sys
import time
import base64
import json
import jwt
import requests
from pathlib import Path
from typing import Optional, Dict
from abc import ABC, abstractmethod

# Add repo root to path for pollo_proxy and laozhang_proxy imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "scripts"))

from pollo_proxy import (
    call_kling_v26_pollo,
    call_kling_o1_pollo,
    call_veo3_pollo,
    poll_kling_pollo,
    poll_veo3_pollo,
    download_video_pollo
)
from laozhang_proxy import call_veo3_laozhang, download_video_laozhang


class VideoProvider(ABC):
    """Abstract base class for video providers."""

    @abstractmethod
    def generate_video(
        self,
        prompt: str,
        start_frame_path: str,
        end_frame_path: Optional[str] = None,
        duration: int = 5,
        generate_audio: bool = True,
        output_path: Optional[str] = None
    ) -> Dict:
        """
        Generate video from frame(s).

        Args:
            prompt: Video generation prompt
            start_frame_path: Path to starting frame
            end_frame_path: Optional path to ending frame (for interpolation)
            duration: Duration in seconds (5 or 10 for Kling, 8 for Veo3)
            generate_audio: Generate audio track (Kling v2.6+ only)
            output_path: Where to save the video

        Returns:
            {
                "status": "success" | "error",
                "task_id": str,
                "video_url": str,
                "video_path": str,
                "duration_actual": float,
                "error": str (if error)
            }
        """
        pass


class KlingDirectProvider(VideoProvider):
    """Kling v2.6 Direct API provider."""

    def __init__(self):
        self.api_base = os.getenv("KLING_API_BASE", "https://api-singapore.klingai.com")
        self.access_key = os.getenv("KLING_ACCESS_KEY")
        self.secret_key = os.getenv("KLING_SECRET_KEY")

        if not self.access_key or not self.secret_key:
            raise ValueError("KLING_ACCESS_KEY and KLING_SECRET_KEY must be set")

    def _generate_jwt_token(self) -> str:
        """Generate JWT token for authentication."""
        payload = {
            "iss": self.access_key,
            "exp": int(time.time()) + 1800,  # 30 min
            "nbf": int(time.time()) - 5  # 5 seconds ago
        }
        return jwt.encode(payload, self.secret_key, algorithm="HS256")

    def _image_to_base64(self, image_path: str) -> str:
        """Convert image to base64 (no data URI prefix)."""
        with open(image_path, 'rb') as f:
            return base64.b64encode(f.read()).decode('utf-8')

    def generate_video(
        self,
        prompt: str,
        start_frame_path: str,
        end_frame_path: Optional[str] = None,
        duration: int = 5,
        generate_audio: bool = True,
        output_path: Optional[str] = None,
        resolution: str = "720p"
    ) -> Dict:
        """Generate video using Kling Direct API."""
        print(f"\n=== Kling Direct Provider ===")
        print(f"Duration: {duration}s")
        print(f"Audio: {'on' if generate_audio else 'off'}")
        print(f"Frames: start" + (" + end" if end_frame_path else ""))
        # Note: Kling Direct doesn't support resolution parameter, ignoring

        # Generate JWT token
        token = self._generate_jwt_token()

        # Encode start frame
        image_base64 = self._image_to_base64(start_frame_path)

        # Build payload
        payload = {
            "model_name": "kling-v2-6",
            "mode": "pro",  # Required for audio
            "image": image_base64,
            "prompt": prompt[:2500],
            "duration": str(duration),  # Must be string
            "sound": "on" if generate_audio else "off",
        }

        # Add end frame if provided
        if end_frame_path:
            image_tail_base64 = self._image_to_base64(end_frame_path)
            payload["image_tail"] = image_tail_base64
            print(f"  End frame: {Path(end_frame_path).name}")

        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }

        # Submit task
        print(f"Submitting to Kling API...")
        response = requests.post(
            f"{self.api_base}/v1/videos/image2video",
            headers=headers,
            json=payload,
            timeout=60
        )

        if response.status_code != 200:
            return {
                "status": "error",
                "error": f"API error {response.status_code}: {response.text[:500]}"
            }

        result = response.json()

        if result.get("code") != 0:
            return {
                "status": "error",
                "error": result.get("message", "Unknown error")
            }

        # Extract task ID
        data = result.get("data", {})
        task_id = data.get("task_id")
        if not task_id:
            return {
                "status": "error",
                "error": f"No task_id in response: {result}"
            }

        print(f"Task created: {task_id}")

        # Poll for completion
        print(f"Polling for completion...")
        video_result = self._poll_status(task_id, timeout=600)

        if video_result["status"] != "success":
            return video_result

        # Download video
        video_url = video_result["video_url"]
        if output_path:
            self._download_video(video_url, output_path)
            return {
                "status": "success",
                "task_id": task_id,
                "video_url": video_url,
                "video_path": output_path,
                "duration_actual": video_result.get("duration", duration)
            }
        else:
            return {
                "status": "success",
                "task_id": task_id,
                "video_url": video_url,
                "duration_actual": video_result.get("duration", duration)
            }

    def _poll_status(self, task_id: str, poll_interval: int = 10, timeout: int = 600) -> Dict:
        """Poll Kling API for task completion."""
        token = self._generate_jwt_token()
        status_url = f"{self.api_base}/v1/videos/image2video/{task_id}"
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }

        start_time = time.time()

        while time.time() - start_time < timeout:
            response = requests.get(status_url, headers=headers, timeout=30)

            if response.status_code != 200:
                return {
                    "status": "error",
                    "error": f"Status check error: {response.status_code}"
                }

            result = response.json()

            if result.get("code") != 0:
                return {
                    "status": "error",
                    "error": result.get("message", "Unknown error")
                }

            data = result.get("data", {})
            task_status = data.get("task_status")

            print(f"  Status: {task_status}")

            if task_status == "succeed":
                # Extract video URL
                task_result = data.get("task_result", {})
                videos = task_result.get("videos", [])
                if not videos:
                    return {
                        "status": "error",
                        "error": "No videos in response"
                    }

                video_info = videos[0]
                return {
                    "status": "success",
                    "video_url": video_info.get("url"),
                    "duration": video_info.get("duration")
                }

            elif task_status == "failed":
                return {
                    "status": "error",
                    "error": data.get("task_status_msg", "Generation failed")
                }

            time.sleep(poll_interval)

        return {
            "status": "error",
            "error": f"Timeout after {timeout}s"
        }

    def _download_video(self, url: str, output_path: str):
        """Download video from URL."""
        print(f"Downloading video...")
        response = requests.get(url, stream=True, timeout=120)
        response.raise_for_status()

        with open(output_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        print(f"Saved: {output_path}")


class KlingO1DirectProvider(VideoProvider):
    """Kling O1 (OmniVideo) Direct API provider with reference image support."""

    def __init__(self):
        self.api_base = os.getenv("KLING_API_BASE", "https://api-singapore.klingai.com")
        self.access_key = os.getenv("KLING_ACCESS_KEY")
        self.secret_key = os.getenv("KLING_SECRET_KEY")

        if not self.access_key or not self.secret_key:
            raise ValueError("KLING_ACCESS_KEY and KLING_SECRET_KEY must be set")

    def _generate_jwt_token(self) -> str:
        """Generate JWT token for authentication."""
        payload = {
            "iss": self.access_key,
            "exp": int(time.time()) + 1800,
            "nbf": int(time.time()) - 5
        }
        return jwt.encode(payload, self.secret_key, algorithm="HS256")

    def _image_to_base64(self, image_path: str) -> str:
        """Convert image to base64 (no data URI prefix)."""
        with open(image_path, 'rb') as f:
            return base64.b64encode(f.read()).decode('utf-8')

    def generate_video(
        self,
        prompt: str,
        start_frame_path: str,
        end_frame_path: Optional[str] = None,
        duration: int = 5,
        generate_audio: bool = False,
        output_path: Optional[str] = None,
        resolution: str = "1080p",
        reference_images: Optional[list] = None
    ) -> Dict:
        """
        Generate video using Kling O1 OmniVideo API.

        Reference images are passed as additional context for character/scene
        consistency. They are referenced in the prompt via <<<image_N>>> syntax.

        Args:
            reference_images: List of file paths to reference images (e.g. characteristic shot).
                              These are injected into image_list WITHOUT type (as references)
                              and referenced in the prompt as <<<image_1>>>, <<<image_2>>>, etc.
        """
        ref_count = len(reference_images) if reference_images else 0
        print(f"\n=== Kling O1 Direct Provider ===")
        print(f"Duration: {duration}s")
        print(f"Frames: start" + (" + end" if end_frame_path else ""))
        print(f"Reference images: {ref_count}")

        token = self._generate_jwt_token()

        # Build image_list: first_frame + optional end_frame + reference images
        image_list = []

        # Start frame (first_frame type — not referenced in prompt)
        start_b64 = self._image_to_base64(start_frame_path)
        image_list.append({
            "image_url": start_b64,
            "type": "first_frame"
        })

        # End frame (if provided)
        if end_frame_path:
            end_b64 = self._image_to_base64(end_frame_path)
            image_list.append({
                "image_url": end_b64,
                "type": "end_frame"
            })

        # Reference images (no type — referenced in prompt as <<<image_N>>>)
        if reference_images:
            for ref_path in reference_images:
                ref_b64 = self._image_to_base64(ref_path)
                image_list.append({
                    "image_url": ref_b64
                })
                print(f"  Reference: {Path(ref_path).name}")

        # Build payload
        payload = {
            "model_name": "kling-video-o1",
            "mode": "pro",
            "prompt": prompt[:2500],
            "image_list": image_list,
            "duration": str(duration),
        }

        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }

        # Submit task
        print(f"Submitting to Kling O1 API...")
        response = requests.post(
            f"{self.api_base}/v1/videos/omni-video",
            headers=headers,
            json=payload,
            timeout=60
        )

        if response.status_code != 200:
            return {
                "status": "error",
                "error": f"API error {response.status_code}: {response.text[:500]}"
            }

        result = response.json()

        if result.get("code") != 0:
            return {
                "status": "error",
                "error": result.get("message", "Unknown error")
            }

        data = result.get("data", {})
        task_id = data.get("task_id")
        if not task_id:
            return {
                "status": "error",
                "error": f"No task_id in response: {result}"
            }

        print(f"Task created: {task_id}")

        # Poll for completion
        print(f"Polling for completion...")
        video_result = self._poll_status(task_id, timeout=600)

        if video_result["status"] != "success":
            return video_result

        # Download video
        video_url = video_result["video_url"]
        if output_path:
            self._download_video(video_url, output_path)
            return {
                "status": "success",
                "task_id": task_id,
                "video_url": video_url,
                "video_path": output_path,
                "duration_actual": video_result.get("duration", duration)
            }
        else:
            return {
                "status": "success",
                "task_id": task_id,
                "video_url": video_url,
                "duration_actual": video_result.get("duration", duration)
            }

    def _poll_status(self, task_id: str, poll_interval: int = 10, timeout: int = 600) -> Dict:
        """Poll Kling O1 API for task completion."""
        token = self._generate_jwt_token()
        status_url = f"{self.api_base}/v1/videos/omni-video/{task_id}"
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }

        start_time = time.time()

        while time.time() - start_time < timeout:
            # Refresh token periodically (every 25 min)
            if time.time() - start_time > 1500:
                token = self._generate_jwt_token()
                headers["Authorization"] = f"Bearer {token}"

            response = requests.get(status_url, headers=headers, timeout=30)

            if response.status_code != 200:
                return {
                    "status": "error",
                    "error": f"Status check error: {response.status_code}"
                }

            result = response.json()

            if result.get("code") != 0:
                return {
                    "status": "error",
                    "error": result.get("message", "Unknown error")
                }

            data = result.get("data", {})
            task_status = data.get("task_status")

            print(f"  Status: {task_status}")

            if task_status == "succeed":
                task_result = data.get("task_result", {})
                videos = task_result.get("videos", [])
                if not videos:
                    return {
                        "status": "error",
                        "error": "No videos in response"
                    }

                video_info = videos[0]
                return {
                    "status": "success",
                    "video_url": video_info.get("url"),
                    "duration": video_info.get("duration")
                }

            elif task_status == "failed":
                return {
                    "status": "error",
                    "error": data.get("task_status_msg", "Generation failed")
                }

            time.sleep(poll_interval)

        return {
            "status": "error",
            "error": f"Timeout after {timeout}s"
        }

    def _download_video(self, url: str, output_path: str):
        """Download video from URL."""
        print(f"Downloading video...")
        response = requests.get(url, stream=True, timeout=120)
        response.raise_for_status()

        with open(output_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        print(f"Saved: {output_path}")


class PolloKlingProvider(VideoProvider):
    """Kling via Pollo.ai proxy."""

    def __init__(self, model: str = "kling-v26"):
        """
        Args:
            model: "kling-v26" or "kling-o1"
        """
        self.model = model
        self.api_key = os.getenv("POLLO_KEY")
        if not self.api_key:
            raise ValueError("POLLO_KEY environment variable not set")

    def generate_video(
        self,
        prompt: str,
        start_frame_path: str,
        end_frame_path: Optional[str] = None,
        duration: int = 5,
        generate_audio: bool = True,
        output_path: Optional[str] = None,
        resolution: str = "720p"
    ) -> Dict:
        """Generate video using Pollo Kling proxy."""
        print(f"\n=== Pollo Kling Provider ({self.model}) ===")
        print(f"Duration: {duration}s")
        print(f"Audio: {generate_audio}")
        # Note: Pollo Kling doesn't support resolution parameter, ignoring
        print(f"Frames: start" + (" + end" if end_frame_path else ""))

        try:
            # Call Pollo API
            if self.model == "kling-v26":
                result = call_kling_v26_pollo(
                    prompt=prompt,
                    api_key=self.api_key,
                    first_frame_path=start_frame_path,
                    last_frame_path=end_frame_path,
                    length=duration,
                    generate_audio=generate_audio
                )
            else:  # kling-o1
                result = call_kling_o1_pollo(
                    prompt=prompt,
                    api_key=self.api_key,
                    first_frame_path=start_frame_path,
                    last_frame_path=end_frame_path,
                    aspect_ratio="9:16",
                    length=duration
                )

            task_id = result.get("task_id")
            if not task_id:
                return {
                    "status": "error",
                    "error": "No task_id in response"
                }

            print(f"Task created: {task_id}")

            # Poll for completion
            print(f"Polling for completion...")
            poll_result = poll_kling_pollo(task_id, self.api_key, timeout=600)

            video_url = poll_result.get("uri")
            if not video_url:
                return {
                    "status": "error",
                    "error": "No video URL in result"
                }

            # Download video
            if output_path:
                download_video_pollo(video_url, output_path)
                return {
                    "status": "success",
                    "task_id": task_id,
                    "video_url": video_url,
                    "video_path": output_path,
                    "duration_actual": duration
                }
            else:
                return {
                    "status": "success",
                    "task_id": task_id,
                    "video_url": video_url,
                    "duration_actual": duration
                }

        except Exception as e:
            return {
                "status": "error",
                "error": str(e)
            }


class PolloKlingO1Provider(VideoProvider):
    """Kling O1 via Pollo.ai proxy with reference image support.

    Uses the `images` array (1-7 URLs) for reference images that O1 uses
    for character/scene consistency. Start frame goes via `image` field.
    """

    def __init__(self):
        self.api_key = os.getenv("POLLO_KEY")
        if not self.api_key:
            raise ValueError("POLLO_KEY environment variable not set")

    def generate_video(
        self,
        prompt: str,
        start_frame_path: str,
        end_frame_path: Optional[str] = None,
        duration: int = 5,
        generate_audio: bool = False,
        output_path: Optional[str] = None,
        resolution: str = "1080p",
        reference_images: Optional[list] = None
    ) -> Dict:
        """Generate video using Pollo Kling O1 with reference images.

        Args:
            reference_images: List of file paths to reference images (characteristic
                shot, person photos). Converted to public URLs. Max 7.
        """
        ref_count = len(reference_images) if reference_images else 0
        print(f"\n=== Pollo Kling O1 Provider ===")
        print(f"Duration: {duration}s")
        print(f"Frames: start" + (" + end" if end_frame_path else ""))
        print(f"Reference images: {ref_count}")

        try:
            result = call_kling_o1_pollo(
                prompt=prompt,
                api_key=self.api_key,
                first_frame_path=start_frame_path,
                last_frame_path=end_frame_path,
                aspect_ratio="9:16",
                length=duration,
                reference_image_paths=reference_images
            )

            task_id = result.get("task_id")
            if not task_id:
                return {
                    "status": "error",
                    "error": "No task_id in response"
                }

            print(f"Task created: {task_id}")

            # Poll for completion
            print(f"Polling for completion...")
            poll_result = poll_kling_pollo(task_id, self.api_key, timeout=600)

            video_url = poll_result.get("uri")
            if not video_url:
                return {
                    "status": "error",
                    "error": "No video URL in result"
                }

            # Download video
            if output_path:
                download_video_pollo(video_url, output_path)
                return {
                    "status": "success",
                    "task_id": task_id,
                    "video_url": video_url,
                    "video_path": output_path,
                    "duration_actual": duration
                }
            else:
                return {
                    "status": "success",
                    "task_id": task_id,
                    "video_url": video_url,
                    "duration_actual": duration
                }

        except Exception as e:
            return {
                "status": "error",
                "error": str(e)
            }


class PolloVeo3Provider(VideoProvider):
    """Google Veo3 via Pollo.ai proxy."""

    def __init__(self):
        self.api_key = os.getenv("POLLO_KEY")
        if not self.api_key:
            raise ValueError("POLLO_KEY environment variable not set")

    def generate_video(
        self,
        prompt: str,
        start_frame_path: str,
        end_frame_path: Optional[str] = None,
        duration: int = 8,
        generate_audio: bool = False,
        output_path: Optional[str] = None,
        resolution: str = "1080p"
    ) -> Dict:
        """Generate video using Pollo Veo3 proxy."""
        print(f"\n=== Pollo Veo3 Provider ===")
        print(f"Duration: {duration}s (fixed)")
        print(f"Resolution: {resolution}")
        print(f"Frames: start" + (" + end" if end_frame_path else ""))

        try:
            # Call Pollo API
            result = call_veo3_pollo(
                prompt=prompt,
                api_key=self.api_key,
                first_frame_path=start_frame_path,
                last_frame_path=end_frame_path,
                aspect_ratio="9:16",
                length=duration,
                resolution=resolution,
                generate_audio=generate_audio
            )

            task_id = result.get("task_id")
            if not task_id:
                return {
                    "status": "error",
                    "error": "No task_id in response"
                }

            print(f"Task created: {task_id}")

            # Poll for completion
            print(f"Polling for completion...")
            poll_result = poll_veo3_pollo(task_id, self.api_key, timeout=600)

            video_url = poll_result.get("uri")
            if not video_url:
                return {
                    "status": "error",
                    "error": "No video URL in result"
                }

            # Download video
            if output_path:
                download_video_pollo(video_url, output_path)
                return {
                    "status": "success",
                    "task_id": task_id,
                    "video_url": video_url,
                    "video_path": output_path,
                    "duration_actual": duration
                }
            else:
                return {
                    "status": "success",
                    "task_id": task_id,
                    "video_url": video_url,
                    "duration_actual": duration
                }

        except Exception as e:
            return {
                "status": "error",
                "error": str(e)
            }


class LaoZhangVeo3Provider(VideoProvider):
    """Google Veo3 via LaoZhang API."""

    def __init__(self):
        self.api_key = os.getenv("LAOZHANG_KEY")
        if not self.api_key:
            raise ValueError("LAOZHANG_KEY environment variable not set")

    def generate_video(
        self,
        prompt: str,
        start_frame_path: str,
        end_frame_path: Optional[str] = None,
        duration: int = 8,
        generate_audio: bool = False,
        output_path: Optional[str] = None,
        resolution: str = "720p"
    ) -> Dict:
        """Generate video using LaoZhang Veo3 API."""
        print(f"\n=== LaoZhang Veo3 Provider ===")
        # Note: LaoZhang Veo3 may support resolution, but not implemented yet
        print(f"Duration: {duration}s")
        print(f"Frames: start" + (" + end" if end_frame_path else ""))

        try:
            # Call LaoZhang API
            result = call_veo3_laozhang(
                prompt=prompt,
                api_key=self.api_key,
                first_frame_path=start_frame_path,
                last_frame_path=end_frame_path,
                model="veo-3.1-fast-fl",
                duration=duration
            )

            video_url = result.get("uri")
            if not video_url:
                return {
                    "status": "error",
                    "error": "No video URL in result"
                }

            print(f"Video URL: {video_url}")

            # Download video
            if output_path:
                download_video_laozhang(video_url, output_path)
                return {
                    "status": "success",
                    "task_id": result.get("id", "unknown"),
                    "video_url": video_url,
                    "video_path": output_path,
                    "duration_actual": duration
                }
            else:
                return {
                    "status": "success",
                    "task_id": result.get("id", "unknown"),
                    "video_url": video_url,
                    "duration_actual": duration
                }

        except Exception as e:
            return {
                "status": "error",
                "error": str(e)
            }


def get_provider(provider_name: str, **kwargs) -> VideoProvider:
    """
    Get video provider by name.

    Args:
        provider_name: One of "kling-direct", "pollo-kling", "pollo-veo3", "laozhang-veo3"
        **kwargs: Provider-specific options (e.g., model for pollo-kling)

    Returns:
        VideoProvider instance

    Raises:
        ValueError: If provider unknown
    """
    if provider_name == "kling-direct":
        return KlingDirectProvider()
    elif provider_name == "kling-o1":
        return KlingO1DirectProvider()
    elif provider_name == "pollo-kling":
        model = kwargs.get("model", "kling-v26")
        return PolloKlingProvider(model=model)
    elif provider_name == "pollo-kling-o1":
        return PolloKlingO1Provider()
    elif provider_name == "pollo-veo3":
        return PolloVeo3Provider()
    elif provider_name == "laozhang-veo3":
        return LaoZhangVeo3Provider()
    else:
        raise ValueError(f"Unknown provider: {provider_name}")
