#!/usr/bin/env python3
"""
Create audio list JSONL file from pickle file for crow_t1 dataset.

This script:
1. Reads the pickle file containing preference data
2. Extracts ALL clip IDs (main clips + condition/edited clips from metadata)
3. For each ID, checks if audio file exists in the audio directory (.m4a, .opus, .mp3)
4. Outputs a JSONL file with entries: {"id": "clip_id", "audio_path": "/full/path/to/audio.ext"}
5. Reports statistics: total IDs, found files, missing files

Usage:
    python create_crow_t1_audio_list.py \
        --input-pkl /home/tony/Data/Preference/crow_t1/interesting_clips_crow_t1_20251030_full_slice.pkl \
        --audio-dir /app2/suno/data/dpo/audios/crow \
        --output /home/tony/Data/Preference/crow_t1/crow_t1_audio_list.jsonl
"""

import argparse
import json
import os
import pandas as pd
from typing import Dict, List, Optional, Set, Tuple

# Try importing tqdm, fall back to no-op if not available
try:
    from tqdm import tqdm
except ImportError:

    def tqdm(iterable, **kwargs):
        return iterable


# Metadata fields that contain clip IDs (same as download_audio_from_s3.py)
SINGLE_CLIP_FIELDS = [
    "artist_clip_id",
    "cover_clip_id",
    "stem_clip_id",
    "infill_clip_id",
    "underpainting_clip_id",
    "overpainting_clip_id",
    "edited_clip_id",
]
MULTI_CLIP_FIELDS = ["playlist_clip_ids", "sample_clip_ids"]


def extract_clip_ids_from_metadata(metadata: Dict) -> List[str]:
    """
    Extract all clip IDs from metadata dictionary.

    Args:
        metadata: Metadata dictionary from dataframe row

    Returns:
        List of unique clip IDs found in metadata
    """
    clip_ids = []

    # Extract from single clip fields
    for field in SINGLE_CLIP_FIELDS:
        if field in metadata and metadata[field]:
            clip_ids.append(metadata[field])

    # Extract from multi-clip fields
    for field in MULTI_CLIP_FIELDS:
        if field in metadata and metadata[field]:
            field_value = metadata[field]
            if isinstance(field_value, list):
                clip_ids.extend(field_value)

    return sorted(list(set(clip_ids)))


def get_row_metadata(row: pd.Series) -> Dict:
    """
    Safely extract metadata from a row.

    Args:
        row: Pandas Series representing a dataframe row

    Returns:
        Metadata dictionary or empty dict if not found
    """
    metadata = row.get("metadata", {})
    return metadata if isinstance(metadata, dict) else {}


def extract_all_clip_ids_from_dataframe(df: pd.DataFrame) -> Set[str]:
    """
    Extract all unique clip IDs from the dataframe.

    This includes:
    - Main clip IDs (s3_id column)
    - All condition/edited clip IDs from metadata

    Args:
        df: Pandas dataframe with preference data

    Returns:
        Set of all unique clip IDs
    """
    all_clip_ids = set()

    print(f"Extracting clip IDs from {len(df)} rows...")
    for idx, row in tqdm(df.iterrows(), total=len(df), desc="Extracting clip IDs"):
        # Add main clip ID
        if "s3_id" in row and row["s3_id"]:
            all_clip_ids.add(row["s3_id"])

        # Add condition/edited clip IDs from metadata
        metadata = get_row_metadata(row)
        condition_ids = extract_clip_ids_from_metadata(metadata)
        all_clip_ids.update(condition_ids)

    return all_clip_ids


def find_audio_file(clip_id: str, audio_dir: str) -> Optional[str]:
    """
    Find audio file for given clip ID in audio directory.

    Checks for .m4a, .opus, and .mp3 extensions.

    Args:
        clip_id: The clip ID to search for
        audio_dir: Directory containing audio files

    Returns:
        Full path to audio file if found, None otherwise
    """
    for ext in [".m4a", ".opus", ".mp3"]:
        filepath = os.path.join(audio_dir, f"{clip_id}{ext}")
        if os.path.exists(filepath):
            return filepath
    return None


