#!/usr/bin/env python3
"""
Filter validation metadata based on quality constraints.

Filters applied (in order):
1. Duration between 30s and 8 minutes (480s)
2. Remove tags containing ":"
3. Remove entries with empty tags after cleaning
4. For English entries with text, keep only those with text_aligned
   (Non-English songs can have text without text_aligned)
"""

import json
import os
from typing import Dict, List, Tuple

try:
    from tqdm import tqdm
except ImportError:
    def tqdm(iterable, *args, **kwargs):
        return iterable


def filter_metadata(
    input_path: str,
    output_path: str,
    min_duration: float = 30.0,
    max_duration: float = 480.0,
) -> Dict[str, int]:
    """
    Filter metadata based on quality constraints.
    
    Args:
        input_path: Path to input metadata JSONL file
        output_path: Path to output filtered metadata JSONL file
        min_duration: Minimum duration in seconds (default: 30s)
        max_duration: Maximum duration in seconds (default: 480s = 8 mins)
        
    Returns:
        Dictionary with filtering statistics
    """
    print(f"Filtering metadata from: {input_path}")
    print(f"Output will be written to: {output_path}")
    print(f"\nFilter criteria:")
    print(f"  - Duration: {min_duration}s to {max_duration}s")
    print(f"  - Remove tags with ':'")
    print(f"  - Remove entries with empty tags after cleaning")
    print(f"  - For English songs with text: require text_aligned")
    print(f"    (Non-English songs can have text without text_aligned)")
    print()
    
    # Count total lines
    print("Counting lines...")
    with open(input_path, "r") as f:
        total_lines = sum(1 for _ in f)
    print(f"Total entries: {total_lines:,}\n")
    
    # Statistics tracking
    stats = {
        "total": total_lines,
        "removed_by_duration": 0,
        "removed_by_empty_tags": 0,
        "removed_by_text_aligned": 0,
        "kept": 0,
    }
    
    # Additional tracking for detailed reporting
    duration_too_short = 0
    duration_too_long = 0
    no_text_aligned = 0
    no_text = 0  # for reference
    
    # Create output directory if needed
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    # Process entries
    print("Processing and filtering entries...")
    with open(input_path, "r") as infile, open(output_path, "w") as outfile:
        for line in tqdm(infile, total=total_lines, desc="Filtering"):
            line = line.strip()
            if not line:
                continue
            
            try:
                meta = json.loads(line)
                
                # Filter 1: Duration check
                duration = meta.get("duration_s")
                if duration is None or duration < min_duration:
                    stats["removed_by_duration"] += 1
                    if duration is not None and duration < min_duration:
                        duration_too_short += 1
                    continue
                
                if duration > max_duration:
                    stats["removed_by_duration"] += 1
                    duration_too_long += 1
                    continue
                
                # Filter 2: Clean tags (remove tags with ":")
                tags = meta.get("tags", [])
                if tags:
                    cleaned_tags = [tag for tag in tags if ":" not in str(tag)]
                    meta["tags"] = cleaned_tags
                    
                    # Filter 3: Remove entries with empty tags after cleaning
                    if not cleaned_tags:
                        stats["removed_by_empty_tags"] += 1
                        continue
                else:
                    # No tags at all - remove
                    stats["removed_by_empty_tags"] += 1
                    continue
                
                # Filter 4: For English entries with text, keep only those with text_aligned
                # Non-English songs can have text without text_aligned
                text = meta.get("text")
                lang = meta.get("lang")
                lang = lang.lower() if lang else ""
                
                if text:
                    text_aligned = meta.get("text_aligned")
                    # Only enforce text_aligned requirement for English songs
                    if lang == "en" and not text_aligned:
                        stats["removed_by_text_aligned"] += 1
                        no_text_aligned += 1
                        continue
                else:
                    # Track entries without text (but these pass the filter)
                    no_text += 1
                
                # All filters passed - write to output
                outfile.write(json.dumps(meta) + "\n")
                stats["kept"] += 1
                
            except json.JSONDecodeError as e:
                print(f"Warning: Skipping malformed JSON line: {e}")
                continue
    
    # Store additional details for reporting
    stats["duration_too_short"] = duration_too_short
    stats["duration_too_long"] = duration_too_long
    stats["no_text_aligned"] = no_text_aligned
    stats["no_text"] = no_text
    
    return stats


