"""
Enhanced Video Accent Analyzer
Supports YouTube, Loom, direct MP4 links, and local video files with improved error handling and features.

"""

import os
import sys
import tempfile
import subprocess
import requests
import json
import warnings
import time
from pathlib import Path
from urllib.parse import urlparse
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
try:
    from IPython.display import display, HTML, Audio
    IPYTHON_AVAILABLE = True
except ImportError:
    IPYTHON_AVAILABLE = False
    # Create dummy display functions
    def display(*args, **kwargs): pass
    def HTML(*args, **kwargs): pass
    def Audio(*args, **kwargs): pass

# Suppress warnings for cleaner output
warnings.filterwarnings('ignore')


def install_if_missing(packages):
    """Install packages if they're not already available in Kaggle"""
    for package in packages:
        try:
            package_name = package.split('==')[0].replace('-', '_')
            if package_name == 'yt_dlp':
                package_name = 'yt_dlp'
            __import__(package_name)
        except ImportError:
            print(f"Installing {package}...")
            subprocess.check_call([sys.executable, "-m", "pip", "install", package, "--quiet"])


# Required packages for Kaggle
required_packages = [
    "yt-dlp",
    "librosa",
    "soundfile",
    "transformers",
    "torch",
    "matplotlib",
    "seaborn"
]

print("🔧 Setting up environment...")
install_if_missing(required_packages)

# Now import the packages
import torch
from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification
import librosa
import soundfile as sf
import yt_dlp


