"""
Investigation Script 3: Subtitle/Transcript Extraction - Granularity Comparison

This script tests different approaches for subtitle extraction:
1. Movie-level (entire video summary)
2. Second-level (time-stamped subtitles per second/phrase)
3. Word-level (individual word timestamps)

Compares two methods:
- Method A: Upload video directly to Gemini
- Method B: Extract audio with ffmpeg, upload audio only

Tests which is faster, more accurate, and uses fewer tokens.
"""

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

# Load environment variables
load_dotenv()

# Initialize Rich console
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 AudioExtractor:
    """Extract audio from video using ffmpeg"""
    
    @staticmethod
    def extract_audio(video_path: str, output_path: str) -> Optional[str]:
        """
        Extract audio from video file using ffmpeg
        
        Args:
            video_path: Path to video file
            output_path: Path for output audio file
            
        Returns:
            Path to extracted audio file, or None if failed
        """
        try:
            console.print(f"[cyan]Extracting audio from video...[/cyan]")
            
            # Use ffmpeg to extract audio as MP3
            cmd = [
                'ffmpeg',
                '-i', video_path,
                '-vn',  # No video
                '-acodec', 'libmp3lame',  # MP3 codec
                '-q:a', '2',  # High quality
                '-y',  # Overwrite output
                output_path
            ]
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True
            )
            
            console.print(f"[green]✓[/green] Audio extracted: {Path(output_path).name}")
            return output_path
            
        except subprocess.CalledProcessError as e:
            console.print(f"[red]❌ FFmpeg error:[/red] {e.stderr}")
            return None
        except FileNotFoundError:
            console.print("[red]❌ FFmpeg not found. Please install ffmpeg.[/red]")
            return None


class SubtitleExtractor:
    """Extract subtitles at different granularity levels"""
    
    def __init__(self):
        self.model = 'gemini-3-pro-preview'
        console.print("[green]✓[/green] Initialized Subtitle Extractor")
    
    def upload_media(self, file_path: str, mime_type: str) -> Any:
        """Upload video or audio file to Gemini"""
        console.print(f"[cyan]Uploading {mime_type.split('/')[0]}:[/cyan] {Path(file_path).name}")
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task = progress.add_task("Uploading to Gemini API...", total=None)
            
            with open(file_path, 'rb') as f:
                upload_config = types.UploadFileConfig(mime_type=mime_type)
                media_file = client.files.upload(file=f, config=upload_config)
            
            # Wait for processing
            while media_file.state == "PROCESSING":
                time.sleep(2)
                media_file = client.files.get(name=media_file.name)
            
            progress.update(task, completed=True)
        
        if media_file.state == "FAILED":
            raise ValueError(f"Media processing failed: {media_file.state}")
        
        console.print(f"[green]✓[/green] Media uploaded successfully")
        return media_file
    
    def extract_movie_level(self, media_file: Any) -> Dict[str, Any]:
        """
        Extract movie-level subtitle (entire video summary)
        
        Returns:
            Dictionary with summary and full transcript
        """
        prompt = """
このビデオ/音声の内容を全体として書き起こしてください。

以下の情報を抽出してください：
1. **全体の要約**: ビデオ全体の内容を2-3文で要約
2. **完全な書き起こし**: すべての発言を時系列で書き起こし
3. **話者**: 誰が話しているか識別

**出力形式**: 必ずJSON形式で返してください：
{
  "summary": "全体要約",
  "full_transcript": "完全な書き起こしテキスト",
  "speakers": ["話者1", "話者2"],
  "duration": "MM:SS"
}
"""
        return self._generate_content(media_file, prompt, "Movie-level")
    
    def extract_second_level(self, media_file: Any) -> Dict[str, Any]:
        """
        Extract second-level subtitles (time-stamped by phrase/sentence)
        
        Returns:
            Dictionary with time-stamped subtitle segments
        """
        prompt = """
この音声/ビデオの内容を時間軸に沿って正確に書き起こし、字幕形式で出力してください。

**要件**:
- 各発言を開始時刻と終了時刻付きでセグメント化
- 話者を識別
- 時刻はMM:SS.mmm形式（秒単位で正確に）

**出力形式**: 必ずJSON形式で返してください：
{
  "subtitles": [
    {
      "index": 1,
      "start_time": "00:00.000",
      "end_time": "00:03.500",
      "speaker": "話者名",
      "text": "発言内容"
    }
  ],
  "total_segments": 0
}
"""
        return self._generate_content(media_file, prompt, "Second-level")
    
    def extract_word_level(self, media_file: Any) -> Dict[str, Any]:
        """
        Extract word-level subtitles (timestamp for each word)
        
        Returns:
            Dictionary with word-by-word timestamps
        """
        prompt = """
この音声/ビデオの内容を単語レベルで書き起こし、各単語のタイムスタンプを付けてください。

**要件**:
- 各単語に開始時刻を付与
- 話者を識別
- 時刻はMM:SS.mmm形式

**出力形式**: 必ずJSON形式で返してください：
{
  "word_timestamps": [
    {
      "word": "単語",
      "start_time": "00:00.000",
      "speaker": "話者名"
    }
  ],
  "total_words": 0
}
"""
        return self._generate_content(media_file, prompt, "Word-level")
    
    def _generate_content(self, media_file: Any, prompt: str, level: str) -> Dict[str, Any]:
        """Generate content with token tracking"""
        console.print(f"[cyan]Extracting {level} subtitles...[/cyan]")
        
        start_time = time.time()
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task = progress.add_task(f"Generating {level} content...", total=None)
            
            response = client.models.generate_content(
                model=self.model,
                contents=[media_file, prompt]
            )
            
            progress.update(task, completed=True)
        
        elapsed_time = time.time() - start_time
        
        # Extract token usage
        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
        }
        
        console.print(f"[cyan]📊 Tokens:[/cyan] {token_info['total_tokens']:,} | [cyan]⏱ Time:[/cyan] {elapsed_time:.2f}s")
        
        # Parse response
        try:
            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"] = {
                "level": level,
                "token_usage": token_info,
                "processing_time_seconds": elapsed_time,
                "model": self.model
            }
            
            console.print(f"[green]✓[/green] {level} extraction complete")
            return result
            
        except json.JSONDecodeError as e:
            console.print(f"[yellow]⚠[/yellow] JSON parsing failed, returning raw response")
            return {
                "raw_response": response.text,
                "error": str(e),
                "_metadata": {
                    "level": level,
                    "token_usage": token_info,
                    "processing_time_seconds": elapsed_time
                }
            }


