#!/usr/bin/env python3
"""
Verify that all records in a vocal stems JSONL file actually contain vocal stems.
Checks that each record has "stems" field with at least one stem containing "vocal" or "vox" (case insensitive).
"""

import json
import argparse
from pathlib import Path
from tqdm import tqdm


def has_vocal_stem(record: dict) -> tuple[bool, list[str]]:
    """
    Check if a record has vocal stems.

    Returns:
        Tuple of (has_vocal_stem: bool, vocal_stems: list[str])
    """
    if "stems" not in record:
        return False, []

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

    # Check for vocal-related stem keys (case insensitive)
    vocal_keywords = ["vocal", "vox"]
    vocal_stems = []

    for stem_name in stems.keys():
        stem_name_lower = stem_name.lower()
        if any(keyword in stem_name_lower for keyword in vocal_keywords):
            vocal_stems.append(stem_name)

    return len(vocal_stems) > 0, vocal_stems


def verify_vocal_stems_file(input_path: Path, show_examples: bool = True, max_examples: int = 10):
    """
    Verify all records in the file have vocal stems.
    """
    total_records = 0
    valid_records = 0
    invalid_records = 0
    invalid_examples = []

    print(f"Verifying vocal stems in: {input_path}")
    print("=" * 60)

    with open(input_path, "r") as infile:
        for line_num, line in enumerate(tqdm(infile, desc="Verifying records", unit=" records"), 1):
            total_records += 1

            try:
                record = json.loads(line.strip())
                has_vocal, vocal_stems = has_vocal_stem(record)

                if has_vocal:
                    valid_records += 1

                    # Show first few examples of valid records
                    if show_examples and valid_records <= max_examples:
                        record_id = record.get("id", f"line_{line_num}")
                        print(
                            f"✅ Valid record {valid_records}: ID={record_id}, vocal stems: {vocal_stems}"
                        )

                else:
                    invalid_records += 1
                    record_id = record.get("id", f"line_{line_num}")
                    stems_list = (
                        list(record.get("stems", {}).keys()) if "stems" in record else ["NO_STEMS_FIELD"]
                    )

                    invalid_example = {
                        "line_number": line_num,
                        "record_id": record_id,
                        "stems": stems_list,
                    }
                    invalid_examples.append(invalid_example)

                    # Show invalid records immediately
                    print(f"❌ INVALID record {invalid_records}: Line {line_num}, ID={record_id}")
                    print(f"   Available stems: {stems_list}")

            except json.JSONDecodeError as e:
                invalid_records += 1
                print(f"❌ JSON decode error at line {line_num}: {e}")
                continue

    # Print summary
    print(f"\n{'='*60}")
    print(f"VERIFICATION RESULTS")
    print(f"{'='*60}")
    print(f"Total records processed: {total_records:,}")
    print(f"Valid vocal records: {valid_records:,}")
    print(f"Invalid records: {invalid_records:,}")

    if total_records > 0:
        valid_percentage = (valid_records / total_records) * 100
        print(f"Valid percentage: {valid_percentage:.2f}%")

    if invalid_records > 0:
        print(f"\n❌ VERIFICATION FAILED: {invalid_records} invalid records found")

        print(f"\nFirst {min(len(invalid_examples), 20)} invalid records:")
        for i, example in enumerate(invalid_examples[:20]):
            print(f"  {i+1}. Line {example['line_number']}: ID={example['record_id']}")
            print(f"     Stems: {example['stems']}")

        if len(invalid_examples) > 20:
            print(f"     ... and {len(invalid_examples) - 20} more invalid records")

        return False
    else:
        print(f"\n✅ VERIFICATION PASSED: All {valid_records:,} records have valid vocal stems")
        return True


def main():
    parser = argparse.ArgumentParser(
        description="Verify that all records in vocal stems JSONL file have vocal stems"
    )
    parser.add_argument(
        "--input-path",
        type=str,
        default="/home/vibert/data/voice_designer/metas_v6_tr_vocal_stems.jsonl",
        help="Path to vocal stems JSONL file to verify",
    )
    parser.add_argument(
        "--no-examples", action="store_true", help="Don't show examples of valid records"
    )
    parser.add_argument(
        "--max-examples", type=int, default=10, help="Maximum number of valid record examples to show"
    )

    args = parser.parse_args()

    input_path = Path(args.input_path)

    if not input_path.exists():
        print(f"Error: Input file {input_path} does not exist")
        return 1

    try:
        is_valid = verify_vocal_stems_file(
            input_path, show_examples=not args.no_examples, max_examples=args.max_examples
        )

        return 0 if is_valid else 1

    except Exception as e:
        print(f"Error during verification: {e}")
        return 1


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