#!/usr/bin/env python3
"""
Accent Detection Script for Voice Designer
Searches through metas_v6 JSONL files to find records with specific accents for accent detection testing.
"""

import json
import argparse
import re
from pathlib import Path
from typing import List, Dict, Set
import logging
from tqdm import tqdm
from collections import defaultdict


# Target accents to search for
TARGET_ACCENTS = {
    "scottish": ["scottish", "scots", "glasgow", "edinburgh", "highland", "aberdeen"],
    "australian": ["australian", "aussie", "sydney", "melbourne", "brisbane", "perth", "adelaide"],
    "german": ["german", "deutsch", "berlin", "munich", "hamburg", "cologne", "frankfurt"],
    "bavarian": ["bavarian", "bayern", "munich", "oktoberfest", "alpine"],
    "chinese": ["chinese", "mandarin", "cantonese", "beijing", "shanghai", "hong kong", "taiwan"],
    "japanese": ["japanese", "tokyo", "osaka", "kyoto", "nihon", "nippon"],
    "korean": ["korean", "seoul", "busan", "hangul", "k-pop", "kpop"],
    "irish": ["irish", "dublin", "belfast", "cork", "galway", "celtic"],
    "french": ["french", "paris", "marseille", "lyon", "quebec", "francais"],
    "italian": ["italy", "italian", "rome", "milan", "naples", "venice", "sicilian"],
    "spanish": ["spanish", "madrid", "barcelona", "seville", "valencia", "andalusian"],
    "russian": ["russian", "moscow", "petersburg", "siberian", "vladimir"],
    "indian": ["indian", "hindi", "mumbai", "delhi", "bangalore", "bollywood"],
    "british": [
        "british",
        "london",
        "manchester",
        "liverpool",
        "birmingham",
        "cockney",
        "posh",
        "received pronunciation",
    ],
    "southern_us": ["southern", "texas", "georgia", "alabama", "mississippi", "louisiana", "drawl"],
    "new_york": ["new york", "brooklyn", "bronx", "manhattan", "queens"],
}


def setup_logging(verbose: bool = False):
    """Setup logging configuration"""
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(level=level, format="%(asctime)s - %(levelname)s - %(message)s")


def contains_accent_keywords(text: str, accent_keywords: List[str]) -> bool:
    """Check if text contains any of the accent keywords (case insensitive)"""
    if not text:
        return False

    text_lower = text.lower()
    return any(keyword.lower() in text_lower for keyword in accent_keywords)


def search_record_for_accents(record: Dict) -> Dict[str, List[str]]:
    """
    Search a record for accent-related keywords in various fields

    Returns:
        Dictionary mapping accent names to list of matching fields/contexts
    """
    found_accents = defaultdict(list)

    # Fields to search in order of priority
    search_fields = [
        ("text", record.get("text", "")),
        ("tags", " ".join(record.get("tags", []))),
        ("artist_info", str(record.get("artist_ids", []))),
        ("playlist_info", str(record.get("playlist_ids", []))),
    ]

    for accent_name, keywords in TARGET_ACCENTS.items():
        for field_name, field_content in search_fields:
            if contains_accent_keywords(field_content, keywords):
                # Find which specific keywords matched
                matched_keywords = [kw for kw in keywords if kw.lower() in field_content.lower()]
                found_accents[accent_name].append(
                    {
                        "field": field_name,
                        "matched_keywords": matched_keywords,
                        "context": field_content[:200] + ("..." if len(field_content) > 200 else ""),
                    }
                )

    return dict(found_accents)


def has_vocal_stems(record: Dict) -> bool:
    """Check if record has vocal-related stems"""
    stems = record.get("stems", {})
    if not stems:
        return False

    vocal_keywords = ["vocal", "vox"]
    return any(
        any(keyword in stem_name.lower() for keyword in vocal_keywords) for stem_name in stems.keys()
    )