def compare_methods(video_path: str, output_dir: Path, granularity: Literal["movie", "second", "word"]):
    """
    Compare video vs audio extraction for a specific granularity level
    
    Args:
        video_path: Path to video file
        output_dir: Output directory for results
        granularity: Level of granularity to test
    """
    extractor = SubtitleExtractor()
    audio_extractor = AudioExtractor()
    results = {}
    
    console.print(f"\n[bold cyan]Testing {granularity.upper()}-level extraction[/bold cyan]")
    console.print("="*60)
    
    # Method A: Direct video upload
    console.print("\n[bold]Method A: Video Upload[/bold]")
    try:
        video_file = extractor.upload_media(video_path, 'video/mp4')
        
        if granularity == "movie":
            result_a = extractor.extract_movie_level(video_file)
        elif granularity == "second":
            result_a = extractor.extract_second_level(video_file)
        else:  # word
            result_a = extractor.extract_word_level(video_file)
        
        results["method_a_video"] = result_a
        
    except Exception as e:
        console.print(f"[red]❌ Method A failed:[/red] {e}")
        results["method_a_video"] = {"error": str(e)}
    
    # Method B: Audio extraction + upload
    console.print("\n[bold]Method B: Audio Extraction + Upload[/bold]")
    try:
        # Extract audio
        audio_path = str(output_dir / f"{Path(video_path).stem}_audio.mp3")
        extracted_audio = audio_extractor.extract_audio(video_path, audio_path)
        
        if extracted_audio:
            audio_file = extractor.upload_media(extracted_audio, 'audio/mpeg')
            
            if granularity == "movie":
                result_b = extractor.extract_movie_level(audio_file)
            elif granularity == "second":
                result_b = extractor.extract_second_level(audio_file)
            else:  # word
                result_b = extractor.extract_word_level(audio_file)
            
            results["method_b_audio"] = result_b
        else:
            results["method_b_audio"] = {"error": "Audio extraction failed"}
            
    except Exception as e:
        console.print(f"[red]❌ Method B failed:[/red] {e}")
        results["method_b_audio"] = {"error": str(e)}
    
    return results


def main():
    """Main investigation workflow"""
    console.print(Panel.fit(
        "[bold cyan]Subtitle Extraction Investigation[/bold cyan]\n"
        "Comparing Video vs Audio | Movie/Second/Word Level Granularity",
        border_style="cyan"
    ))
    
    # Setup paths
    input_dir = Path("assets/0_input")
    output_dir = Path("assets/1_processed/subtitle_investigation")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Get first video for testing
    video_files = list(input_dir.glob("*.mp4"))
    
    if not video_files:
        console.print("[red]❌ No video files found[/red]")
        return
    
    test_video = video_files[0]  # Use first video for comparison
    console.print(f"\n[cyan]Test Video:[/cyan] {test_video.name}")
    
    # Test all granularity levels
    all_results = {}
    
    for granularity in ["movie", "second", "word"]:
        results = compare_methods(str(test_video), output_dir, granularity)
        all_results[granularity] = results
        
        # Save individual results
        result_file = output_dir / f"{test_video.stem}_{granularity}_comparison.json"
        with open(result_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        console.print(f"[green]✓[/green] Results saved: {result_file.name}")
    
    # Generate comparison summary
    console.print("\n" + "="*60)
    console.print("[bold green]Comparison Summary[/bold green]\n")
    
    table = Table(title="Performance Comparison", show_header=True)
    table.add_column("Granularity", style="cyan")
    table.add_column("Method", style="yellow")
    table.add_column("Tokens", justify="right", style="magenta")
    table.add_column("Time (s)", justify="right", style="blue")
    
    for granularity, results in all_results.items():
        for method_name, data in results.items():
            if "_metadata" in data:
                meta = data["_metadata"]
                table.add_row(
                    granularity.upper(),
                    method_name.replace("_", " ").title(),
                    f"{meta['token_usage']['total_tokens']:,}",
                    f"{meta['processing_time_seconds']:.2f}"
                )
    
    console.print(table)
    
    # Save complete comparison
    summary_file = output_dir / "complete_comparison.json"
    with open(summary_file, 'w', encoding='utf-8') as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    
    console.print(f"\n[green]✓[/green] Complete comparison saved: {summary_file}")


if __name__ == "__main__":
    main()