#!/usr/bin/env python3
"""
Script to download audio files from S3 based on a pickle file containing audio metadata.
Processes rows in pairs (even/odd indices) where pairs should have matching edited clip IDs.

Usage:
    python download_audio_from_s3.py <pickle_file_path> [--dest-folder <folder>]

Example:
# 30b t1 v24: 130026
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/30b_v0/interesting_clips_v4_t_1_20240808_v23_slice.pkl \
    --dest-folder /app2/suno/data/dpo/audios/30b \
    --num-workers 32

# Auk t0: (may not be the final version) 156310
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/auk_t0/fully_merged_auk_t0_final_slice.pkl \
    --dest-folder /app2/suno/data/dpo/audios/auk_t0 \
    --num-workers 16

# 30b t6 v35: 107104
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/30b_v6/interesting_clips_v4_h_t_6_20250222_full_v34.pkl \
    --dest-folder /app2/suno/data/dpo/audios/v4 \
    --num-workers 32

# 13b s32 v34: 449352
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/13b_v32/interesting_clips_v4_h_s_32_20250330_full_long_bluejay_r2.pkl\
    --dest-folder /app2/suno/data/dpo/audios/v4 \
    --num-workers 32

# Auk mixed 13b: 153414
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/13b_v32/interesting_clips_v4_h_s_32_20250411_full_long_slice.pkl \
    --dest-folder /app2/suno/data/dpo/audios/v4 \
    --num-workers 32

# Auk mixed 30b: 253726
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/30b_v6/interesting_clips_v4_h_t_6_20250411_full_long_slice.pkl \
    --dest-folder /app2/suno/data/dpo/audios/v4 \
    --num-workers 32

# Auk mixed 2 13b: 201960
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/13b_v32/interesting_clips_v4_h_s_32_20250412_20250501_full_long_slice.pkl \
    --dest-folder /app2/suno/data/dpo/audios/v4 \
    --num-workers 32

# Auk mixed 2 30b: 104826
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/30b_v6/interesting_clips_v4_h_t_6_20250426_20250501_full_long_slice.pkl \
    --dest-folder /app2/suno/data/dpo/audios/v4 \
    --num-workers 32

# Auk t1 v6: 453768
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/auk_t1/interesting_clips_auk_t1_20250520_v6_slice.pkl \
    --dest-folder /app2/suno/data/dpo/audios/auk_t1 \
    --num-workers 32

# Auk t1 v7: 401248
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/auk_t1/interesting_clips_auk_t1_20250520_v7_slice.pkl \
    --dest-folder /app2/suno/data/dpo/audios/auk_t1 \
    --num-workers 32

# Auk t1 v19: 676032
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/auk_t1/interesting_clips_auk_t1_20250602_v19_slice.pkl \
    --dest-folder /app2/suno/data/dpo/audios/auk_t1 \
    --num-workers 32

# Auk t1 v29: 703108 (note that we can't reproduce this exactly...)
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/auk_t1/interesting_clips_auk_t1_20250613_v29_slice.pkl\
    --dest-folder /app2/suno/data/dpo/audios/auk_t1 \
    --num-workers 32

# Crow t1 & t2: 2733682
python /home/tony/Work/tony/Preference/download_audio_from_s3.py \
    /home/tony/Data/Preference/crow_t1/interesting_clips_crow_t1_20251030_full_slice.pkl\
    --dest-folder /app2/suno/data/dpo/audios/crow \
    --num-workers 32

Note: The dataframe is expected to have paired rows where:
- Even and odd rows form pairs (rows 0-1, 2-3, 4-5, etc.)
- Pairs should have different clip_ids but same edited clip IDs
- Both rows in a pair must download successfully for the pair to be marked successful
"""

import argparse
import boto3
import os
import pandas as pd
from typing import Optional, Tuple, Dict, List
from tqdm import tqdm
import sys
import json
from multiprocessing import Pool, cpu_count
from dataclasses import dataclass

