"""
TimeLapseForge - Video Generation Providers
Supports Text-to-Video and Image-to-Video via multiple APIs.
All imports are lazy - no extra packages required at startup.
"""

import os
import io
import time
import base64
import tempfile
import requests
from PIL import Image
from typing import Optional, Dict, List, Any
from abc import ABC, abstractmethod


def _safe_import(package_name):
    try:
        import importlib
        return importlib.import_module(package_name)
    except ImportError:
        return None


def _require_import(package_name, pip_name=None):
    mod = _safe_import(package_name)
    if mod is None:
        pip = pip_name or package_name
        raise ImportError(
            "Package '" + pip + "' is not installed. "
            "Add it to requirements.txt or use a different provider."
        )
    return mod


class BaseVideoProvider(ABC):
    name = "base"
    display_name = "Base Video Provider"
    website = ""
    supports_t2v = False
    supports_i2v = False
    default_model = ""
    available_models = []
    requires_package = ""

    def __init__(self, api_key=""):
        self.api_key = api_key.strip()

    def text_to_video(self, prompt, width=1280, height=720,
                      duration=4, seed=None, model=None, **kwargs):
        raise NotImplementedError("This provider does not support text-to-video")

    def image_to_video(self, image, prompt="", width=1280, height=720,
                       duration=4, seed=None, model=None, **kwargs):
        raise NotImplementedError("This provider does not support image-to-video")

    @staticmethod
    def _image_to_base64(img, fmt="PNG"):
        buf = io.BytesIO()
        img.save(buf, format=fmt)
        return base64.b64encode(buf.getvalue()).decode("utf-8")

    @staticmethod
    def _image_to_bytes(img, fmt="PNG"):
        buf = io.BytesIO()
        img.save(buf, format=fmt)
        buf.seek(0)
        return buf

    @staticmethod
    def _save_from_url(url, suffix=".mp4"):
        resp = requests.get(url, timeout=600, stream=True)
        resp.raise_for_status()
        tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
        for chunk in resp.iter_content(chunk_size=8192):
            tmp.write(chunk)
        tmp.close()
        return tmp.name

    @staticmethod
    def _save_from_bytes(data, suffix=".mp4"):
        tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
        tmp.write(data)
        tmp.close()
        return tmp.name


# ============================================
# 1. STABILITY AI VIDEO
# ============================================

class StabilityVideoProvider(BaseVideoProvider):
    name = "stability_video"
    display_name = "Stability AI (Stable Video Diffusion)"
    website = "https://platform.stability.ai/account/keys"
    supports_t2v = False
    supports_i2v = True
    default_model = "svd"
    available_models = ["svd", "svd-xt"]
    requires_package = ""

    API_BASE = "https://api.stability.ai/v2beta"

    def image_to_video(self, image, prompt="", width=1024, height=576,
                       duration=4, seed=None, model=None, **kwargs):
        headers = {"Authorization": "Bearer " + self.api_key}
        buf = self._image_to_bytes(image.resize((1024, 576), Image.LANCZOS))

        data = {"seed": seed or 0, "cfg_scale": 2.5, "motion_bucket_id": 40}
        files = {"image": ("frame.png", buf, "image/png")}

        resp = requests.post(
            self.API_BASE + "/image-to-video",
            headers=headers, files=files, data=data, timeout=60,
        )
        resp.raise_for_status()
        gen_id = resp.json().get("id", "")

        for _ in range(60):
            time.sleep(10)
            poll = requests.get(
                self.API_BASE + "/image-to-video/result/" + gen_id,
                headers={
                    "Authorization": "Bearer " + self.api_key,
                    "Accept": "video/*",
                },
                timeout=30,
            )
            if poll.status_code == 200:
                return self._save_from_bytes(poll.content)
            elif poll.status_code != 202:
                poll.raise_for_status()

        raise TimeoutError("Stability video generation timed out")


# ============================================
# 2. REPLICATE VIDEO
# ============================================