class VideoAccentAnalyzer:
    def __init__(self, model_name="dima806/multiple_accent_classification"):
        """Initialize the accent analyzer for Kaggle environment"""
        self.model_name = model_name
        # Enhanced accent labels with better mapping
        self.accent_labels = [
            "british", "canadian", "us", "indian", "australian", "neutral"
        ]
        self.accent_display_names = {
            'british': '🇬🇧 British English',
            'us': '🇺🇸 American English',
            'australian': '🇦🇺 Australian English',
            'canadian': '🇨🇦 Canadian English',
            'indian': '🇮🇳 Indian English',
            'neutral': '🌐 Neutral English'
        }
        self.temp_dir = "/tmp/accent_analyzer"
        os.makedirs(self.temp_dir, exist_ok=True)
        self.model_loaded = False
        self._load_model()





    def _load_model(self):
        """Load the accent classification model with error handling"""
        print("🤖 Loading accent classification model...")
        try:
            self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(self.model_name)
            self.model = Wav2Vec2ForSequenceClassification.from_pretrained(self.model_name)
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            self.model.to(self.device)
            self.model.eval()  # Set to evaluation mode
            self.model_loaded = True
            print(f"✅ Model loaded successfully on {self.device}")
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            print("💡 Tip: Check your internet connection and Kaggle environment setup")
            raise

    def _validate_url(self, url):
        """Validate and normalize URL"""
        if not url or not isinstance(url, str):
            return False, "Invalid URL format"

        url = url.strip()
        if not url.startswith(('http://', 'https://')):
            return False, "URL must start with http:// or https://"

        return True, url

    def trim_video(self, input_path, output_path, duration):
        try:
            cmd = ['ffmpeg', '-i', input_path, '-t', str(duration), '-c', 'copy', output_path, '-y']
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
            if result.returncode == 0:
                print(f"✂️ Trimmed video to {duration} seconds")
                return output_path
            else:
                print(f"❌ Trimming failed: {result.stderr}")
                return input_path  # fallback to original
        except Exception as e:
            print(f"❌ Trimming exception: {e}")
            return input_path

    def _get_youtube_cookies(self):
        """Get YouTube cookies from browser"""
        import browser_cookie3

        try:
            # Try Firefox first
            cookies = browser_cookie3.firefox(domain_name='.youtube.com')
        except:
            try:
                # Try Chrome as fallback
                cookies = browser_cookie3.chrome(domain_name='.youtube.com')
            except:
                print("⚠️ Could not get cookies from browser")
                return None

        return cookies





    def download_video(self, url, max_duration=None):
        """Download video using yt-dlp with cookie support"""
        is_valid, result = self._validate_url(url)
        if not is_valid:
            print(f"❌ {result}")
            return None

        url = result
        output_path = os.path.join(self.temp_dir, "video.%(ext)s")

        # Enhanced yt-dlp options
        ydl_opts = {
            'outtmpl': output_path,
            'format': 'worst[ext=mp4]/worst',
            'quiet': False,
            'no_warnings': False,
            'socket_timeout': 60,
            'retries': 5
        }

        # Add cookies for YouTube URLs
        if 'youtube.com' in url or 'youtu.be' in url:
            cookies = self._get_youtube_cookies()
            if cookies:
                cookie_file = os.path.join(self.temp_dir, 'cookies.txt')
                with open(cookie_file, 'w') as f:
                    f.write('# Netscape HTTP Cookie File\n')
                    for cookie in cookies:
                        f.write(f'.youtube.com\tTRUE\t/\tFALSE\t{cookie.expires}\t{cookie.name}\t{cookie.value}\n')
                ydl_opts['cookiefile'] = cookie_file

        if max_duration:
            ydl_opts['match_filter'] = lambda info: None if info.get('duration', 0) <= 200000 else "Video too long"

        try:
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                print(f"📥 Downloading video from: {url}")
                start_time = time.time()

                # Get video info first
                try:
                    info = ydl.extract_info(url, download=False)
                    print(f"📺 Found video: {info.get('title', 'Unknown')} ({info.get('duration', 0)}s)")
                except Exception as e:
                    print(f"⚠️ Could not extract video info: {e}")

                # Download the video
                ydl.download([url])
                download_time = time.time() - start_time

                # Find downloaded file (try multiple patterns)
                video_path = None
                for file in os.listdir(self.temp_dir):
                    if file.startswith("video.") and os.path.getsize(
                            os.path.join(self.temp_dir, file)) > 1000:  # At least 1KB
                        potential_path = os.path.join(self.temp_dir, file)
                        print(f"📁 Found downloaded file: {file} ({os.path.getsize(potential_path) / 1024:.1f}KB)")

                        # Try basic validation - if ffprobe fails, still try to extract audio
                        if self._is_valid_video(potential_path):
                            print(f"✅ Video validation passed: {file}")
                            video_path = potential_path
                            break
                        else:
                            print(f"⚠️ Video validation failed, but continuing with: {file}")
                            video_path = potential_path  # Still try to use it
                            break

                if video_path:
                    print(f"✅ Downloaded video: {os.path.basename(video_path)} ({download_time:.1f}s)")
                    return video_path
                else:
                    print("❌ No video files found after download")
                    return None

        except Exception as e:
            print(f"⚠️ yt-dlp failed: {e}")
            return self._try_direct_download(url)

    def _is_valid_video(self, file_path):
        """Verify video file has valid structure (more lenient)"""
        try:
            # First check if file exists and has reasonable size
            if not os.path.exists(file_path) or os.path.getsize(file_path) < 1000:
                return False

            # Try ffprobe with more lenient settings
            result = subprocess.run(
                ['ffprobe', '-v', 'quiet', '-print_format', 'json', '-show_format', file_path],
                capture_output=True, text=True, timeout=15
            )

            if result.returncode == 0:
                try:
                    # Try to parse the JSON output
                    info = json.loads(result.stdout)
                    # Check if we have format information
                    if 'format' in info and 'duration' in info['format']:
                        return True
                except json.JSONDecodeError:
                    pass

            # If ffprobe fails, try a simpler check - just see if ffmpeg can read it
            result2 = subprocess.run(
                ['ffmpeg', '-i', file_path, '-t', '1', '-f', 'null', '-', '-v', 'quiet'],
                capture_output=True, text=True, timeout=15
            )

            return result2.returncode == 0

        except subprocess.TimeoutExpired:
            print("⚠️ Video validation timed out, assuming valid")
            return True  # If validation times out, assume it's valid
        except Exception as e:
            print(f"⚠️ Video validation error: {e}, assuming valid")
            return True  # If validation fails, assume it's valid and let audio extraction handle it

    def _try_direct_download(self, url):
        """Enhanced fallback for direct video URLs"""
        try:
            print("🔄 Trying direct download...")
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
            }

            response = requests.get(url, stream=True, timeout=60, headers=headers)
            response.raise_for_status()

            content_type = response.headers.get("Content-Type", "")
            if "text/html" in content_type:
                print("⚠️ Received HTML instead of video - check URL access")
                return None

            video_path = os.path.join(self.temp_dir, "video.mp4")
            file_size = 0

            with open(video_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        file_size += len(chunk)

            print(f"📁 Downloaded {file_size / (1024 * 1024):.1f} MB")

            if self._is_valid_video(video_path):
                print("✅ Direct download successful")
                return video_path
            else:
                print("❌ Downloaded file is not a valid video")
                return None

        except Exception as e:
            print(f"❌ Direct download failed: {e}")
            return None

    def extract_audio(self, video_path, max_duration=None):
        """Extract audio with improved error handling and progress"""
        audio_path = os.path.join(self.temp_dir, "audio.wav")

        # Enhanced ffmpeg command with better error handling
        cmd = ['ffmpeg', '-i', video_path, '-vn', '-acodec', 'pcm_s16le',
               '-ar', '16000', '-ac', '1', '-y', '-v', 'warning']

        if max_duration:
            cmd.extend(['-t', str(max_duration)])
        cmd.append(audio_path)

        try:
            print(f"🎵 Extracting audio (max {max_duration}s)...")
            start_time = time.time()

            # Run ffmpeg with more detailed output for debugging
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=180)
            extraction_time = time.time() - start_time

            if result.returncode == 0 and os.path.exists(audio_path) and os.path.getsize(audio_path) > 1000:
                file_size = os.path.getsize(audio_path) / (1024 * 1024)
                print(f"✅ Audio extracted successfully ({extraction_time:.1f}s, {file_size:.1f}MB)")
                return audio_path
            else:
                print(f"❌ FFmpeg stderr: {result.stderr}")
                print(f"❌ FFmpeg stdout: {result.stdout}")

                # Try alternative extraction method
                print("🔄 Trying alternative audio extraction...")
                cmd_alt = ['ffmpeg', '-i', video_path, '-vn', '-acodec', 'libmp3lame',
                           '-ar', '16000', '-ac', '1', '-y', '-v', 'warning']
                if max_duration:
                    cmd_alt.extend(['-t', str(max_duration)])

                audio_path_alt = os.path.join(self.temp_dir, "audio.mp3")
                cmd_alt.append(audio_path_alt)

                result_alt = subprocess.run(cmd_alt, capture_output=True, text=True, timeout=180)

                if result_alt.returncode == 0 and os.path.exists(audio_path_alt):
                    # Convert mp3 to wav
                    cmd_convert = ['ffmpeg', '-i', audio_path_alt, '-ar', '16000', '-ac', '1',
                                   audio_path, '-y', '-v', 'quiet']
                    result_convert = subprocess.run(cmd_convert, capture_output=True, text=True, timeout=60)

                    if result_convert.returncode == 0 and os.path.exists(audio_path):
                        file_size = os.path.getsize(audio_path) / (1024 * 1024)
                        print(f"✅ Alternative extraction successful ({file_size:.1f}MB)")
                        return audio_path

                raise Exception(f"Both extraction methods failed. FFmpeg error: {result.stderr}")

        except subprocess.TimeoutExpired:
            print("❌ Audio extraction timed out")
            return None
        except Exception as e:
            print(f"❌ Audio extraction failed: {e}")
            return None

    def classify_accent(self, audio_path):
        """Enhanced accent classification with better preprocessing"""
        if not self.model_loaded:
            print("❌ Model not loaded properly")
            return None

        try:
            print("🔍 Loading and preprocessing audio...")
            audio, sr = librosa.load(audio_path, sr=16000)

            # Enhanced preprocessing
            if len(audio) == 0:
                print("❌ Empty audio file")
                return None

            # Remove silence from beginning and end
            audio_trimmed, _ = librosa.effects.trim(audio, top_db=20)

            # Use multiple chunks for better accuracy if audio is long
            chunk_size = 16000 * 20  # 20 seconds chunks
            chunks = []

            if len(audio_trimmed) > chunk_size:
                # Split into overlapping chunks
                step_size = chunk_size // 2
                for i in range(0, len(audio_trimmed) - chunk_size + 1, step_size):
                    chunks.append(audio_trimmed[i:i + chunk_size])
                if len(audio_trimmed) % step_size != 0:
                    chunks.append(audio_trimmed[-chunk_size:])
            else:
                chunks = [audio_trimmed]

            print(f"🎯 Analyzing {len(chunks)} audio chunk(s)...")

            all_predictions = []

            for i, chunk in enumerate(chunks[:3]):  # Limit to 3 chunks for efficiency
                inputs = self.feature_extractor(
                    chunk,
                    sampling_rate=16000,
                    return_tensors="pt",
                    padding=True,
                    max_length=16000 * 20,
                    truncation=True
                )
                inputs = {k: v.to(self.device) for k, v in inputs.items()}

                with torch.no_grad():
                    outputs = self.model(**inputs)
                    logits = outputs.logits
                    probabilities = torch.nn.functional.softmax(logits, dim=-1)
                    all_predictions.append(probabilities[0].cpu().numpy())

            # Average predictions across chunks
            avg_probabilities = sum(all_predictions) / len(all_predictions)
            predicted_idx = avg_probabilities.argmax()
            predicted_idx = min(predicted_idx, len(self.accent_labels) - 1)

            # Calculate English confidence (exclude 'neutral' for this calculation)
            english_accents = ["british", "canadian", "us", "australian", "indian"]
            english_confidence = sum(
                avg_probabilities[i] * 100
                for i, label in enumerate(self.accent_labels)
                if label in english_accents
            )

            results = {
                'predicted_accent': self.accent_labels[predicted_idx],
                'accent_confidence': avg_probabilities[predicted_idx] * 100,
                'english_confidence': english_confidence,
                'audio_duration': len(audio) / 16000,
                'processed_duration': len(audio_trimmed) / 16000,
                'chunks_analyzed': len(all_predictions),
                'all_probabilities': {
                    self.accent_labels[i]: avg_probabilities[i] * 100
                    for i in range(len(self.accent_labels))
                },
                'is_english_likely': english_confidence > 60,
                'audio_quality_score': self._assess_audio_quality(audio_trimmed)
            }

            print(f"✅ Classification complete ({results['chunks_analyzed']} chunks)")
            return results

        except Exception as e:
            print(f"❌ Classification failed: {e}")
            return None

    def _assess_audio_quality(self, audio):
        """Assess audio quality for better result interpretation"""
        try:
            # Simple quality metrics
            rms_energy = librosa.feature.rms(y=audio)[0].mean()
            zero_crossing_rate = librosa.feature.zero_crossing_rate(audio)[0].mean()

            # Normalize to 0-100 scale
            quality_score = min(100, (rms_energy * 1000 + (1 - zero_crossing_rate) * 50))
            return max(0, quality_score)
        except:
            return 50  # Default moderate quality

    def analyze_video_url(self, url, max_duration=30):
        """Complete pipeline with enhanced error handling"""
        print(f"🎬 Starting analysis of: {url}")
        print(f"⏱️ Max duration: {max_duration} seconds")

        video_path = self.download_video(url, max_duration)
        if not video_path:
            return {"error": "Failed to download video", "url": url}

        audio_path = self.extract_audio(video_path, max_duration)
        if not audio_path:
            return {"error": "Failed to extract audio", "url": url}

        results = self.classify_accent(audio_path)
        if not results:
            return {"error": "Failed to classify accent", "url": url}

        results.update({
            'source_url': url,
            'video_file': os.path.basename(video_path),
            'audio_file': os.path.basename(audio_path),
            'analysis_timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
        })

        return results

    def analyze_local_video(self, file_path, max_duration=30):
        """Enhanced local video analysis"""
        print(f"🎬 Starting analysis of local file: {file_path}")
        print(f"⏱️ Max duration: {max_duration} seconds")

        if not os.path.isfile(file_path):
            return {"error": f"File not found: {file_path}"}

        # Check file size
        file_size = os.path.getsize(file_path) / (1024 * 1024)  # MB
        print(f"📁 File size: {file_size:.1f} MB")

        video_filename = os.path.basename(file_path)
        print(f"✅ Using local video: {video_filename}")

        audio_path = self.extract_audio(file_path, max_duration)
        if not audio_path:
            return {"error": "Failed to extract audio"}

        results = self.classify_accent(audio_path)
        if not results:
            return {"error": "Failed to classify accent"}

        results.update({
            'source_file': file_path,
            'video_file': video_filename,
            'audio_file': os.path.basename(audio_path),
            'file_size_mb': file_size,
            'is_local': True,
            'analysis_timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
        })

        return results


