"""
Investigation Script 1: Video Upload and Timeline/Dialogue Extraction using Gemini

This script demonstrates:
1. Uploading video files to Google Gemini API
2. Extracting dialogue/transcription from video
3. Identifying timestamps for key moments
4. Structuring the output for video editing pipeline

Target: Process interview videos (№１_基本情報.mp4, №２_ライフイベント.mp4)
"""

import os
import json
import time
from pathlib import Path
from typing import Dict, Any, List
from dotenv import load_dotenv
from google import genai
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.panel import Panel
from rich.json import JSON

# Load environment variables
load_dotenv()

# Initialize Rich console for beautiful output
console = Console()

# Configure Gemini API
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    console.print("[red]❌ GEMINI_API_KEY not found in .env file[/red]")
    exit(1)

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


class VideoExtractor:
    """Extract timeline and dialogue from video using Gemini"""
    
    def __init__(self):
        # Use Gemini 3 Pro for video understanding
        self.model = 'gemini-3-pro-preview'  # Latest model supporting video
        console.print("[green]✓[/green] Initialized Gemini 3 Pro model")
    
    def upload_video(self, video_path: str) -> Any:
        """
        Upload video file to Gemini API
        
        Args:
            video_path: Path to video file
            
        Returns:
            Uploaded file object
        """
        console.print(f"[cyan]Uploading video:[/cyan] {Path(video_path).name}")
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task = progress.add_task("Uploading to Gemini API...", total=None)
            
            # Upload the video file with config
            from google.genai import types
            with open(video_path, 'rb') as f:
                upload_config = types.UploadFileConfig(
                    mime_type='video/mp4'
                )
                video_file = client.files.upload(
                    file=f,
                    config=upload_config
                )
            
            # Wait for processing to complete
            while video_file.state == "PROCESSING":
                time.sleep(2)
                video_file = client.files.get(name=video_file.name)
            
            progress.update(task, completed=True)
        
        if video_file.state == "FAILED":
            raise ValueError(f"Video processing failed: {video_file.state}")
        
        console.print(f"[green]✓[/green] Video uploaded successfully: {video_file.name}")
        return video_file
    
    def extract_timeline_and_dialogue(self, video_file: Any, video_name: str) -> Dict[str, Any]:
        """
        Extract structured timeline and dialogue from video
        
        Args:
            video_file: Uploaded video file object
            video_name: Name of the video for context
            
        Returns:
            Dictionary containing timeline segments with timestamps and transcription
        """
        console.print("\n[cyan]Analyzing video content...[/cyan]")
        
        # Prompt for structured extraction (in Japanese for better understanding)
        prompt = f"""
このビデオは保険コンサルタントと顧客のインタビュー映像です。
ビデオのタイトル: {video_name}

以下の情報を抽出してください：

1. **全体の要約**: ビデオの内容を2-3文で要約
2. **話者情報**: 誰が話しているか（例：コンサルタント、顧客）
3. **タイムライン**: 重要な場面ごとに以下の情報を含むセグメントに分割
   - 開始時刻（MM:SS形式）
   - 終了時刻（MM:SS形式）  
   - 発言内容（日本語の書き起こし）
   - トピック/キーワード（例：「基本情報」「家族構成」「将来の不安」）
   - 感情/トーン（例：「心配」「前向き」「真剣」）
   - 映像使用の推奨度（1-5、5が最も重要）

4. **キーワード**: B-Roll素材検索用のキーワードを5-10個抽出（日本語）

**出力形式**: 必ずJSON形式で返してください。以下の構造に従ってください：

{{
  "summary": "ビデオ全体の要約",
  "speakers": ["話者1", "話者2"],
  "duration": "MM:SS",
  "segments": [
    {{
      "start_time": "00:00",
      "end_time": "00:30",
      "speaker": "コンサルタント",
      "transcript": "発言内容の書き起こし",
      "topic": "トピック名",
      "emotion": "感情/トーン",
      "importance": 4,
      "keywords": ["キーワード1", "キーワード2"]
    }}
  ],
  "broll_keywords": ["素材検索用キーワード1", "キーワード2", "..."]
}}
"""
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task = progress.add_task("Extracting content with Gemini...", total=None)
            
            # Generate content from video - use file object directly with prompt
            response = client.models.generate_content(
                model=self.model,
                contents=[
                    video_file,  # Pass the file object directly
                    prompt
                ]
            )
            
            progress.update(task, completed=True)
        
        # Parse the response
        try:
            # Extract token usage information
            token_info = {
                "input_tokens": response.usage_metadata.prompt_token_count if hasattr(response, 'usage_metadata') else 0,
                "output_tokens": response.usage_metadata.candidates_token_count if hasattr(response, 'usage_metadata') else 0,
                "total_tokens": response.usage_metadata.total_token_count if hasattr(response, 'usage_metadata') else 0
            }
            
            # Display token usage
            console.print(f"[cyan]📊 Token Usage:[/cyan] Input: {token_info['input_tokens']:,} | Output: {token_info['output_tokens']:,} | Total: {token_info['total_tokens']:,}")
            
            # Extract JSON from response (handle markdown code blocks if present)
            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()
            
            result = json.loads(response_text)
            result["_metadata"] = {
                "token_usage": token_info,
                "model": self.model,
                "video_name": video_name
            }
            console.print("[green]✓[/green] Successfully extracted timeline and dialogue")
            return result
        except json.JSONDecodeError as e:
            console.print(f"[yellow]⚠[/yellow] Failed to parse JSON, returning raw response")
            return {
                "raw_response": response.text,
                "error": str(e),
                "_metadata": {
                    "token_usage": token_info if 'token_info' in locals() else {},
                    "model": self.model
                }
            }
    
    def save_results(self, results: Dict[str, Any], output_path: str):
        """Save extraction results to JSON file"""
        output_file = Path(output_path)
        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(f"[green]✓[/green] Results saved to: {output_file}")


