#!/usr/bin/env python3
"""
YouTube PO Token Handler
Extracts and manages PO tokens for yt-dlp YouTube access
"""

import os
import json
import time
import subprocess
import logging
from pathlib import Path
from typing import Optional, Dict, Any
import sqlite3
import tempfile
from dotenv import load_dotenv

# Load environment variables
load_dotenv()


class YouTubePOTokenHandler:
    """Handles PO token extraction and management for YouTube."""
    
    def __init__(self, logger: Optional[logging.Logger] = None):
        self.logger = logger or logging.getLogger(__name__)
        self.token_cache = {}
        self.token_expiry = {}
        
        # Firefox profile detection
        self.firefox_profile_path = self._find_firefox_profile()
        
        # Token types we can extract
        self.token_types = ['mweb.gvs', 'mweb.subs', 'web.gvs', 'web.subs']
        
    def _find_firefox_profile(self) -> Optional[Path]:
        """Find the active Firefox profile directory."""
        try:
            # Common Firefox profile locations
            profile_paths = [
                Path.home() / ".mozilla/firefox",
                Path.home() / "snap/firefox/common/.mozilla/firefox",  # Snap in home
                Path("/snap/firefox/common/.mozilla/firefox"),  # Snap system
                Path("/var/lib/snapd/desktop/firefox/.mozilla/firefox")  # Snap alt
            ]
            
            for base_path in profile_paths:
                if not base_path.exists():
                    continue
                
                self.logger.debug(f"Checking Firefox path: {base_path}")
                
                # Look for profiles.ini
                profiles_ini = base_path / "profiles.ini"
                if profiles_ini.exists():
                    # Parse profiles.ini to find default profile
                    content = profiles_ini.read_text()
                    for line in content.split('\n'):
                        if 'Path=' in line and 'default' in line.lower():
                            profile_name = line.split('=')[1].strip()
                            profile_path = base_path / profile_name
                            if profile_path.exists():
                                self.logger.info(f"Found Firefox profile via profiles.ini: {profile_path}")
                                return profile_path
                    
                    # Fallback: find any .default profile
                    for item in base_path.iterdir():
                        if item.is_dir() and 'default' in item.name:
                            self.logger.info(f"Found Firefox profile via .default search: {item}")
                            return item
                else:
                    # No profiles.ini, look for .default directories directly
                    for item in base_path.iterdir():
                        if item.is_dir() and 'default' in item.name:
                            self.logger.info(f"Found Firefox profile directly: {item}")
                            return item
            
            self.logger.warning("Firefox profile not found in any standard locations")
            return None
            
        except Exception as e:
            self.logger.error(f"Error finding Firefox profile: {e}")
            return None
    
    def _extract_token_from_network_log(self) -> Optional[Dict[str, str]]:
        """Extract PO token from Firefox network activity (requires manual browser session)."""
        
        # This is a placeholder for the manual token extraction process
        # In practice, users would need to:
        # 1. Open YouTube in Firefox
        # 2. Open Developer Tools -> Network tab
        # 3. Filter by 'player' or 'v1/player'
        # 4. Find requests with PO tokens in payload
        # 5. Copy the token values
        
        self.logger.info("PO Token extraction requires manual browser session:")
        self.logger.info("1. Open YouTube in Firefox (signed in as benreed1987@gmail.com)")
        self.logger.info("2. Open Developer Tools (F12) -> Network tab")
        self.logger.info("3. Filter by 'player' or search for 'v1/player' requests")
        self.logger.info("4. Look for 'serviceIntegrityDimensions.poToken' in request payload")
        
        return None
    
    def _check_token_cache(self, token_type: str) -> Optional[str]:
        """Check if we have a valid cached token."""
        if token_type not in self.token_cache:
            return None
        
        # Check if token has expired (tokens typically last 1-6 hours)
        if token_type in self.token_expiry:
            if time.time() > self.token_expiry[token_type]:
                self.logger.debug(f"Token {token_type} has expired")
                del self.token_cache[token_type]
                del self.token_expiry[token_type]
                return None
        
        return self.token_cache[token_type]
    
    def _save_token_to_cache(self, token_type: str, token: str, ttl_hours: int = 2):
        """Save token to memory cache with TTL."""
        self.token_cache[token_type] = token
        self.token_expiry[token_type] = time.time() + (ttl_hours * 3600)
        self.logger.debug(f"Cached token {token_type} for {ttl_hours} hours")
    
    def get_po_token(self, token_type: str = 'mweb.gvs') -> Optional[str]:
        """Get a PO token for the specified type."""
        
        # Check cache first
        cached_token = self._check_token_cache(token_type)
        if cached_token:
            self.logger.debug(f"Using cached token for {token_type}")
            return cached_token
        
        # Try environment variable first (manual override)
        env_var = f"YOUTUBE_PO_TOKEN_{token_type.replace('.', '_').upper()}"
        env_token = os.getenv(env_var)
        if env_token:
            self.logger.info(f"Using PO token from environment: {env_var}")
            self._save_token_to_cache(token_type, env_token)
            return env_token
        
        # Try to extract from browser (requires manual process)
        self.logger.warning(f"No PO token found for {token_type}")
        self.logger.info("To obtain PO tokens manually:")
        self.logger.info("1. Visit https://music.youtube.com in Firefox")
        self.logger.info("2. Open Developer Tools (F12)")
        self.logger.info("3. Go to Network tab, filter by 'player'")
        self.logger.info("4. Play any video and look for v1/player requests")
        self.logger.info("5. Find 'serviceIntegrityDimensions.poToken' in request payload")
        self.logger.info(f"6. Set environment variable: export {env_var}='your_token_here'")
        
        return None
    
    def test_token_validity(self, token: str, token_type: str = 'mweb.gvs') -> bool:
        """Test if a PO token is valid by attempting a simple yt-dlp request."""
        try:
            # Create a simple test video URL
            test_url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ"  # Rick Roll (reliable test)
            
            # Build yt-dlp command with PO token
            cmd = [
                "yt-dlp",
                "--cookies-from-browser", "firefox",
                "--extractor-args", f"youtube:po_token={token_type}+{token}",
                "--simulate",  # Don't download, just test access
                "--quiet",
                test_url
            ]
            
            # Run test with timeout
            result = subprocess.run(
                cmd, 
                capture_output=True, 
                text=True, 
                timeout=30,
                cwd=Path.cwd()
            )
            
            if result.returncode == 0:
                self.logger.info(f"PO token {token_type} is valid")
                return True
            else:
                self.logger.warning(f"PO token {token_type} validation failed: {result.stderr}")
                return False
                
        except subprocess.TimeoutExpired:
            self.logger.warning("PO token validation timed out")
            return False
        except Exception as e:
            self.logger.error(f"Error testing PO token: {e}")
            return False
    
    def get_ytdlp_args(self, include_po_token: bool = True) -> Dict[str, Any]:
        """Get yt-dlp configuration with PO token support."""
        
        base_args = {
            'cookiesfrombrowser': ('firefox',),  # Use Firefox cookies
            'quiet': False,
            'no_warnings': False,
            'extract_flat': False,
        }
        
        if include_po_token:
            # Try to get a valid PO token
            token = self.get_po_token('mweb.gvs')  # Primary token type
            
            if token:
                # Add PO token to extractor args - correct format: "CLIENT.CONTEXT+TOKEN"
                extractor_args = {
                    'youtube': {
                        'po_token': f'mweb.gvs+{token}',
                        'player_client': 'default,mweb'
                    }
                }
                base_args['extractor_args'] = extractor_args
                self.logger.info("PO token configured for yt-dlp")
            else:
                self.logger.warning("No PO token available - transcript extraction may fail")
                # Still use cookies for best-effort access
                extractor_args = {
                    'youtube': {
                        'player_client': 'default,mweb'
                    }
                }
                base_args['extractor_args'] = extractor_args
        
        return base_args
    
    def extract_subtitle_with_token(self, video_id: str) -> Optional[str]:
        """Extract subtitle using yt-dlp with PO token."""
        try:
            video_url = f"https://www.youtube.com/watch?v={video_id}"
            
            # Get yt-dlp configuration with PO token
            ytdl_opts = self.get_ytdlp_args(include_po_token=True)
            
            # Add subtitle-specific options
            ytdl_opts.update({
                'writesubtitles': True,
                'writeautomaticsub': True,
                'subtitleslangs': ['en'],
                'skip_download': True,
                'subtitlesformat': 'vtt/srt/json3',
            })
            
            import yt_dlp
            
            with yt_dlp.YoutubeDL(ytdl_opts) as ydl:
                # Extract video info including subtitles
                info = ydl.extract_info(video_url, download=False)
                
                if not info:
                    return None
                
                # Check for subtitles
                subtitles = info.get('subtitles', {})
                auto_captions = info.get('automatic_captions', {})
                
                # Prefer manual subtitles over auto-generated
                captions_data = subtitles.get('en') or auto_captions.get('en')
                
                if not captions_data:
                    return None
                
                # Find best subtitle format
                best_subtitle = None
                for subtitle in captions_data:
                    if subtitle.get('ext') in ['vtt', 'srt', 'json3']:
                        best_subtitle = subtitle
                        break
                
                if not best_subtitle:
                    best_subtitle = captions_data[0]
                
                # Fetch subtitle content
                subtitle_url = best_subtitle.get('url')
                if subtitle_url:
                    import urllib.request
                    with urllib.request.urlopen(subtitle_url) as response:
                        content = response.read().decode('utf-8')
                    
                    # Simple VTT parsing (extract text only)
                    if best_subtitle.get('ext') == 'vtt':
                        lines = content.split('\n')
                        text_parts = []
                        for line in lines:
                            line = line.strip()
                            if (line and 
                                not line.startswith('WEBVTT') and
                                not line.startswith('NOTE') and
                                '-->' not in line and
                                not line.isdigit()):
                                # Remove HTML tags
                                import re
                                clean_line = re.sub(r'<[^>]+>', '', line)
                                if clean_line:
                                    text_parts.append(clean_line)
                        
                        return ' '.join(text_parts) if text_parts else None
                    
                    return content  # Return raw content for other formats
                
        except Exception as e:
            self.logger.error(f"Error extracting subtitle with PO token for {video_id}: {e}")
        
        return None
    
    def get_status(self) -> Dict[str, Any]:
        """Get status of PO token handler."""
        return {
            'firefox_profile_found': self.firefox_profile_path is not None,
            'firefox_profile_path': str(self.firefox_profile_path) if self.firefox_profile_path else None,
            'cached_tokens': list(self.token_cache.keys()),
            'token_types_supported': self.token_types,
            'env_tokens_available': [
                env_var for env_var in [
                    'YOUTUBE_PO_TOKEN_MWEB_GVS',
                    'YOUTUBE_PO_TOKEN_MWEB_SUBS', 
                    'YOUTUBE_PO_TOKEN_WEB_GVS',
                    'YOUTUBE_PO_TOKEN_WEB_SUBS'
                ] if os.getenv(env_var)
            ]
        }