# Constants
S3_BUCKET = "suno-data-uploads"
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"]
# All metadata fields we actually need (for memory optimization)
ALL_METADATA_FIELDS = SINGLE_CLIP_FIELDS + MULTI_CLIP_FIELDS
DOWNLOAD_ATTEMPTS = [
    ("m4a", "studio/uploads"),
    ("m4a", "studio/deleted"),
    ("opus", "studio/uploads"),
    ("opus", "studio/deleted"),
    ("mp3", "studio/uploads"),
    ("mp3", "studio/deleted"),
]

# Global S3 client for worker processes
s3_client = None


def init_worker():
    """Initialize global S3 client in worker process"""
    global s3_client
    try:
        s3_client = boto3.client("s3")
    except Exception as e:
        print(f"Error initializing S3 client in worker: {e}", file=sys.stderr)


@dataclass
class ProcessingResult:
    """Result of processing a row pair"""

    clip_ids: List[str]
    success: bool
    pair_index: int
    edited_ids_match: bool
    files_skipped: int = 0
    files_downloaded: int = 0


@dataclass
class DownloadSummary:
    """Summary of download operation"""

    total_pairs: int
    successful_pairs: int
    failed_pairs: int
    mismatched_pairs: int
    successful_clip_ids: List[str]
    total_files_skipped: int = 0
    total_files_downloaded: int = 0


def check_file_exists(s3_id: str, local_folder: str) -> bool:
    """Check if audio file already exists locally in any supported format"""
    # Check for m4a, opus, and mp3 formats
    return (
        os.path.exists(f"{local_folder}/{s3_id}.m4a")
        or os.path.exists(f"{local_folder}/{s3_id}.opus")
        or os.path.exists(f"{local_folder}/{s3_id}.mp3")
    )


def download_audio_to_local_from_s3(s3_id: str, local_folder: str, client) -> bool:
    """
    Download audio file from S3 to local folder with atomic writes.

    Downloads to a temporary file first, then renames to prevent partial files
    from being treated as complete downloads if the process is interrupted.

    Args:
        s3_id: The S3 ID of the audio file
        local_folder: Local folder to save the file
        client: Boto3 S3 client instance

    Returns:
        True if download successful or file already exists, False otherwise
    """
    # Check if file already exists - if so, skip all download attempts
    if check_file_exists(s3_id, local_folder):
        return True

    if client is None:
        print(f"Error: S3 client not initialized for {s3_id}", file=sys.stderr)
        return False

    # Only try downloading if file doesn't exist
    for file_format, s3_path in DOWNLOAD_ATTEMPTS:
        temp_filepath = None
        try:
            local_filepath = f"{local_folder}/{s3_id}.{file_format}"
            temp_filepath = f"{local_folder}/{s3_id}.{file_format}.tmp"

            # Download to temp file first (atomic write)
            client.download_file(
                S3_BUCKET, f"{s3_path}/{s3_id}.{file_format}", temp_filepath
            )

            # Only rename to final name if download completed successfully
            os.rename(temp_filepath, local_filepath)
            return True
        except Exception:
            # Clean up temp file if download failed
            if temp_filepath and os.path.exists(temp_filepath):
                try:
                    os.remove(temp_filepath)
                except Exception:
                    pass
            continue

    # If all attempts failed
    # print(f"Failed to download {s3_id} from all locations", file=sys.stderr)
    return False


def extract_clip_ids_from_metadata(metadata: Dict) -> List[str]:
    """Extract all clip IDs from metadata dictionary"""
    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 download_clips(
    clip_ids: List[str], dest_folder: str, client, verbose: bool
) -> Tuple[List[bool], bool, int, int]:
    """
    Download multiple clips and return individual and overall success status

    Returns:
        Tuple of (list of individual download results, all successful, number of skipped files, number of actual downloads)
    """
    results = []
    skipped = 0
    downloaded = 0

    for clip_id in clip_ids:
        # Check if file exists before attempting download
        if check_file_exists(clip_id, dest_folder):
            results.append(True)
            skipped += 1
            continue

        success = download_audio_to_local_from_s3(clip_id, dest_folder, client)
        results.append(success)
        if success:
            downloaded += 1
        elif verbose:
            pass
            # print(f"Failed to download clip: {clip_id}")

    return results, all(results), skipped, downloaded


