#!/usr/bin/env python3
"""
Extract accent audio samples with pure English lyrics from captioned vocal stems dataset.

This script performs the following operations:
1. Extract vocal stem records with specific accents and pure English lyrics
2. Export matching records as individual JSON files
3. Copy corresponding audio files (opus format) from source
4. Convert opus audio files to MP3 format
5. Clean up intermediate opus files
6. Create zip archives of the MP3 collections

Usage:
    python extract_accent_audio_samples.py \
        --input-jsonl /path/to/captioned_dataset.jsonl \
        --output-dir /path/to/output \
        --accents "german,british,scottish" \
        --samples-per-accent 15 \
        --convert-to-mp3 \
        --create-zip
"""

import argparse
import json
import os
import re
import shutil
import subprocess
import zipfile
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

try:
    import langdetect
except ImportError:
    print("Warning: langdetect not installed. Language detection will be disabled.")
    print("Install with: pip install langdetect")
    langdetect = None


class AccentAudioExtractor:
    """Extract and process accent audio samples with pure English validation."""

    # Common English words for validation
    COMMON_ENGLISH_WORDS = {
        "the",
        "be",
        "to",
        "of",
        "and",
        "a",
        "in",
        "that",
        "have",
        "i",
        "it",
        "for",
        "not",
        "on",
        "with",
        "he",
        "as",
        "you",
        "do",
        "at",
        "this",
        "but",
        "his",
        "by",
        "from",
        "they",
        "we",
        "say",
        "her",
        "she",
        "or",
        "an",
        "will",
        "my",
        "one",
        "all",
        "would",
        "there",
        "their",
        "what",
        "so",
        "up",
        "out",
        "if",
        "about",
        "who",
        "get",
        "which",
        "go",
        "me",
        "when",
        "make",
        "can",
        "like",
        "time",
        "no",
        "just",
        "him",
        "know",
    }

    # Accent keyword mappings
    ACCENT_KEYWORDS = {
        "german": ["german accent"],
        "british": ["british accent", "uk accent", "english accent"],
        "scottish": ["scottish accent", "scots accent"],
        "indian": ["indian accent", "hindi accent"],
        "japanese": ["japanese accent"],
        "korean": ["korean accent"],
        "slavic": ["slavic accent", "russian accent", "eastern european accent"],
        "irish": ["irish accent"],
        "australian": ["australian accent", "aussie accent"],
        "french": ["french accent"],
        "italian": ["italian accent"],
        "spanish": ["spanish accent", "hispanic accent"],
        "chinese": ["chinese accent", "mandarin accent"],
        "american": ["american accent", "us accent"],
    }

    def __init__(self, input_jsonl: str, output_dir: str):
        """Initialize the extractor.

        Args:
            input_jsonl: Path to input JSONL file with captioned vocal stems
            output_dir: Base output directory for all extracted files
        """
        self.input_jsonl = input_jsonl
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def is_strictly_english(self, text: str) -> bool:
        """Check if text contains only English lyrics.

        Args:
            text: Lyrics text to validate

        Returns:
            True if text is strictly English, False otherwise
        """
        if not text:
            return False

        # Remove annotations like [Verse], [Chorus], etc.
        cleaned_text = re.sub(r"\[.*?\]", "", text)
        cleaned_text = cleaned_text.strip()

        if not cleaned_text:
            return False

        # Use langdetect if available
        if langdetect:
            try:
                detected_lang = langdetect.detect(cleaned_text)
                if detected_lang != "en":
                    return False
            except:
                # If detection fails, continue with other checks
                pass

        # Check for non-Latin characters
        if re.search(r"[а-яА-ЯёЁ]", cleaned_text):  # Cyrillic
            return False
        if re.search(r"[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff]", cleaned_text):  # CJK
            return False
        if re.search(r"[\u0600-\u06ff\u0750-\u077f]", cleaned_text):  # Arabic
            return False
        if re.search(r"[\u0590-\u05ff]", cleaned_text):  # Hebrew
            return False

        # Check for common English words
        words = set(cleaned_text.lower().split())
        english_word_count = len(words & self.COMMON_ENGLISH_WORDS)

        # Require at least 10 common English words for longer texts
        if len(words) > 50 and english_word_count < 10:
            return False

        return True

    def has_accent(self, record: Dict, accent: str) -> bool:
        """Check if record has the specified accent in vocal captions.

        Args:
            record: Record dictionary with stems_captions
            accent: Accent type to search for

        Returns:
            True if accent is found in vocal captions
        """
        stems_captions = record.get("stems_captions", {})
        vocal_captions = stems_captions.get("Vocals", [])

        # Get accent keywords for this accent type
        keywords = self.ACCENT_KEYWORDS.get(accent, [accent + " accent"])

        for caption_entry in vocal_captions:
            if caption_entry.get("prompt_type") == "voice_description_keywords":
                caption = caption_entry.get("caption", "").lower()
                for keyword in keywords:
                    if keyword in caption:
                        return True

        return False

    def extract_accent_samples(
        self, accents: List[str], samples_per_accent: int = 15, max_records_to_scan: Optional[int] = None
    ) -> Dict[str, List[Dict]]:
        """Extract samples for specified accents with pure English lyrics.

        Args:
            accents: List of accent types to extract
            samples_per_accent: Number of samples to extract per accent
            max_records_to_scan: Maximum records to scan (None for all)

        Returns:
            Dictionary mapping accent to list of extracted records
        """
        accent_samples = {accent: [] for accent in accents}
        accent_counts = {accent: 0 for accent in accents}

        print(f"Extracting samples for accents: {', '.join(accents)}")
        print(f"Target: {samples_per_accent} samples per accent")
        print("=" * 60)

        records_scanned = 0
        records_rejected = 0

        with open(self.input_jsonl, "r") as f:
            for line_num, line in enumerate(f, 1):
                if max_records_to_scan and records_scanned >= max_records_to_scan:
                    break

                if not line.strip():
                    continue

                records_scanned += 1

                # Check if we have enough samples for all accents
                if all(accent_counts[a] >= samples_per_accent for a in accents):
                    break

                try:
                    record = json.loads(line.strip())

                    # Check if record has English lyrics
                    text = record.get("text", "")
                    if not self.is_strictly_english(text):
                        records_rejected += 1
                        continue

                    # Check each accent type
                    for accent in accents:
                        if accent_counts[accent] >= samples_per_accent:
                            continue

                        if self.has_accent(record, accent):
                            accent_samples[accent].append(record)
                            accent_counts[accent] += 1
                            print(
                                f"  Found {accent}: {record.get('id')} "
                                f"({accent_counts[accent]}/{samples_per_accent})"
                            )
                            break  # Each record only counts for one accent

                except json.JSONDecodeError:
                    continue

                # Progress update
                if records_scanned % 10000 == 0:
                    print(
                        f"  Scanned {records_scanned:,} records, "
                        f"rejected {records_rejected:,} non-English..."
                    )
                    for accent in accents:
                        if accent_counts[accent] < samples_per_accent:
                            print(f"    {accent}: {accent_counts[accent]}/{samples_per_accent}")

        print("=" * 60)
        print(f"Extraction complete! Scanned {records_scanned:,} records")
        print(f"Rejected {records_rejected:,} records with non-English lyrics")

        for accent in accents:
            print(f"  {accent}: {len(accent_samples[accent])} samples extracted")

        return accent_samples

    def export_samples_to_json(self, accent_samples: Dict[str, List[Dict]]) -> Dict[str, Path]:
        """Export accent samples to individual JSON files.

        Args:
            accent_samples: Dictionary mapping accent to records

        Returns:
            Dictionary mapping accent to output directory path
        """
        output_dirs = {}

        for accent, samples in accent_samples.items():
            # Create accent-specific directory
            accent_dir = self.output_dir / f"{accent}_accent_pure_english"
            accent_dir.mkdir(exist_ok=True)
            output_dirs[accent] = accent_dir

            print(f"\nExporting {accent} samples to {accent_dir}")

            for idx, record in enumerate(samples, 1):
                record_id = record.get("id", f"record_{idx}")
                output_file = accent_dir / f"{accent}_{idx:03d}_{record_id}.json"

                with open(output_file, "w") as f:
                    json.dump(record, f, indent=2, ensure_ascii=False)

                print(f"  Exported: {output_file.name}")

        return output_dirs

    def copy_audio_files(
        self, accent_samples: Dict[str, List[Dict]], output_dirs: Dict[str, Path]
    ) -> Dict[str, Path]:
        """Copy audio files for each accent to audio directories.

        Args:
            accent_samples: Dictionary mapping accent to records
            output_dirs: Dictionary mapping accent to JSON output directories

        Returns:
            Dictionary mapping accent to audio output directory
        """
        audio_dirs = {}

        for accent, samples in accent_samples.items():
            # Create audio directory
            audio_dir = self.output_dir / f"{accent}_accent_pure_english_audio"
            audio_dir.mkdir(exist_ok=True)
            audio_dirs[accent] = audio_dir

            print(f"\nCopying audio files for {accent} to {audio_dir}")

            copied_count = 0
            for record in samples:
                local_filepath = record.get("local_filepath")
                if local_filepath and os.path.exists(local_filepath):
                    record_id = record.get("id", "unknown")
                    dest_path = audio_dir / f"{record_id}.opus"

                    shutil.copy2(local_filepath, dest_path)
                    copied_count += 1
                    print(f"  Copied: {record_id}.opus")
                else:
                    print(f"  Warning: Could not find audio file for {record.get('id')}")

            print(f"  Total: {copied_count} audio files copied")

        return audio_dirs

    def convert_to_mp3(self, audio_dirs: Dict[str, Path], bitrate: str = "192k") -> None:
        """Convert opus files to MP3 format.

        Args:
            audio_dirs: Dictionary mapping accent to audio directories
            bitrate: MP3 bitrate (default: 192k)
        """
        print(f"\nConverting opus files to MP3 (bitrate: {bitrate})...")

        for accent, audio_dir in audio_dirs.items():
            opus_files = list(audio_dir.glob("*.opus"))
            print(f"\n{accent}: Converting {len(opus_files)} files...")

            for opus_file in opus_files:
                mp3_file = opus_file.with_suffix(".mp3")

                cmd = [
                    "ffmpeg",
                    "-i",
                    str(opus_file),
                    "-acodec",
                    "libmp3lame",
                    "-b:a",
                    bitrate,
                    str(mp3_file),
                    "-y",
                    "-loglevel",
                    "error",
                ]

                subprocess.run(cmd, check=False)
                print(f"  Converted: {mp3_file.name}")

            mp3_count = len(list(audio_dir.glob("*.mp3")))
            print(f"  Completed: {mp3_count} MP3 files created")

    def cleanup_opus_files(self, audio_dirs: Dict[str, Path]) -> None:
        """Remove opus files after MP3 conversion.

        Args:
            audio_dirs: Dictionary mapping accent to audio directories
        """
        print("\nCleaning up opus files...")

        total_removed = 0
        for accent, audio_dir in audio_dirs.items():
            opus_files = list(audio_dir.glob("*.opus"))

            for opus_file in opus_files:
                opus_file.unlink()

            print(f"  {accent}: Removed {len(opus_files)} opus files")
            total_removed += len(opus_files)

        print(f"Total removed: {total_removed} opus files")

    def create_zip_archives(
        self, audio_dirs: Dict[str, Path], create_individual: bool = True, create_combined: bool = True
    ) -> List[Path]:
        """Create zip archives of audio collections.

        Args:
            audio_dirs: Dictionary mapping accent to audio directories
            create_individual: Create individual zip for each accent
            create_combined: Create combined zip with all accents

        Returns:
            List of created zip file paths
        """
        zip_files = []

        if create_individual:
            print("\nCreating individual zip files...")

            for accent, audio_dir in audio_dirs.items():
                zip_name = f"{accent}_accent_audio.zip"
                zip_path = self.output_dir / zip_name

                mp3_files = list(audio_dir.glob("*.mp3"))

                with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
                    for mp3_file in mp3_files:
                        zipf.write(mp3_file, mp3_file.name)

                size_mb = zip_path.stat().st_size / (1024 * 1024)
                print(f"  {accent}: {zip_name} ({size_mb:.1f} MB, {len(mp3_files)} files)")
                zip_files.append(zip_path)

        if create_combined:
            print("\nCreating combined zip file...")

            combined_zip = self.output_dir / "all_accent_audio_samples.zip"

            with zipfile.ZipFile(combined_zip, "w", zipfile.ZIP_DEFLATED) as zipf:
                total_files = 0

                for accent, audio_dir in audio_dirs.items():
                    mp3_files = list(audio_dir.glob("*.mp3"))

                    for mp3_file in mp3_files:
                        # Preserve directory structure in combined zip
                        arcname = f"{audio_dir.name}/{mp3_file.name}"
                        zipf.write(mp3_file, arcname)
                        total_files += 1

                size_mb = combined_zip.stat().st_size / (1024 * 1024)
                print(f"  Combined: {combined_zip.name} ({size_mb:.1f} MB, {total_files} files)")
                zip_files.append(combined_zip)

        return zip_files