class ReplicateVideoProvider(BaseVideoProvider):
    name = "replicate_video"
    display_name = "Replicate (Kling / CogVideo / AnimateDiff)"
    website = "https://replicate.com/account/api-tokens"
    supports_t2v = True
    supports_i2v = True
    default_model = "minimax/video-01"
    available_models = [
        "minimax/video-01",
        "minimax/video-01-live",
        "tencent/hunyuan-video",
        "kwaivgi/kling-v1.6-pro",
        "luosiallen/latent-consistency-model",
        "lucataco/animate-diff",
        "zsxkib/cogvideo-5b",
    ]
    requires_package = "replicate"

    def text_to_video(self, prompt, width=1280, height=720,
                      duration=4, seed=None, model=None, **kwargs):
        replicate_mod = _require_import("replicate")
        client = replicate_mod.Client(api_token=self.api_key)
        model_id = model or self.default_model

        input_params = {"prompt": prompt}
        if seed is not None:
            input_params["seed"] = seed

        output = client.run(model_id, input=input_params)
        if isinstance(output, list):
            url = str(output[0])
        elif hasattr(output, "url"):
            url = str(output.url)
        else:
            url = str(output)

        return self._save_from_url(url)

    def image_to_video(self, image, prompt="", width=1280, height=720,
                       duration=4, seed=None, model=None, **kwargs):
        replicate_mod = _require_import("replicate")
        client = replicate_mod.Client(api_token=self.api_key)
        model_id = model or self.default_model

        buf = self._image_to_bytes(image)
        input_params = {"image": buf}
        if prompt:
            input_params["prompt"] = prompt
        if seed is not None:
            input_params["seed"] = seed

        output = client.run(model_id, input=input_params)
        if isinstance(output, list):
            url = str(output[0])
        elif hasattr(output, "url"):
            url = str(output.url)
        else:
            url = str(output)

        return self._save_from_url(url)


# ============================================
# 3. FAL.AI VIDEO
# ============================================

class FalVideoProvider(BaseVideoProvider):
    name = "fal_video"
    display_name = "Fal.ai (Kling / Veo / Luma / MiniMax)"
    website = "https://fal.ai/dashboard/keys"
    supports_t2v = True
    supports_i2v = True
    default_model = "fal-ai/kling-video/v2.0/master"
    available_models = [
        "fal-ai/kling-video/v2.0/master",
        "fal-ai/kling-video/v1.5/pro/image-to-video",
        "fal-ai/minimax-video/image-to-video",
        "fal-ai/minimax-video",
        "fal-ai/luma-dream-machine",
        "fal-ai/luma-dream-machine/image-to-video",
        "fal-ai/veo2",
        "fal-ai/hunyuan-video",
        "fal-ai/cogvideox-5b",
        "fal-ai/wan/v2.1/1.3b",
        "fal-ai/ltx-video",
    ]
    requires_package = "fal_client"

    def text_to_video(self, prompt, width=1280, height=720,
                      duration=4, seed=None, model=None, **kwargs):
        fal_client = _require_import("fal_client", "fal-client")
        os.environ["FAL_KEY"] = self.api_key
        model_id = model or self.default_model

        arguments = {"prompt": prompt}
        if seed is not None:
            arguments["seed"] = seed
        if duration:
            arguments["duration"] = str(duration) + "s"

        result = fal_client.subscribe(model_id, arguments=arguments)
        video_data = result.get("video", result.get("data", {}))
        if isinstance(video_data, dict):
            url = video_data.get("url", "")
        elif isinstance(video_data, str):
            url = video_data
        else:
            raise ValueError("No video URL in response")

        return self._save_from_url(url)

    def image_to_video(self, image, prompt="", width=1280, height=720,
                       duration=4, seed=None, model=None, **kwargs):
        fal_client = _require_import("fal_client", "fal-client")
        os.environ["FAL_KEY"] = self.api_key

        b64 = self._image_to_base64(image)
        data_uri = "data:image/png;base64," + b64

        model_id = model or "fal-ai/kling-video/v1.5/pro/image-to-video"
        arguments = {"image_url": data_uri}
        if prompt:
            arguments["prompt"] = prompt
        if seed is not None:
            arguments["seed"] = seed
        if duration:
            arguments["duration"] = str(duration) + "s"

        result = fal_client.subscribe(model_id, arguments=arguments)
        video_data = result.get("video", result.get("data", {}))
        if isinstance(video_data, dict):
            url = video_data.get("url", "")
        elif isinstance(video_data, str):
            url = video_data
        else:
            raise ValueError("No video URL in response")

        return self._save_from_url(url)


# ============================================
# 4. RUNWAY
# ============================================

