"""
Investigation Script 2: Keyword Translation and Pexels Media Fetcher

This script demonstrates:
1. Translating Japanese keywords to English using Gemini
2. Searching Pexels API for relevant B-Roll footage
3. Downloading videos to local assets folder
4. Caching translations to avoid repeated API calls

Target: Find stock footage for insurance/financial video production
"""

import os
import json
import time
from pathlib import Path
from typing import List, Dict, Any, Optional
from dotenv import load_dotenv
from google import genai
import requests
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, DownloadColumn, TransferSpeedColumn
from rich.panel import Panel
from rich.table import Table

# Load environment variables
load_dotenv()

# Initialize Rich console
console = Console()

# Configure APIs
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
PEXELS_API_KEY = os.getenv("PEXELS_API_KEY")

if not GEMINI_API_KEY or not PEXELS_API_KEY:
    console.print("[red]❌ API keys not found in .env file[/red]")
    console.print("Required: GEMINI_API_KEY, PEXELS_API_KEY")
    exit(1)

# Initialize Gemini client
client = genai.Client(api_key=GEMINI_API_KEY)


class KeywordTranslator:
    """Translate Japanese keywords to English using Gemini"""
    
    def __init__(self, cache_file: str = "assets/1_processed/translation_cache.json"):
        self.model = 'gemini-2.0-flash-exp'  # Latest model for translation
        self.cache_file = Path(cache_file)
        self.cache = self._load_cache()
        console.print("[green]✓[/green] Initialized Gemini translator with cache")
    
    def _load_cache(self) -> Dict[str, str]:
        """Load translation cache from file"""
        if self.cache_file.exists():
            with open(self.cache_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        return {}
    
    def _save_cache(self):
        """Save translation cache to file"""
        self.cache_file.parent.mkdir(parents=True, exist_ok=True)
        with open(self.cache_file, 'w', encoding='utf-8') as f:
            json.dump(self.cache, f, ensure_ascii=False, indent=2)
    
    def translate(self, japanese_keywords: List[str]) -> Dict[str, str]:
        """
        Translate Japanese keywords to English
        
        Args:
            japanese_keywords: List of Japanese keywords
            
        Returns:
            Dictionary mapping Japanese to English keywords
        """
        translations = {}
        keywords_to_translate = []
        
        # Check cache first
        for keyword in japanese_keywords:
            if keyword in self.cache:
                translations[keyword] = self.cache[keyword]
            else:
                keywords_to_translate.append(keyword)
        
        if keywords_to_translate:
            console.print(f"[cyan]Translating {len(keywords_to_translate)} new keywords...[/cyan]")
            
            # Batch translation for efficiency
            prompt = f"""
Translate the following Japanese keywords to English for searching stock footage/videos.
Provide natural, searchable English terms that would work well for video search.

Japanese keywords: {', '.join(keywords_to_translate)}

Return ONLY a JSON object mapping Japanese to English, like this:
{{
  "日本語キーワード1": "english keyword 1",
  "日本語キーワード2": "english keyword 2"
}}
"""
            
            response = client.models.generate_content(
                model=self.model,
                contents=prompt
            )
            
            # Display token usage
            if hasattr(response, 'usage_metadata'):
                token_info = {
                    "input_tokens": response.usage_metadata.prompt_token_count,
                    "output_tokens": response.usage_metadata.candidates_token_count,
                    "total_tokens": response.usage_metadata.total_token_count
                }
                console.print(f"[cyan]📊 Translation Tokens:[/cyan] {token_info['total_tokens']:,} (in: {token_info['input_tokens']:,}, out: {token_info['output_tokens']:,})")
            
            try:
                # Parse response
                response_text = response.text
                if "```json" in response_text:
                    json_start = response_text.find("```json") + 7
                    json_end = response_text.find("```", json_start)
                    response_text = response_text[json_start:json_end].strip()
                elif "```" in response_text:
                    json_start = response_text.find("```") + 3
                    json_end = response_text.find("```", json_start)
                    response_text = response_text[json_start:json_end].strip()
                
                new_translations = json.loads(response_text)
                
                # Update cache and results
                self.cache.update(new_translations)
                translations.update(new_translations)
                self._save_cache()
                
                console.print(f"[green]✓[/green] Translated {len(new_translations)} keywords")
                
            except json.JSONDecodeError as e:
                console.print(f"[yellow]⚠[/yellow] Translation parsing failed: {e}")
                # Fallback: use romanized version
                for keyword in keywords_to_translate:
                    translations[keyword] = keyword
        
        return translations


class PexelsFetcher:
    """Fetch B-Roll videos from Pexels API"""
    
    def __init__(self, api_key: str, output_dir: str = "assets/1_processed/broll"):
        self.api_key = api_key
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.base_url = "https://api.pexels.com/videos/search"
        console.print("[green]✓[/green] Initialized Pexels API client")
    
    def search_videos(self, query: str, max_results: int = 3, min_duration: int = 5) -> List[Dict[str, Any]]:
        """
        Search for videos on Pexels
        
        Args:
            query: Search query (in English)
            max_results: Maximum number of videos to return
            min_duration: Minimum video duration in seconds
            
        Returns:
            List of video metadata dictionaries
        """
        headers = {
            "Authorization": self.api_key
        }
        
        params = {
            "query": query,
            "per_page": max_results,
            "orientation": "landscape",  # Better for editing
            "size": "medium"  # Balance quality and file size
        }
        
        try:
            response = requests.get(self.base_url, headers=headers, params=params)
            response.raise_for_status()
            
            data = response.json()
            videos = []
            
            for video in data.get("videos", []):
                # Find best quality HD file
                video_files = video.get("video_files", [])
                hd_file = None
                
                for vf in video_files:
                    if vf.get("quality") == "hd" and vf.get("width") == 1920:
                        hd_file = vf
                        break
                
                # Fallback to first available file
                if not hd_file and video_files:
                    hd_file = video_files[0]
                
                if hd_file and video.get("duration", 0) >= min_duration:
                    videos.append({
                        "id": video["id"],
                        "url": video["url"],
                        "duration": video["duration"],
                        "width": video["width"],
                        "height": video["height"],
                        "download_url": hd_file["link"],
                        "file_size": hd_file.get("size", 0),
                        "photographer": video["user"]["name"]
                    })
            
            return videos[:max_results]
            
        except requests.exceptions.RequestException as e:
            console.print(f"[red]❌ Pexels API error:[/red] {e}")
            return []
    
    def download_video(self, video: Dict[str, Any], keyword: str) -> Optional[str]:
        """
        Download a video file
        
        Args:
            video: Video metadata dictionary
            keyword: Keyword used for search (for filename)
            
        Returns:
            Path to downloaded file, or None if failed
        """
        # Create safe filename
        safe_keyword = "".join(c for c in keyword if c.isalnum() or c in (' ', '-', '_')).strip()
        safe_keyword = safe_keyword.replace(' ', '_')
        filename = f"{safe_keyword}_{video['id']}.mp4"
        filepath = self.output_dir / filename
        
        # Skip if already downloaded
        if filepath.exists():
            console.print(f"[yellow]⊘[/yellow] Already exists: {filename}")
            return str(filepath)
        
        try:
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                DownloadColumn(),
                TransferSpeedColumn(),
                console=console
            ) as progress:
                task = progress.add_task(f"Downloading {filename[:40]}...", total=video.get('file_size'))
                
                response = requests.get(video['download_url'], stream=True)
                response.raise_for_status()
                
                with open(filepath, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        f.write(chunk)
                        progress.update(task, advance=len(chunk))
            
            console.print(f"[green]✓[/green] Downloaded: {filename}")
            return str(filepath)
            
        except Exception as e:
            console.print(f"[red]❌ Download failed:[/red] {e}")
            if filepath.exists():
                filepath.unlink()  # Remove partial file
            return None


def main():
    """Main investigation workflow"""
    console.print(Panel.fit(
        "[bold cyan]Pexels Media Fetch Investigation Script[/bold cyan]\n"
        "Testing keyword translation and B-Roll video download",
        border_style="cyan"
    ))
    
    # Initialize components
    translator = KeywordTranslator()
    fetcher = PexelsFetcher(PEXELS_API_KEY)
    
    # Test keywords (typical for insurance/life planning videos)
    test_keywords_ja = [
        "家族",
        "老後",
        "貯蓄",
        "住宅",
        "教育",
        "保険",
        "将来",
        "不安",
        "安心",
        "計画"
    ]
    
    console.print("\n[bold]Test Keywords (Japanese):[/bold]")
    console.print(", ".join(test_keywords_ja))
    
    # Translate keywords
    console.print("\n" + "="*60)
    translations = translator.translate(test_keywords_ja)
    
    # Display translations
    table = Table(title="Translation Results", show_header=True, header_style="bold cyan")
    table.add_column("Japanese", style="yellow")
    table.add_column("English", style="green")
    
    for ja, en in translations.items():
        table.add_row(ja, en)
    
    console.print(table)
    
    # Search and download videos
    console.print("\n" + "="*60)
    console.print("[bold]Searching Pexels for B-Roll footage...[/bold]\n")
    
    results = {}
    
    for ja_keyword, en_keyword in list(translations.items())[:5]:  # Limit to 5 for testing
        console.print(f"\n[cyan]Searching:[/cyan] '{en_keyword}' (from '{ja_keyword}')")
        
        videos = fetcher.search_videos(en_keyword, max_results=2)
        
        if videos:
            console.print(f"  Found {len(videos)} video(s)")
            
            downloaded = []
            for video in videos:
                filepath = fetcher.download_video(video, en_keyword)
                if filepath:
                    downloaded.append({
                        "filepath": filepath,
                        "duration": video["duration"],
                        "photographer": video["photographer"]
                    })
            
            results[ja_keyword] = {
                "english": en_keyword,
                "videos": downloaded
            }
        else:
            console.print(f"  [yellow]No videos found[/yellow]")
    
    # Save results summary
    output_file = Path("assets/1_processed/broll_fetch_results.json")
    output_file.parent.mkdir(parents=True, exist_ok=True)
    
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    
    console.print("\n" + "="*60)
    console.print(f"[bold green]✓ Investigation Complete![/bold green]")
    console.print(f"Downloaded videos saved to: {fetcher.output_dir}")
    console.print(f"Results summary saved to: {output_file}")


if __name__ == "__main__":
    main()