#!/usr/bin/env python3
"""
Sample and display records that contain vocal stems from metas_v6_tr.jsonl
"""

import json
import random
import argparse
from pathlib import Path
from typing import Dict, Any, List
from tqdm import tqdm
from file_read_backwards import FileReadBackwards


def has_vocal_stem(record: Dict[str, Any]) -> bool:
    """
    Check if a record has vocal stems
    """
    if "stems" not in record:
        return False

    stems = record["stems"]
    if not isinstance(stems, dict):
        return False

    # Check for vocal-related stem keys
    vocal_keys = ["vocal", "voice", "vox", "backing_vocal"]
    return any(any(vocal_key in key.lower() for vocal_key in vocal_keys) for key in stems.keys())


def find_vocal_stem_records(
    filepath: Path,
    max_samples: int = 10,
    random_seed: int = 42,
    skip_lines: int = 0,
    reverse: bool = False,
) -> List[Dict[str, Any]]:
    """
    Find records with vocal stems

    Args:
        filepath: Path to JSONL file
        max_samples: Maximum number of vocal records to find
        random_seed: Random seed for reproducibility
        skip_lines: Skip N lines between each processed line (0 = process every line)
        reverse: If True, traverse the file backwards
    """
    random.seed(random_seed)
    vocal_records = []

    if reverse:
        # Read file backwards using FileReadBackwards
        with FileReadBackwards(filepath, encoding="utf-8") as f:
            line_count = 0
            for line in tqdm(f, desc="Searching backwards for vocal stem records", unit=" lines"):
                # Skip lines if skip_lines > 0
                if skip_lines > 0 and line_count % (skip_lines + 1) != 0:
                    line_count += 1
                    continue

                try:
                    record = json.loads(line.strip())
                    # Add reverse line position info
                    record["_line_number_from_end"] = line_count

                    if has_vocal_stem(record):
                        vocal_records.append(record)
                        if len(vocal_records) >= max_samples:
                            break
                except json.JSONDecodeError:
                    pass

                line_count += 1
    else:
        # Forward traversal with line counting
        with open(filepath, "r") as f:
            line_count = 0
            for line in tqdm(f, desc="Searching for vocal stem records", unit=" lines"):
                # Skip lines if skip_lines > 0
                if skip_lines > 0 and line_count % (skip_lines + 1) != 0:
                    line_count += 1
                    continue

                try:
                    record = json.loads(line.strip())
                    # Add line number to record
                    record["_line_number"] = line_count

                    if has_vocal_stem(record):
                        vocal_records.append(record)
                        if len(vocal_records) >= max_samples:
                            break
                except json.JSONDecodeError:
                    pass

                line_count += 1

    return vocal_records


def display_vocal_record(record: Dict[str, Any], index: int):
    """
    Display information about a vocal stem record
    """
    print(f"\n{'='*60}")
    print(f"VOCAL RECORD #{index + 1}")
    print(f"{'='*60}")

    # Line position info
    line_info = ""
    if "_line_number" in record:
        line_info = f" (Line: {record['_line_number']})"
    elif "_line_number_from_end" in record:
        line_info = f" (Line from end: {record['_line_number_from_end']})"

    # Basic info
    print(f"ID: {record.get('id', 'N/A')}{line_info}")
    print(f"Duration: {record.get('duration_s', 'N/A')} seconds")
    print(f"Language: {record.get('lang', 'N/A')}")

    # Stems info
    stems = record.get("stems", {})
    print(f"\nSTEMS ({len(stems)} total):")
    for stem_name, stem_path in stems.items():
        is_vocal = any(
            vocal_key in stem_name.lower() for vocal_key in ["vocal", "voice", "vox", "backing"]
        )
        marker = "🎤" if is_vocal else "🎵"
        print(f"  {marker} {stem_name}: {Path(stem_path).name}")

    # Text/lyrics (truncated)
    text = record.get("text", "")
    if text:
        text_preview = text[:200] + "..." if len(text) > 200 else text
        print(f"\nTEXT/LYRICS:")
        print(f"  {text_preview}")

    # Tags
    tags = record.get("tags", [])
    if tags:
        print(f"\nTAGS ({len(tags)}):")
        # Show first 10 tags
        displayed_tags = tags[:10]
        print(f"  {', '.join(str(tag) for tag in displayed_tags)}")
        if len(tags) > 10:
            print(f"  ... and {len(tags) - 10} more")

    # Stems captions if available
    stems_captions = record.get("stems_captions", {})
    if stems_captions:
        print(f"\nSTEM CAPTIONS:")
        for stem_name, captions in stems_captions.items():
            is_vocal_stem = any(
                vocal_key in stem_name.lower() for vocal_key in ["vocal", "voice", "vox", "backing"]
            )
            if is_vocal_stem and isinstance(captions, list) and captions:
                print(f"  🎤 {stem_name}:")
                for caption_data in captions[:2]:  # Show first 2 captions
                    if isinstance(caption_data, dict):
                        caption_text = caption_data.get("caption", "N/A")
                        prompt_type = caption_data.get("prompt_type", "unknown")
                        print(f"    [{prompt_type}]: {caption_text[:150]}...")


