#!/usr/bin/env python3
"""
Genesis Veo Generator
=====================
Integrates Veo 3.1 for high-quality video generation.

Features:
- Text-to-video generation
- Asynchronous operation polling
- Optional audio generation
- Cost tracking

Usage:
    from veo_generator import VeoAgent
    agent = VeoAgent()
    result = agent.generate("A cinematic drone shot of a futuristic city")
"""

import os
import sys
import json
import time
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Any, List

# Add genesis-system to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from google import genai
from google.genai import types
from budget_manager import BudgetManager

# Load config
CONFIG_PATH = os.path.join(os.path.dirname(__file__), "genesis_config.json")
with open(CONFIG_PATH) as f:
    CONFIG = json.load(f)

# Pricing (Estimated for preview)
VEO_COST_PER_VIDEO = 0.50  # Placeholder cost for high-res video generation

# Output directory
OUTPUT_DIR = Path(os.path.dirname(__file__)) / "generated_videos"
OUTPUT_DIR.mkdir(exist_ok=True)


class VeoAgent:
    """
    Agent for generating videos using Veo 3.1.
    Handles the asynchronous polling process for video completion.
    """
    
    def __init__(self, 
                 model: str = "veo-3.1-generate-preview",
                 budget_limit: float = 10.0):
        self.model_name = model
        self.budget = BudgetManager(daily_limit=budget_limit)
        
        # Configure API
        api_key = CONFIG["gemini"]["api_key"]
        self.client = genai.Client(api_key=api_key)
        
        print(f"[OK] Veo Agent initialized")
        print(f"[OK] Model: {model}")
        print(f"[OK] Base Cost: ${VEO_COST_PER_VIDEO}/video")
    
    def generate(self,
                 prompt: str,
                 negative_prompt: Optional[str] = None,
                 generate_audio: bool = True,
                 save_to_disk: bool = True,
                 wait_for_completion: bool = True) -> Dict[str, Any]:
        """
        Generate a video from a text prompt.
        
        Args:
            prompt: Description of the video
            negative_prompt: What to avoid in the video
            generate_audio: Whether to include native audio
            save_to_disk: Save locally upon completion
            wait_for_completion: Sync/Async toggle
            
        Returns:
            Dict with 'operation_id', 'status', 'path' (if sync), 'cost'
        """
        # Budget check
        if not self.budget.is_within_budget():
            return {
                "error": "budget_exceeded",
                "status": "failed"
            }
        
        try:
            # Start generation operation
            operation = self.client.models.generate_videos(
                model=self.model_name,
                prompt=prompt,
                config=types.GenerateVideosConfig(
                    negative_prompt=negative_prompt,
                    generate_audio=generate_audio
                )
            )
            
            operation_id = operation.name
            print(f"[VEO] Operation started: {operation_id}")
            
            if not wait_for_completion:
                return {
                    "operation_id": operation_id,
                    "status": "pending",
                    "cost": VEO_COST_PER_VIDEO
                }
            
            # Polling for completion
            print("[VEO] Polling for completion (this may take several minutes)...")
            start_time = time.time()
            while not operation.done:
                time.sleep(10)
                operation = self.client.operations.get(operation_id)
                elapsed = int(time.time() - start_time)
                print(f"[VEO] Status: {operation.metadata.get('status', 'running')} ({elapsed}s elapsed)")
            
            if operation.error:
                return {
                    "error": str(operation.error),
                    "status": "failed"
                }
            
            # Extract video data
            video_data = operation.result.generated_videos[0].video.video_bytes
            
            # Log cost
            self.budget.log_cost(VEO_COST_PER_VIDEO, self.model_name)
            
            # Save to disk
            path = None
            if save_to_disk:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"veo_{timestamp}.mp4"
                filepath = OUTPUT_DIR / filename
                
                with open(filepath, "wb") as f:
                    f.write(video_data)
                
                path = str(filepath)
                print(f"[SAVED] {filepath}")
            
            return {
                "operation_id": operation_id,
                "status": "completed",
                "path": path,
                "cost": VEO_COST_PER_VIDEO,
                "prompt": prompt,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            return {
                "error": str(e),
                "status": "failed",
                "timestamp": datetime.now().isoformat()
            }
    
    def get_stats(self) -> Dict[str, Any]:
        """Get agent statistics."""
        return {
            "model": self.model_name,
            "budget_spent": self.budget.get_current_spend(),
            "budget_limit": self.budget.daily_limit,
            "within_budget": self.budget.is_within_budget(),
            "output_dir": str(OUTPUT_DIR)
        }


# CLI for testing
if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Veo Generator")
    parser.add_argument("--generate", "-g", type=str, help="Generate video from prompt")
    parser.add_argument("--no-audio", action="store_true", help="Disable audio generation")
    parser.add_argument("--test", action="store_true", help="Run self-test")
    args = parser.parse_args()
    
    agent = VeoAgent()
    
    if args.test:
        print("\n[TEST] Veo Generator Self-Test")
        print(f"[OK] Agent initialized: {agent.get_stats()}")
        print("\n[DONE] Self-test complete")
        
    elif args.generate:
        result = agent.generate(args.generate, generate_audio=not args.no_audio)
        if result.get("error"):
            print(f"[ERROR] {result['error']}")
        elif result.get("status") == "completed":
            print(f"[OK] Video generated: {result['path']}")
            print(f"[COST] ${result['cost']:.2f}")
    else:
        print("Usage: python veo_generator.py --generate 'your prompt'")
        print("       python veo_generator.py --test")