def process_row_pair(
    args_tuple: Tuple[int, str, str, List[str], List[str], str, bool],
) -> ProcessingResult:
    """
    Process a pair of rows using pre-extracted IDs.

    Args:
        args_tuple: Tuple of (pair_idx, clip_id1, clip_id2, edited_ids1, edited_ids2, dest_folder, verbose)

    Returns:
        ProcessingResult with download status for the pair
    """
    (
        pair_idx,
        clip_id1,
        clip_id2,
        edited_ids1,
        edited_ids2,
        dest_folder,
        verbose,
    ) = args_tuple

    global s3_client

    # Check if edited IDs match
    edited_ids_match = edited_ids1 == edited_ids2
    if not edited_ids_match and verbose:
        print(f"Warning: Pair {pair_idx} has mismatched edited clip IDs!")
        print(f"  Row {pair_idx * 2}: {edited_ids1}")
        print(f"  Row {pair_idx * 2 + 1}: {edited_ids2}")

    # Download main clips
    main_clips = [clip_id1, clip_id2]
    main_results, main_success, main_skipped, main_downloaded = download_clips(
        main_clips, dest_folder, s3_client, verbose
    )

    # Download edited clips (union of both sets)
    all_edited_ids = sorted(list(set(edited_ids1 + edited_ids2)))
    edited_success = True
    edited_skipped = 0
    edited_downloaded = 0
    if all_edited_ids:
        _, edited_success, edited_skipped, edited_downloaded = download_clips(
            all_edited_ids, dest_folder, s3_client, verbose
        )

    # Determine overall success
    pair_success = main_success and edited_success
    total_skipped = main_skipped + edited_skipped
    total_downloaded = main_downloaded + edited_downloaded

    # Collect ALL successfully downloaded clip IDs (main + edited)
    all_clip_ids = [clip_id1, clip_id2] + all_edited_ids

    if verbose and not pair_success:
        # Only print failure details to reduce spam
        print(
            f"Pair {pair_idx} failed: "
            f"main_success={main_success}, edited_success={edited_success}, "
            f"skipped={total_skipped}, downloaded={total_downloaded}"
        )

    return ProcessingResult(
        clip_ids=all_clip_ids,  # Include BOTH main and edited clips
        success=pair_success,
        pair_index=pair_idx,
        edited_ids_match=edited_ids_match,
        files_skipped=total_skipped,
        files_downloaded=total_downloaded,
    )


def clean_metadata_fields(metadata: Dict) -> Dict:
    """Clean metadata dict to only include needed fields and remove NaN values"""
    if not isinstance(metadata, dict):
        return {}

    cleaned = {}
    for key in ALL_METADATA_FIELDS:
        if key not in metadata:
            continue
        value = metadata[key]
        if isinstance(value, list):
            cleaned_list = [v for v in value if not pd.isna(v)]
            if cleaned_list:
                cleaned[key] = cleaned_list
        elif not pd.isna(value):
            cleaned[key] = value
    return cleaned