def main():
    parser = argparse.ArgumentParser(description="Sample vocal stem records from v6 metadata")
    parser.add_argument(
        "--data-path",
        type=str,
        default="/app2/suno/data/auk_v0/metas_v6_tr.jsonl",
        help="Path to the JSONL file",
    )
    parser.add_argument(
        "--max-samples", type=int, default=10, help="Maximum number of vocal records to find"
    )
    parser.add_argument(
        "--output-file", type=str, default=None, help="Optional: save records to JSON file"
    )
    parser.add_argument("--random-seed", type=int, default=42, help="Random seed for reproducibility")
    parser.add_argument(
        "--skip-lines",
        type=int,
        default=0,
        help="Skip N lines between each processed line (0 = process every line, 1000 = process every 1001st line)",
    )
    parser.add_argument(
        "--reverse", action="store_true", help="Traverse the file backwards from end to beginning"
    )

    args = parser.parse_args()

    data_path = Path(args.data_path)

    print(f"Searching for vocal stem records in {data_path}")
    print(f"Looking for up to {args.max_samples} records...")

    if args.skip_lines > 0:
        print(
            f"Processing every {args.skip_lines + 1} lines (skipping {args.skip_lines} lines between each)"
        )

    if args.reverse:
        print("🔄 Traversing file backwards from end to beginning")

    vocal_records = find_vocal_stem_records(
        data_path, args.max_samples, args.random_seed, args.skip_lines, args.reverse
    )

    print(f"\n🎤 Found {len(vocal_records)} vocal stem records!")

    if not vocal_records:
        print("No vocal stem records found.")
        return

    # Display each record
    for i, record in enumerate(vocal_records):
        display_vocal_record(record, i)

    # Save to file if requested
    if args.output_file:
        output_path = Path(args.output_file)
        with open(output_path, "w") as f:
            json.dump(vocal_records, f, indent=2)
        print(f"\n💾 Saved {len(vocal_records)} records to {output_path}")

    # Summary stats
    print(f"\n{'='*60}")
    print("SUMMARY STATISTICS")
    print(f"{'='*60}")

    total_stems = sum(len(record.get("stems", {})) for record in vocal_records)
    avg_stems = total_stems / len(vocal_records) if vocal_records else 0

    languages = [record.get("lang") for record in vocal_records if record.get("lang")]
    lang_counts = {}
    for lang in languages:
        lang_counts[lang] = lang_counts.get(lang, 0) + 1

    print(f"Total vocal records found: {len(vocal_records)}")
    print(f"Average stems per record: {avg_stems:.1f}")
    print(f"Languages: {dict(sorted(lang_counts.items(), key=lambda x: x[1], reverse=True))}")

    # Count specific vocal stem types
    vocal_stem_types = {}
    for record in vocal_records:
        stems = record.get("stems", {})
        for stem_name in stems.keys():
            if any(vocal_key in stem_name.lower() for vocal_key in ["vocal", "voice", "vox", "backing"]):
                vocal_stem_types[stem_name] = vocal_stem_types.get(stem_name, 0) + 1

    if vocal_stem_types:
        print(
            f"Vocal stem types: {dict(sorted(vocal_stem_types.items(), key=lambda x: x[1], reverse=True))}"
        )


if __name__ == "__main__":
    main()