import aiohttp
from flask import current_app
from . import BaseVideoProvider


class VeoProvider(BaseVideoProvider):
    BASE_URL = "https://generativelanguage.googleapis.com/v1beta"

    def _get_api_key(self):
        from .key_helper import get_provider_key
        key, _ = get_provider_key("veo", "GOOGLE_GEMINI_API_KEY")
        return key

    async def generate_video(self, prompt: str, duration: int = 8, **kwargs) -> dict:
        api_key = self._get_api_key()
        if not api_key:
            return {"error": "Google Gemini API key not configured", "status": "error"}

        # Use fast generate endpoint for quicker results
        url = f"{self.BASE_URL}/models/veo-3.0-generate-001:predictLongRunning"
        payload = {
            "instances": [{"prompt": prompt}],
            "parameters": {
                "sampleCount": 1,
                "aspectRatio": "16:9",
                # Note: durationSeconds parameter causes API errors, omitting it
                # The API will use default duration
            },
        }
        headers = {"Content-Type": "application/json"}

        async with aiohttp.ClientSession() as session:
            async with session.post(f"{url}?key={api_key}", json=payload, headers=headers) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    operation_name = data.get("name", "")
                    return {"status": "processing", "operation_id": operation_name, "provider": "veo"}
                else:
                    error_text = await resp.text()
                    return {"status": "error", "error": error_text}

    async def check_status(self, operation_id: str) -> dict:
        api_key = self._get_api_key()
        url = f"{self.BASE_URL}/{operation_id}"

        async with aiohttp.ClientSession() as session:
            async with session.get(f"{url}?key={api_key}") as resp:
                data = await resp.json()
                if data.get("done"):
                    # Check for errors
                    if data.get("error"):
                        return {"status": "error", "error": data.get("error", {}).get("message", "Unknown error")}

                    videos = data.get("response", {}).get("predictions", [])
                    if videos and videos[0].get("videoUri"):
                        video_uri = videos[0].get("videoUri")

                        # Download the video bytes to avoid expired signed URLs
                        try:
                            async with session.get(video_uri) as video_resp:
                                if video_resp.status == 200:
                                    video_bytes = await video_resp.read()
                                    return {"status": "completed", "video_bytes": video_bytes}
                                else:
                                    return {"status": "error", "error": f"Failed to download video: {video_resp.status}"}
                        except Exception as e:
                            # If download fails, return the URI anyway
                            return {"status": "completed", "video_url": video_uri}

                    # Handle case where video might be in a different response structure
                    if data.get("response", {}).get("video"):
                        return {"status": "completed", "video_url": data["response"]["video"]}

                    return {"status": "error", "error": "No video generated"}
                # Still processing
                return {"status": "processing"}