def display_results(self, results):
    """Display results in text format"""
    if 'error' in results:
        print(f"❌ {results['error']}")
        return

    accent = results['predicted_accent']
    confidence = results['accent_confidence']
    english_conf = results['english_confidence']
    duration = results['audio_duration']
    processed_duration = results.get('processed_duration', duration)
    quality_score = results.get('audio_quality_score', 50)

    accent_display = self.accent_display_names.get(accent, accent.title())

    print(f"\n=== Accent Analysis Results ===")
    print(f"Predicted Accent: {accent_display}")
    print(f"Confidence: {confidence:.1f}%")
    print(f"English Confidence: {english_conf:.1f}%")
    print(f"Audio Duration: {duration:.1f}s")
    print(f"Processed Duration: {processed_duration:.1f}s")
    print(f"Audio Quality: {quality_score:.0f}/100")
    print(f"Chunks Analyzed: {results.get('chunks_analyzed', 1)}")


def _plot_probabilities(self, probabilities):
    """Create a visualization of accent probabilities"""
    try:
        plt.figure(figsize=(10, 6))

        accents = [self.accent_display_names.get(acc, acc.title()) for acc in probabilities.keys()]
        probs = list(probabilities.values())

        # Create color map
        colors = ['#4CAF50' if p == max(probs) else '#2196F3' if p >= 20 else '#FFC107' if p >= 10 else '#9E9E9E'
                  for p in probs]

        bars = plt.bar(accents, probs, color=colors, alpha=0.8, edgecolor='black', linewidth=0.5)

        plt.title('Accent Classification Probabilities', fontsize=16, fontweight='bold', pad=20)
        plt.xlabel('Accent Type', fontsize=12)
        plt.ylabel('Probability (%)', fontsize=12)
        plt.xticks(rotation=45, ha='right')
        plt.grid(axis='y', alpha=0.3)

        # Add value labels on bars
        for bar, prob in zip(bars, probs):
            height = bar.get_height()
            plt.text(bar.get_x() + bar.get_width() / 2., height + 0.5,
                     f'{prob:.1f}%', ha='center', va='bottom', fontweight='bold')

        plt.tight_layout()
        plt.show()

    except Exception as e:
        print(f"⚠️ Could not create visualization: {e}")