def create_row_pair_data(
    df: pd.DataFrame, dest_folder: str, verbose: bool
) -> List[Tuple[int, str, str, List[str], List[str], str, bool]]:
    """
    Create row pairs with pre-extracted IDs to minimize memory usage in workers.
    Extracts necessary data from DataFrame and returns lightweight list of tuples.

    Returns list of tuples: (pair_idx, clip_id1, clip_id2, edited_ids1, edited_ids2, dest_folder, verbose)
    """
    # Extract only the columns we need as numpy arrays
    s3_ids = df["s3_id"].values if "s3_id" in df.columns else [None] * len(df)
    ids = df["id"].values
    metadatas = df["metadata"].values if "metadata" in df.columns else [{}] * len(df)

    tasks = []
    # We iterate here to perform extraction once, allowing the heavy DF/metadata to be freed
    # before multiprocessing starts.
    for i in tqdm(range(0, len(df) - 1, 2), desc="Preparing tasks"):
        # Row 1
        s3_id1 = s3_ids[i]
        clip_id1 = s3_id1 if pd.notna(s3_id1) else ids[i]
        meta1 = metadatas[i]
        edited_ids1 = extract_clip_ids_from_metadata(
            meta1 if isinstance(meta1, dict) else {}
        )

        # Row 2
        s3_id2 = s3_ids[i + 1]
        clip_id2 = s3_id2 if pd.notna(s3_id2) else ids[i + 1]
        meta2 = metadatas[i + 1]
        edited_ids2 = extract_clip_ids_from_metadata(
            meta2 if isinstance(meta2, dict) else {}
        )

        tasks.append(
            (
                i // 2,
                clip_id1,
                clip_id2,
                edited_ids1,
                edited_ids2,
                dest_folder,
                verbose,
            )
        )

    return tasks


def print_summary(summary: DownloadSummary, total_rows: int, output_json: str):
    """Print download summary"""
    print(f"\n{'=' * 50}")
    print("Download Summary:")
    print(f"Total pairs processed: {summary.total_pairs}")
    print(f"Total rows processed: {total_rows - (total_rows % 2)}")
    print(f"Successful pairs: {summary.successful_pairs}")
    print(f"Failed pairs: {summary.failed_pairs}")
    print(f"Pairs with mismatched edited IDs: {summary.mismatched_pairs}")
    print(f"Total successful clip IDs: {len(summary.successful_clip_ids)}")
    print(f"Files already existed (skipped): {summary.total_files_skipped}")
    print(f"Files actually downloaded: {summary.total_files_downloaded}")
    print(f"Output JSON: {output_json}")
    print(f"{'=' * 50}")


def cleanup_temp_files(folder: str):
    """Remove any leftover .tmp files from interrupted downloads"""
    tmp_count = 0
    if not os.path.exists(folder):
        return
    for filename in os.listdir(folder):
        if filename.endswith(".tmp"):
            try:
                os.remove(os.path.join(folder, filename))
                tmp_count += 1
            except Exception:
                pass
    if tmp_count > 0:
        print(
            f"Cleaned up {tmp_count} temporary files from previous interrupted downloads"
        )


def validate_aws_credentials():
    """Validate AWS credentials are properly configured"""
    try:
        test_client = boto3.client("s3")
        test_client.list_buckets()
    except Exception as e:
        print(f"Error initializing S3 client: {str(e)}", file=sys.stderr)
        print("Please ensure AWS credentials are properly configured.", file=sys.stderr)
        sys.exit(1)


def get_output_json_path(pickle_file: str, output_json: Optional[str]) -> str:
    """Determine output JSON file path"""
    if output_json:
        return output_json

    base_name = os.path.splitext(os.path.basename(pickle_file))[0]
    pickle_dir = os.path.dirname(pickle_file)
    return os.path.join(pickle_dir, f"{base_name}_successful_downloads.json")