def process_jsonl_file(
    input_file: str,
    output_file: str,
    target_accents: List[str] = None,
    vocal_stems_only: bool = True,
    max_per_accent: int = 10,
    limit: int = None,
) -> Dict[str, int]:
    """
    Process JSONL file to find records with target accents

    Args:
        input_file: Path to input JSONL file
        output_file: Path to output JSONL file
        target_accents: List of accent names to search for (None = all)
        vocal_stems_only: Only include records with vocal stems
        max_per_accent: Maximum records per accent type
        limit: Maximum total records to process

    Returns:
        Dictionary with counts of found accents
    """
    input_path = Path(input_file)
    output_path = Path(output_file)

    if not input_path.exists():
        raise FileNotFoundError(f"Input file not found: {input_file}")

    # Create output directory
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Track findings
    accent_counts = defaultdict(int)
    found_records = []
    processed_count = 0

    # Filter target accents if specified
    search_accents = target_accents or list(TARGET_ACCENTS.keys())

    logging.info(f"Searching for accents: {', '.join(search_accents)}")
    logging.info(f"Input: {input_file}")
    logging.info(f"Output: {output_file}")
    logging.info(f"Vocal stems only: {vocal_stems_only}")
    logging.info(f"Max per accent: {max_per_accent}")

    with open(input_path, "r") as f:
        # Get total lines for progress bar
        total_lines = sum(1 for _ in f)
        f.seek(0)

        with tqdm(total=min(total_lines, limit) if limit else total_lines, desc="Processing") as pbar:
            for line_num, line in enumerate(f):
                if limit and processed_count >= limit:
                    break

                try:
                    record = json.loads(line.strip())
                    processed_count += 1
                    pbar.update(1)

                    # Skip if no stems when vocal_stems_only is True
                    if vocal_stems_only and not has_vocal_stems(record):
                        continue

                    # Search for accents
                    found_accents = search_record_for_accents(record)

                    if found_accents:
                        # Check if we need any of these accents
                        relevant_accents = {
                            acc: matches
                            for acc, matches in found_accents.items()
                            if acc in search_accents and accent_counts[acc] < max_per_accent
                        }

                        if relevant_accents:
                            # Add metadata about found accents
                            record["_accent_analysis"] = {
                                "found_accents": relevant_accents,
                                "search_priority": min(
                                    accent_counts[acc] for acc in relevant_accents.keys()
                                ),
                                "line_number": line_num + 1,
                            }

                            found_records.append(record)

                            # Update counts
                            for accent in relevant_accents.keys():
                                accent_counts[accent] += 1

                            logging.debug(
                                f"Found record {record.get('id', 'unknown')} with accents: {list(relevant_accents.keys())}"
                            )

                except json.JSONDecodeError as e:
                    logging.warning(f"Skipping malformed JSON at line {line_num + 1}: {e}")
                    continue
                except Exception as e:
                    logging.error(f"Error processing line {line_num + 1}: {e}")
                    continue

    # Sort records by accent priority and diversity
    found_records.sort(
        key=lambda r: (
            r["_accent_analysis"]["search_priority"],  # Prefer less common accents
            len(r["_accent_analysis"]["found_accents"]),  # Prefer records with multiple accents
        )
    )

    # Write results
    with open(output_path, "w") as f:
        for record in found_records:
            f.write(json.dumps(record) + "\n")

    logging.info(f"Found {len(found_records)} records with target accents")
    logging.info("Accent distribution:")
    for accent, count in sorted(accent_counts.items()):
        logging.info(f"  {accent}: {count} records")

    return dict(accent_counts)


def main():
    parser = argparse.ArgumentParser(
        description="Find records with specific accents for voice analysis testing",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Find all accents in validation set
  python find_accent_records.py --input /app2/suno/data/auk_v0/metas_v6_val.jsonl --output accent_test_records.jsonl

  # Find only Scottish and Australian accents
  python find_accent_records.py --input metas_v6_val.jsonl --output scottish_aussie.jsonl --accents scottish australian

  # Find German accents, allow non-vocal stems, get up to 20 per accent
  python find_accent_records.py --input metas_v6_val.jsonl --output german_test.jsonl --accents german bavarian --no-vocal-only --max-per-accent 20

Available accents: scottish, australian, german, bavarian, chinese, japanese, korean, irish, french, italian, spanish, russian, indian, british, southern_us, new_york
        """,
    )

    parser.add_argument(
        "--input", "-i", type=str, required=True, help="Input JSONL file (e.g., metas_v6_val.jsonl)"
    )

    parser.add_argument(
        "--output", "-o", type=str, required=True, help="Output JSONL file for found records"
    )

    parser.add_argument(
        "--accents",
        nargs="+",
        choices=list(TARGET_ACCENTS.keys()),
        default=None,
        help="Specific accents to search for (default: all)",
    )

    parser.add_argument(
        "--vocal-only",
        action="store_true",
        default=True,
        help="Only include records with vocal stems (default: True)",
    )

    parser.add_argument(
        "--no-vocal-only", action="store_true", help="Include records without vocal stems"
    )

    parser.add_argument(
        "--max-per-accent", type=int, default=10, help="Maximum records per accent type (default: 10)"
    )

    parser.add_argument(
        "--limit", type=int, default=None, help="Maximum records to process from input (for testing)"
    )

    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose logging")

    args = parser.parse_args()

    setup_logging(args.verbose)

    # Handle vocal stems logic
    vocal_stems_only = args.vocal_only and not args.no_vocal_only

    try:
        accent_counts = process_jsonl_file(
            input_file=args.input,
            output_file=args.output,
            target_accents=args.accents,
            vocal_stems_only=vocal_stems_only,
            max_per_accent=args.max_per_accent,
            limit=args.limit,
        )

        print(f"\n✅ Accent detection complete!")
        print(f"📁 Results saved to: {args.output}")
        print(f"📊 Found {sum(accent_counts.values())} total records")

        if accent_counts:
            print("\n🌍 Accent distribution:")
            for accent, count in sorted(accent_counts.items()):
                print(f"   {accent.replace('_', ' ').title()}: {count} records")
        else:
            print("⚠️  No records found with target accents")

        return 0

    except Exception as e:
        logging.error(f"Script failed: {e}")
        return 1


if __name__ == "__main__":
    exit(main())