class RunwayVideoProvider(BaseVideoProvider):
    name = "runway_video"
    display_name = "Runway (Gen-3 Alpha / Gen-4)"
    website = "https://dev.runwayml.com/"
    supports_t2v = True
    supports_i2v = True
    default_model = "gen4_turbo"
    available_models = ["gen3a_turbo", "gen4_turbo"]
    requires_package = ""

    API_BASE = "https://api.dev.runwayml.com/v1"

    def _poll_task(self, task_id):
        for _ in range(120):
            time.sleep(5)
            resp = requests.get(
                self.API_BASE + "/tasks/" + task_id,
                headers={"Authorization": "Bearer " + self.api_key},
                timeout=30,
            )
            resp.raise_for_status()
            data = resp.json()
            status = data.get("status", "")
            if status == "SUCCEEDED":
                output = data.get("output", [])
                if output:
                    return self._save_from_url(output[0])
                raise ValueError("Task succeeded but no output URL")
            elif status in ("FAILED", "CANCELLED"):
                raise RuntimeError("Runway task " + status + ": " + str(data.get("failure", "")))
        raise TimeoutError("Runway task timed out")

    def text_to_video(self, prompt, width=1280, height=720,
                      duration=5, seed=None, model=None, **kwargs):
        headers = {
            "Authorization": "Bearer " + self.api_key,
            "Content-Type": "application/json",
            "X-Runway-Version": "2024-11-06",
        }
        payload = {
            "model": model or self.default_model,
            "promptText": prompt,
            "duration": min(duration, 10),
            "ratio": "16:9",
        }
        if seed is not None:
            payload["seed"] = seed

        resp = requests.post(
            self.API_BASE + "/image_to_video",
            headers=headers, json=payload, timeout=60,
        )
        resp.raise_for_status()
        task_id = resp.json().get("id", "")
        return self._poll_task(task_id)

    def image_to_video(self, image, prompt="", width=1280, height=720,
                       duration=5, seed=None, model=None, **kwargs):
        headers = {
            "Authorization": "Bearer " + self.api_key,
            "Content-Type": "application/json",
            "X-Runway-Version": "2024-11-06",
        }
        b64 = self._image_to_base64(image)
        data_uri = "data:image/png;base64," + b64

        payload = {
            "model": model or self.default_model,
            "promptImage": data_uri,
            "duration": min(duration, 10),
            "ratio": "16:9",
        }
        if prompt:
            payload["promptText"] = prompt
        if seed is not None:
            payload["seed"] = seed

        resp = requests.post(
            self.API_BASE + "/image_to_video",
            headers=headers, json=payload, timeout=60,
        )
        resp.raise_for_status()
        task_id = resp.json().get("id", "")
        return self._poll_task(task_id)


# ============================================
# 5. LUMA DREAM MACHINE
# ============================================

class LumaVideoProvider(BaseVideoProvider):
    name = "luma_video"
    display_name = "Luma Dream Machine (Ray 2)"
    website = "https://lumalabs.ai/api"
    supports_t2v = True
    supports_i2v = True
    default_model = "ray-2"
    available_models = ["ray-2", "ray-2-flash"]
    requires_package = ""

    API_BASE = "https://api.lumalabs.ai/dream-machine/v1"

    def _poll_generation(self, gen_id):
        for _ in range(120):
            time.sleep(5)
            resp = requests.get(
                self.API_BASE + "/generations/" + gen_id,
                headers={"Authorization": "Bearer " + self.api_key},
                timeout=30,
            )
            resp.raise_for_status()
            data = resp.json()
            state = data.get("state", "")
            if state == "completed":
                assets = data.get("assets", {})
                video_url = assets.get("video", "")
                if video_url:
                    return self._save_from_url(video_url)
                raise ValueError("Completed but no video URL")
            elif state == "failed":
                raise RuntimeError("Luma generation failed: " + str(data.get("failure_reason", "")))
        raise TimeoutError("Luma generation timed out")

    def text_to_video(self, prompt, width=1280, height=720,
                      duration=5, seed=None, model=None, **kwargs):
        headers = {
            "Authorization": "Bearer " + self.api_key,
            "Content-Type": "application/json",
        }
        payload = {
            "prompt": prompt,
            "model": model or self.default_model,
            "resolution": "720p",
        }

        resp = requests.post(
            self.API_BASE + "/generations",
            headers=headers, json=payload, timeout=60,
        )
        resp.raise_for_status()
        gen_id = resp.json().get("id", "")
        return self._poll_generation(gen_id)

    def image_to_video(self, image, prompt="", width=1280, height=720,
                       duration=5, seed=None, model=None, **kwargs):
        headers = {
            "Authorization": "Bearer " + self.api_key,
            "Content-Type": "application/json",
        }
        b64 = self._image_to_base64(image)
        data_uri = "data:image/png;base64," + b64

        payload = {
            "prompt": prompt or "animate this image naturally",
            "model": model or self.default_model,
            "keyframes": {
                "frame0": {"type": "image", "url": data_uri}
            },
        }

        resp = requests.post(
            self.API_BASE + "/generations",
            headers=headers, json=payload, timeout=60,
        )
        resp.raise_for_status()
        gen_id = resp.json().get("id", "")
        return self._poll_generation(gen_id)


# ============================================
# 6. MINIMAX (HAILUO)
# ============================================