def batch_analyze(self, urls, max_duration=30):
    """Analyze multiple videos with progress tracking"""
    results = []
    failed_count = 0

    print(f"🚀 Starting batch analysis of {len(urls)} videos")

    for i, url in enumerate(urls, 1):
        print(f"\n{'=' * 60}")
        print(f"Processing video {i}/{len(urls)}")

        result = self.analyze_video_url(url, max_duration)
        result['video_index'] = i

        if 'error' in result:
            failed_count += 1
            print(f"❌ Failed: {result['error']}")
        else:
            print(f"✅ Success: {result['predicted_accent']} ({result['accent_confidence']:.1f}%)")

        results.append(result)
        self.display_results(result)

        # Small delay to prevent overwhelming servers
        if i < len(urls):
            time.sleep(1)

    # Summary
    success_count = len(urls) - failed_count
    print(f"\n📈 Batch Analysis Summary:")
    print(f"   ✅ Successful: {success_count}/{len(urls)}")
    print(f"   ❌ Failed: {failed_count}/{len(urls)}")

    return results


def export_results(self, results, filename="accent_analysis_results.json"):
    """Export results to JSON file"""
    try:
        with open(filename, 'w') as f:
            json.dump(results, f, indent=2, default=str)
        print(f"💾 Results exported to {filename}")
    except Exception as e:
        print(f"❌ Export failed: {e}")