def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description="Download audio files from S3 based on a pickle file (processing paired rows)"
    )
    parser.add_argument(
        "pickle_file",
        type=str,
        help="Path to the pickle file containing audio metadata",
    )
    parser.add_argument(
        "--dest-folder",
        type=str,
        default="/app2/suno/data/dpo/audios/auk_t0",
        help="Destination folder for downloaded audio files (default: /app2/suno/data/dpo/audios/auk_t0)",
    )
    parser.add_argument(
        "--verbose", action="store_true", help="Print detailed progress for each file"
    )
    parser.add_argument(
        "--output-json",
        type=str,
        help="Path to output JSON file with successful downloads (default: <pickle_file_name>_successful_downloads.json)",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=None,
        help="Number of parallel workers (default: number of CPU cores)",
    )

    args = parser.parse_args()

    # Validate pickle file exists
    if not os.path.exists(args.pickle_file):
        print(f"Error: Pickle file not found: {args.pickle_file}", file=sys.stderr)
        sys.exit(1)

    # Set output paths
    args.output_json = get_output_json_path(args.pickle_file, args.output_json)

    # Create destination folder
    os.makedirs(args.dest_folder, exist_ok=True)

    # Clean up any leftover .tmp files from previous interrupted runs
    cleanup_temp_files(args.dest_folder)

    # Load the dataframe
    print(f"Loading data from: {args.pickle_file}")
    try:
        df = pd.read_pickle(args.pickle_file)
        print(f"Loaded preference data shape: {df.shape}")

        # Validate even number of rows
        if len(df) % 2 != 0:
            print(
                f"Warning: Dataframe has odd number of rows ({len(df)}). Last row will be skipped.",
                file=sys.stderr,
            )

    except Exception as e:
        print(f"Error loading pickle file: {str(e)}", file=sys.stderr)
        sys.exit(1)

    # Validate AWS credentials
    validate_aws_credentials()

    # Set number of workers
    num_workers = args.num_workers or cpu_count()
    print(f"\nUsing {num_workers} parallel workers for downloads")

    # Process all rows in pairs
    print(f"Downloading audio files to: {args.dest_folder}")
    print(f"Processing {len(df) // 2} pairs of rows")

    # Save df length before we delete it (needed for summary)
    total_rows = len(df)

    # Extract minimal data from rows to minimize memory usage
    print("Extracting minimal row data (IDs) for parallel processing...")
    row_pair_args = create_row_pair_data(df, args.dest_folder, args.verbose)

    num_pairs = len(row_pair_args)

    # Free up memory - we no longer need the full DataFrame
    # Because row_pair_args now only contains strings/IDs and not references to metadata dicts,
    # the GC can actually reclaim the heavy dataframe memory here.
    del df

    # Initialize accumulators for streaming results (avoid holding all results in memory)
    successful_clip_ids = []
    successful_pairs = 0
    failed_pairs = 0
    mismatched_pairs = 0
    total_files_skipped = 0
    total_files_downloaded = 0

    # Process in parallel with streaming result processing
    # Use initializer to share S3 client
    with Pool(processes=num_workers, initializer=init_worker) as pool:
        # Process with progress bar, handling results as they come in
        for result in tqdm(
            pool.imap_unordered(process_row_pair, row_pair_args),
            total=num_pairs,
            desc="Processing row pairs",
        ):
            # Process result immediately to avoid accumulating in memory
            if result.success:
                successful_pairs += 1
                successful_clip_ids.extend(result.clip_ids)
            else:
                failed_pairs += 1

            if not result.edited_ids_match:
                mismatched_pairs += 1

            total_files_skipped += result.files_skipped
            total_files_downloaded += result.files_downloaded

    # Free row_pair_args after processing
    del row_pair_args

    # Create summary from accumulated stats
    summary = DownloadSummary(
        total_pairs=num_pairs,
        successful_pairs=successful_pairs,
        failed_pairs=failed_pairs,
        mismatched_pairs=mismatched_pairs,
        successful_clip_ids=successful_clip_ids,
        total_files_skipped=total_files_skipped,
        total_files_downloaded=total_files_downloaded,
    )

    # Write successful downloads to JSON
    print(f"\nWriting successful downloads to: {args.output_json}")
    try:
        with open(args.output_json, "w") as f:
            json.dump(summary.successful_clip_ids, f, indent=2)
        print(
            f"Successfully wrote {len(summary.successful_clip_ids)} clip IDs to JSON file"
        )
    except Exception as e:
        print(f"Error writing JSON file: {str(e)}", file=sys.stderr)

    # Print summary
    print_summary(summary, total_rows, args.output_json)

    # Exit with error code if there were failures
    if summary.failed_pairs > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()