def print_filter_report(stats: Dict[str, int]) -> None:
    """
    Print detailed filtering report.
    
    Args:
        stats: Statistics dictionary from filtering
    """
    print("\n" + "=" * 80)
    print("FILTERING REPORT")
    print("=" * 80)
    print()
    
    total = stats["total"]
    
    print(f"Starting entries:                {total:>10,}")
    print()
    
    # Filter 1: Duration
    removed_duration = stats["removed_by_duration"]
    pct_duration = 100 * removed_duration / total
    print(f"Filter 1 - Duration (30s - 480s):")
    print(f"  Removed:                       {removed_duration:>10,}  ({pct_duration:>5.2f}%)")
    print(f"    - Too short (< 30s):         {stats['duration_too_short']:>10,}")
    print(f"    - Too long (> 480s):         {stats['duration_too_long']:>10,}")
    remaining_after_duration = total - removed_duration
    print(f"  Remaining after filter:        {remaining_after_duration:>10,}")
    print()
    
    # Filter 2 & 3: Tags
    removed_tags = stats["removed_by_empty_tags"]
    pct_tags = 100 * removed_tags / total
    print(f"Filter 2 & 3 - Clean tags & remove empty:")
    print(f"  Removed (empty after cleaning):{removed_tags:>10,}  ({pct_tags:>5.2f}%)")
    remaining_after_tags = remaining_after_duration - removed_tags
    print(f"  Remaining after filter:        {remaining_after_tags:>10,}")
    print()
    
    # Filter 4: Text aligned
    removed_text = stats["removed_by_text_aligned"]
    pct_text = 100 * removed_text / total
    print(f"Filter 4 - Text with text_aligned (English only):")
    print(f"  Removed (English + text but no alignment): {removed_text:>10,}  ({pct_text:>5.2f}%)")
    print(f"  Note: Entries without text:    {stats['no_text']:>10,}  (these pass through)")
    print(f"  Note: Non-English songs:       (pass through even without alignment)")
    remaining_after_text = remaining_after_tags - removed_text
    print(f"  Remaining after filter:        {remaining_after_text:>10,}")
    print()
    
    # Final stats
    print("-" * 80)
    kept = stats["kept"]
    pct_kept = 100 * kept / total
    total_removed = total - kept
    pct_removed = 100 * total_removed / total
    
    print(f"FINAL RESULTS:")
    print(f"  Total entries kept:            {kept:>10,}  ({pct_kept:>5.2f}%)")
    print(f"  Total entries removed:         {total_removed:>10,}  ({pct_removed:>5.2f}%)")
    print()
    
    # Breakdown of removals
    print("Removal breakdown by filter:")
    print(f"  Duration filter:               {removed_duration:>10,}  ({100*removed_duration/total_removed:>5.2f}% of removed)")
    print(f"  Empty tags filter:             {removed_tags:>10,}  ({100*removed_tags/total_removed:>5.2f}% of removed)")
    print(f"  Text aligned filter:           {removed_text:>10,}  ({100*removed_text/total_removed:>5.2f}% of removed)")
    
    print("=" * 80)


def main():
    """Main function to run filtering."""
    # Configuration
    input_path = os.path.expanduser("~/Data/Preference/RealGen/metas_v5_val_filtered.jsonl")
    output_dir = "/home/tony/Work/tony/RealGen"
    output_filename = "metas_v5_val_filtered_clean.jsonl"
    output_path = os.path.join(output_dir, output_filename)
    
    # Filter parameters
    min_duration = 30.0  # 30 seconds
    max_duration = 480.0  # 8 minutes = 480 seconds
    
    print("🎵 Validation Metadata Filtering Script")
    print(f"Input:  {input_path}")
    print(f"Output: {output_path}")
    print()
    
    # Check if input file exists
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"Input file not found: {input_path}")
    
    # Filter metadata
    stats = filter_metadata(input_path, output_path, min_duration, max_duration)
    
    # Print report
    print_filter_report(stats)
    
    print("\n✅ Filtering completed successfully!")
    print(f"📁 Output file: {output_path}")
    
    # Verify output file
    if os.path.exists(output_path):
        file_size = os.path.getsize(output_path)
        print(f"📊 Output file size: {file_size / (1024**3):.2f} GB")


if __name__ == "__main__":
    main()