def main():
    """Main investigation workflow"""
    console.print(Panel.fit(
        "[bold cyan]Video Extraction Investigation Script[/bold cyan]\n"
        "Testing Gemini 1.5 Pro for video timeline and dialogue extraction",
        border_style="cyan"
    ))
    
    # Initialize extractor
    extractor = VideoExtractor()
    
    # Define input and output paths
    input_dir = Path("assets/0_input")
    output_dir = Path("assets/1_processed/video_extractions")
    
    # Get all MP4 files
    video_files = list(input_dir.glob("*.mp4"))
    
    if not video_files:
        console.print("[red]❌ No video files found in assets/0_input/[/red]")
        return
    
    console.print(f"\n[cyan]Found {len(video_files)} video(s) to process:[/cyan]")
    for video in video_files:
        console.print(f"  • {video.name}")
    
    # Process each video
    all_results = {}
    total_tokens_used = 0
    
    for video_path in video_files:
        console.print(f"\n{'='*60}")
        console.print(f"[bold]Processing: {video_path.name}[/bold]")
        console.print(f"{'='*60}\n")
        
        try:
            # Upload video
            video_file = extractor.upload_video(str(video_path))
            
            # Extract timeline and dialogue
            results = extractor.extract_timeline_and_dialogue(video_file, video_path.name)
            
            # Display results
            console.print("\n[bold green]Extraction Results:[/bold green]")
            console.print(JSON(json.dumps(results, ensure_ascii=False)))
            
            # Save results
            output_file = output_dir / f"{video_path.stem}_extraction.json"
            extractor.save_results(results, str(output_file))
            
            all_results[video_path.name] = results
            
            # Track total tokens
            if "_metadata" in results and "token_usage" in results["_metadata"]:
                total_tokens_used += results["_metadata"]["token_usage"].get("total_tokens", 0)
            
        except Exception as e:
            console.print(f"[red]❌ Error processing {video_path.name}:[/red] {e}")
            continue
    
    # Save combined results
    if all_results:
        combined_output = output_dir / "all_extractions.json"
        extractor.save_results(all_results, str(combined_output))
        
        console.print("\n" + "="*60)
        console.print("[bold green]✓ Investigation Complete![/bold green]")
        console.print(f"Processed {len(all_results)} video(s)")
        console.print(f"Results saved to: {output_dir}")
        console.print(f"\n[cyan]💰 Total Tokens Consumed: {total_tokens_used:,}[/cyan]")


if __name__ == "__main__":
    main()