def create_audio_list(
    input_pkl: str, audio_dir: str, output_jsonl: str, verbose: bool = False
) -> Tuple[int, int, int]:
    """
    Create audio list JSONL from pickle file.

    Args:
        input_pkl: Path to pickle file with preference data
        audio_dir: Directory containing downloaded audio files
        output_jsonl: Path to output JSONL file
        verbose: Whether to print detailed progress

    Returns:
        Tuple of (total_ids, found_files, missing_files)
    """
    # Load the pickle file
    print(f"Loading preference data from: {input_pkl}")
    df = pd.read_pickle(input_pkl)
    print(f"Loaded dataframe shape: {df.shape}")

    # Extract all clip IDs from dataframe (main + condition clips)
    all_clip_ids = extract_all_clip_ids_from_dataframe(df)
    clip_ids_list = sorted(list(all_clip_ids))

    print(f"\nTotal unique clip IDs extracted: {len(clip_ids_list):,}")
    print(f"  - Main clip IDs (rows): {len(df):,}")
    print(f"  - Total including conditions: {len(clip_ids_list):,}")

    # Find audio files for each ID
    found_entries = []
    missing_ids = []

    print(f"\nSearching for audio files in: {audio_dir}")
    for clip_id in tqdm(clip_ids_list, desc="Finding audio files"):
        audio_path = find_audio_file(clip_id, audio_dir)

        if audio_path:
            found_entries.append({"id": clip_id, "audio_path": audio_path})
        else:
            missing_ids.append(clip_id)
            if verbose:
                print(f"Warning: Audio file not found for ID: {clip_id}")

    # Write output JSONL
    print(f"\nWriting audio list to: {output_jsonl}")
    os.makedirs(os.path.dirname(output_jsonl), exist_ok=True)

    with open(output_jsonl, "w") as f:
        for entry in found_entries:
            f.write(json.dumps(entry) + "\n")

    # Report statistics
    total_ids = len(clip_ids_list)
    found_files = len(found_entries)
    missing_files = len(missing_ids)

    print(f"\n{'='*60}")
    print("Audio List Creation Summary:")
    print(f"{'='*60}")
    print(f"Total clip IDs:           {total_ids:>10,}")
    print(
        f"Audio files found:        {found_files:>10,} ({found_files/total_ids*100:.2f}%)"
    )
    print(
        f"Audio files missing:      {missing_files:>10,} ({missing_files/total_ids*100:.2f}%)"
    )
    print(f"{'='*60}")
    print(f"Output written to: {output_jsonl}")

    if missing_files > 0:
        missing_log = output_jsonl.replace(".jsonl", "_missing_ids.json")
        with open(missing_log, "w") as f:
            json.dump(missing_ids, f, indent=2)
        print(f"Missing IDs logged to: {missing_log}")

    return total_ids, found_files, missing_files


def main():
    parser = argparse.ArgumentParser(
        description="Create audio list JSONL from pickle file for crow_t1 dataset"
    )
    parser.add_argument(
        "--input-pkl",
        type=str,
        default="/home/tony/Data/Preference/crow_t1/interesting_clips_crow_t1_20251030_full_slice.pkl",
        help="Path to pickle file with preference data",
    )
    parser.add_argument(
        "--audio-dir",
        type=str,
        default="/app2/suno/data/dpo/audios/crow",
        help="Directory containing downloaded audio files",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="/home/tony/Data/Preference/crow_t1/crow_t1_audio_list.jsonl",
        help="Path to output JSONL file",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print detailed progress including missing files",
    )

    args = parser.parse_args()

    # Validate input files
    if not os.path.exists(args.input_pkl):
        print(f"Error: Input pickle file not found: {args.input_pkl}")
        return 1

    if not os.path.exists(args.audio_dir):
        print(f"Error: Audio directory not found: {args.audio_dir}")
        return 1

    # Create audio list
    total, found, missing = create_audio_list(
        args.input_pkl, args.audio_dir, args.output, args.verbose
    )

    # Exit with error code if too many files are missing
    if missing > total * 0.01:  # More than 1% missing
        print(f"\nWarning: {missing/total*100:.2f}% of files are missing!")
        return 1

    return 0


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