class MinimaxVideoProvider(BaseVideoProvider):
    name = "minimax_video"
    display_name = "MiniMax / Hailuo AI"
    website = "https://platform.minimaxi.com/"
    supports_t2v = True
    supports_i2v = False
    default_model = "T2V-01"
    available_models = ["T2V-01", "T2V-01-HD"]
    requires_package = ""

    API_BASE = "https://api.minimaxi.chat/v1"

    def text_to_video(self, prompt, width=1280, height=720,
                      duration=4, seed=None, model=None, **kwargs):
        headers = {
            "Authorization": "Bearer " + self.api_key,
            "Content-Type": "application/json",
        }
        payload = {
            "model": model or self.default_model,
            "prompt": prompt,
        }

        resp = requests.post(
            self.API_BASE + "/video_generation",
            headers=headers, json=payload, timeout=60,
        )
        resp.raise_for_status()
        task_id = resp.json().get("task_id", "")

        for _ in range(120):
            time.sleep(5)
            poll = requests.get(
                self.API_BASE + "/query/video_generation?task_id=" + task_id,
                headers=headers, timeout=30,
            )
            poll.raise_for_status()
            data = poll.json()
            status = data.get("status", "")
            if status == "Success":
                file_id = data.get("file_id", "")
                if file_id:
                    dl = requests.get(
                        self.API_BASE + "/files/retrieve?file_id=" + file_id,
                        headers=headers, timeout=30,
                    )
                    dl.raise_for_status()
                    dl_data = dl.json()
                    download_url = dl_data.get("file", {}).get("download_url", "")
                    if download_url:
                        return self._save_from_url(download_url)
            elif status == "Fail":
                raise RuntimeError("MiniMax video failed")

        raise TimeoutError("MiniMax video timed out")


# ============================================
# 7. HUGGINGFACE VIDEO
# ============================================

class HuggingFaceVideoProvider(BaseVideoProvider):
    name = "huggingface_video"
    display_name = "HuggingFace Inference (CogVideoX / Wan)"
    website = "https://huggingface.co/settings/tokens"
    supports_t2v = True
    supports_i2v = False
    default_model = "THUDM/CogVideoX-5b"
    available_models = [
        "THUDM/CogVideoX-5b",
        "THUDM/CogVideoX-2b",
        "Wan-AI/Wan2.1-T2V-14B",
        "ali-vilab/text-to-video-ms-1.7b",
    ]
    requires_package = ""

    API_BASE = "https://api-inference.huggingface.co/models"

    def text_to_video(self, prompt, width=512, height=512,
                      duration=4, seed=None, model=None, **kwargs):
        model_id = model or self.default_model
        url = self.API_BASE + "/" + model_id
        headers = {"Authorization": "Bearer " + self.api_key}
        payload = {"inputs": prompt}

        resp = requests.post(url, headers=headers, json=payload, timeout=300)
        if resp.status_code == 503:
            time.sleep(30)
            resp = requests.post(url, headers=headers, json=payload, timeout=300)

        resp.raise_for_status()

        ct = resp.headers.get("Content-Type", "")
        if "video" in ct or "octet" in ct:
            return self._save_from_bytes(resp.content)

        raise ValueError("HuggingFace did not return video data")


# ============================================
# VIDEO PROVIDER REGISTRY
# ============================================

VIDEO_PROVIDERS = {
    "stability_video": StabilityVideoProvider,
    "replicate_video": ReplicateVideoProvider,
    "fal_video": FalVideoProvider,
    "runway_video": RunwayVideoProvider,
    "luma_video": LumaVideoProvider,
    "minimax_video": MinimaxVideoProvider,
    "huggingface_video": HuggingFaceVideoProvider,
}

VIDEO_PROVIDER_DISPLAY_NAMES = {
    cls.display_name: key for key, cls in VIDEO_PROVIDERS.items()
}


def get_video_provider(provider_name, api_key):
    cls = VIDEO_PROVIDERS.get(provider_name)
    if cls is None:
        raise ValueError("Unknown video provider: " + str(provider_name))
    return cls(api_key=api_key)


def get_video_provider_info():
    info = []
    for key, cls in VIDEO_PROVIDERS.items():
        info.append({
            "name": key,
            "display_name": cls.display_name,
            "website": cls.website,
            "supports_t2v": cls.supports_t2v,
            "supports_i2v": cls.supports_i2v,
            "default_model": cls.default_model,
            "available_models": cls.available_models,
        })
    return info


def get_video_models_for_provider(provider_name):
    cls = VIDEO_PROVIDERS.get(provider_name)
    return cls.available_models if cls else []