#!/usr/bin/env python3
"""
Extract records with vocal stems from metas JSONL files.
Filters records that have "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) -> bool:
    """
    Check if a record has vocal stems.

    Returns True if:
    1. Record has "stems" field
    2. At least one stem key contains "vocal" or "vox" (case insensitive)
    """
    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"]

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

    return False


def extract_vocal_records(input_path: Path, output_path: Path, verbose: bool = True):
    """
    Extract all records with vocal stems from input JSONL to output JSONL.
    """
    vocal_count = 0
    total_count = 0

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

    with open(input_path, "r") as infile, open(output_path, "w") as outfile:
        for line in tqdm(infile, desc="Processing records", unit=" records"):
            total_count += 1

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

                if has_vocal_stem(record):
                    vocal_count += 1
                    # Write the record to output file
                    outfile.write(json.dumps(record) + "\n")

                    if verbose and vocal_count <= 5:
                        # Show first few vocal records found
                        stems = record.get("stems", {})
                        vocal_stems = [
                            name
                            for name in stems.keys()
                            if any(kw in name.lower() for kw in ["vocal", "vox"])
                        ]
                        print(
                            f"Found vocal record {vocal_count}: ID={record.get('id', 'N/A')}, "
                            f"vocal stems: {vocal_stems}"
                        )

            except json.JSONDecodeError as e:
                if verbose:
                    print(f"Warning: Skipping malformed JSON at line {total_count}: {e}")
                continue

    return total_count, vocal_count


def main():
    parser = argparse.ArgumentParser(description="Extract records with vocal stems from JSONL files")
    parser.add_argument(
        "--input-path",
        type=str,
        default="/app2/suno/data/auk_v0/metas_v6_tr.jsonl",
        help="Path to input JSONL file",
    )
    parser.add_argument(
        "--output-path",
        type=str,
        default="/home/vibert/data/voice_designer/metas_v6_tr_vocal_stems.jsonl",
        help="Path to output JSONL file",
    )
    parser.add_argument("--quiet", action="store_true", help="Suppress verbose output")

    args = parser.parse_args()

    input_path = Path(args.input_path)
    output_path = Path(args.output_path)

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

    print(f"Extracting vocal stem records...")
    print(f"Input:  {input_path}")
    print(f"Output: {output_path}")
    print()

    try:
        total_count, vocal_count = extract_vocal_records(input_path, output_path, verbose=not args.quiet)

        print(f"\n{'='*60}")
        print(f"EXTRACTION COMPLETE")
        print(f"{'='*60}")
        print(f"Total records processed: {total_count:,}")
        print(f"Records with vocal stems: {vocal_count:,}")
        print(f"Vocal percentage: {vocal_count / total_count * 100:.2f}%")
        print(f"Output saved to: {output_path}")

        return 0

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


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