def cleanup(self):
    """Clean up temporary files"""
    try:
        import shutil
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir, ignore_errors=True)
            print("🧹 Cleaned up temporary files")
    except Exception as e:
        print(f"⚠️ Cleanup warning: {e}")


# Helper Functions
def show_examples():
    """Show usage examples"""
    examples = {
        "Loom": "https://www.loom.com/share/abc123def456",
        "Direct MP4": "https://example.com/video.mp4",
        "Local File": "/local/input/dataset/video.mp4"
    }

    print("\n🎯 Supported Video Formats:")
    for platform, example in examples.items():
        print(f"  {platform:12}: {example}")

    print("\n💡 Usage Tips:")
    print("  • Keep videos under 2 minutes for best results")
    print("  • Ensure clear audio quality")
    print("  • Multiple speakers may affect accuracy")
    print("  • Model works best with sustained speech")


def quick_test_url():
    """Interactive test for video URLs"""
    print("🔍 Quick Test Mode for Video URLs")
    print("🎯 Supported:  Loom, Direct MP4 links")
    print("💡 Examples:")
    print("   Loom: https://www.loom.com/share/VIDEO_ID")
    print("   Direct: https://example.com/video.mp4")

    url = input("\n📎 Enter your video URL (Loom, MP4 , etc.): ").strip()
    if not url:
        print("❌ No URL provided.")
        return None

    max_duration = input("⏱️ Max duration in seconds (default 20): ").strip()
    try:
        max_duration = int(max_duration) if max_duration else 20
    except ValueError:
        max_duration = 20
        print(f"⚠️ Invalid duration, using {max_duration} seconds")

    analyzer = VideoAccentAnalyzer()
    try:
        print(f"\n🚀 Starting analysis...")
        results = analyzer.analyze_video_url(url, max_duration=max_duration)
        analyzer.display_results(results)
        return results
    finally:
        analyzer.cleanup()


def demo_analysis():
    """Demo function with example usage"""
    print("🎬 Video Accent Analyzer Demo")
    print("=" * 50)

    # Initialize analyzer
    analyzer = VideoAccentAnalyzer()

    # Example analysis (replace with actual video URL)
    example_url = "https://example.com/video.mp4"  # Replace with real URL
    print(f"\n🎯 Example: Analyzing {example_url}")

    # Uncomment to run actual analysis
    # results = analyzer.analyze_video_url(example_url, max_duration=30)
    # analyzer.display_results(results)
    # analyzer.cleanup()

    print("\n📚 To use the analyzer:")
    print("1. analyzer = VideoAccentAnalyzer()")
    print("2. results = analyzer.analyze_video_url('your-url', max_duration=30)")
    print("3. analyzer.display_results(results)")
    print("4. analyzer.cleanup()  # Clean up temporary files")


# Show examples on import
show_examples()