def main():
    """Main execution function."""
    parser = argparse.ArgumentParser(
        description="Extract accent audio samples with pure English lyrics",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        "--input-jsonl", required=True, help="Path to input JSONL file with captioned vocal stems"
    )

    parser.add_argument(
        "--output-dir", required=True, help="Base output directory for all extracted files"
    )

    parser.add_argument(
        "--accents",
        default="german,british,scottish,indian,japanese,korean,slavic",
        help="Comma-separated list of accent types to extract",
    )

    parser.add_argument(
        "--samples-per-accent",
        type=int,
        default=15,
        help="Number of samples to extract per accent (default: 15)",
    )

    parser.add_argument("--max-scan", type=int, help="Maximum number of records to scan (default: all)")

    parser.add_argument("--convert-to-mp3", action="store_true", help="Convert opus files to MP3 format")

    parser.add_argument(
        "--mp3-bitrate", default="192k", help="MP3 bitrate for conversion (default: 192k)"
    )

    parser.add_argument(
        "--cleanup-opus", action="store_true", help="Remove opus files after MP3 conversion"
    )

    parser.add_argument(
        "--create-zip", action="store_true", help="Create zip archives of audio collections"
    )

    parser.add_argument(
        "--zip-individual", action="store_true", help="Create individual zip file for each accent"
    )

    parser.add_argument(
        "--zip-combined", action="store_true", help="Create combined zip file with all accents"
    )

    args = parser.parse_args()

    # Parse accent list
    accents = [a.strip() for a in args.accents.split(",")]

    # Initialize extractor
    extractor = AccentAudioExtractor(args.input_jsonl, args.output_dir)

    # Extract samples
    print("\n" + "=" * 60)
    print("STEP 1: Extracting accent samples with pure English lyrics")
    print("=" * 60)
    accent_samples = extractor.extract_accent_samples(accents, args.samples_per_accent, args.max_scan)

    # Export to JSON
    print("\n" + "=" * 60)
    print("STEP 2: Exporting samples to JSON files")
    print("=" * 60)
    json_dirs = extractor.export_samples_to_json(accent_samples)

    # Copy audio files
    print("\n" + "=" * 60)
    print("STEP 3: Copying audio files")
    print("=" * 60)
    audio_dirs = extractor.copy_audio_files(accent_samples, json_dirs)

    # Convert to MP3 if requested
    if args.convert_to_mp3:
        print("\n" + "=" * 60)
        print("STEP 4: Converting audio to MP3")
        print("=" * 60)
        extractor.convert_to_mp3(audio_dirs, args.mp3_bitrate)

        # Cleanup opus files if requested
        if args.cleanup_opus:
            print("\n" + "=" * 60)
            print("STEP 5: Cleaning up opus files")
            print("=" * 60)
            extractor.cleanup_opus_files(audio_dirs)

    # Create zip archives if requested
    if args.create_zip or args.zip_individual or args.zip_combined:
        print("\n" + "=" * 60)
        print("STEP 6: Creating zip archives")
        print("=" * 60)

        create_individual = args.create_zip or args.zip_individual
        create_combined = args.create_zip or args.zip_combined

        zip_files = extractor.create_zip_archives(audio_dirs, create_individual, create_combined)

        print(f"\nCreated {len(zip_files)} zip file(s)")

    print("\n" + "=" * 60)
    print("✅ All operations completed successfully!")
    print("=" * 60)

    print(f"\nOutput directory: {extractor.output_dir}")
    print("\nCreated directories:")
    for accent in accents:
        if accent in json_dirs:
            print(f"  JSON:  {json_dirs[accent]}")
        if accent in audio_dirs:
            print(f"  Audio: {audio_dirs[accent]}")


if __name__ == "__main__":
    main()