#!/usr/bin/env python3
"""
Process stems_captions data and append to existing metadata files
Following Victor's pattern from add_extreme.ipynb
"""

import os
import json
import random
from tqdm import tqdm

# =============================================================================
# CONFIGURATION CONSTANTS - Modify these as needed
# =============================================================================

# Input data file
INPUT_STEMS_CAPTIONS_FILE = "/home/vibert/shared/stems_captions_c686c20e.jsonl"

# Output directories (using ~/tmp for intermediate processing)
OUTPUT_DIR = "/home/vibert/tmp/stems_captions_output"
METADATA_DIR = "/home/vibert/tmp/new_metadata"

# Source metadata files (v3 -> v5) - will need to be copied from restricted dirs
SOURCE_TRAIN_FILE = "/app2/suno/data/auk_v0/metas_v3_tr.jsonl"
SOURCE_VAL_FILE = "/app2/suno/data/auk_v0/metas_v3_val.jsonl"

# Target metadata files (v5) - will be created in ~/tmp
TARGET_TRAIN_FILE = f"{METADATA_DIR}/metas_v5_tr.jsonl"
TARGET_VAL_FILE = f"{METADATA_DIR}/metas_v5_val.jsonl"

# Output split files
OUTPUT_TRAIN_FILE = f"{OUTPUT_DIR}/train.jsonl"
OUTPUT_VAL_FILE = f"{OUTPUT_DIR}/val.jsonl"

# Split configuration
TRAIN_VAL_SPLIT_RATIO = 0.95  # 95% train, 5% validation
RANDOM_SEED = 42

# Processing options
REQUIRE_BOTH_STEMS_AND_CAPTIONS = False  # Only process records with both stems and stems_captions

# =============================================================================
# MAIN PROCESSING FUNCTIONS
# =============================================================================


def load_stems_captions_data(input_file):
    """Load and filter stems_captions data"""
    print(f"Loading stems_captions data from: {input_file}")

    if not os.path.exists(input_file):
        raise FileNotFoundError(f"Input file not found: {input_file}")

    records = []
    total_count = 0
    filtered_count = 0

    with open(input_file, "r") as f:
        for line in tqdm(f, desc="Loading records"):
            total_count += 1
            try:
                record = json.loads(line.strip())

                # Filter records based on requirements
                has_stems = record.get("stems") and len(record["stems"]) > 0
                has_stems_captions = record.get("stems_captions") and len(record["stems_captions"]) > 0

                if REQUIRE_BOTH_STEMS_AND_CAPTIONS:
                    if has_stems and has_stems_captions:
                        records.append(record)
                        filtered_count += 1
                else:
                    records.append(record)
                    filtered_count += 1

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

    print(f"Loaded {filtered_count} records out of {total_count} total records")
    if REQUIRE_BOTH_STEMS_AND_CAPTIONS:
        print(f"Filtered for records with both stems and stems_captions")

    return records


def split_train_val(records, split_ratio=TRAIN_VAL_SPLIT_RATIO, random_seed=RANDOM_SEED):
    """Split records into train and validation sets"""
    print(f"Splitting {len(records)} records with ratio {split_ratio:.2f}")

    # Shuffle with fixed seed for reproducibility
    random.seed(random_seed)
    shuffled_records = records.copy()
    random.shuffle(shuffled_records)

    # Calculate split index
    split_index = int(len(shuffled_records) * split_ratio)

    train_records = shuffled_records[:split_index]
    val_records = shuffled_records[split_index:]

    print(f"Split into {len(train_records)} train and {len(val_records)} val records")

    return train_records, val_records


def save_split_files(train_records, val_records, output_dir):
    """Save train and val splits to output directory"""
    print(f"Saving split files to: {output_dir}")

    # Create output directory
    os.makedirs(output_dir, exist_ok=True)

    # Save train file
    train_file = f"{output_dir}/train.jsonl"
    with open(train_file, "w") as f:
        for record in tqdm(train_records, desc="Writing train file"):
            f.write(json.dumps(record) + "\n")
    print(f"Saved {len(train_records)} train records to: {train_file}")

    # Save val file
    val_file = f"{output_dir}/val.jsonl"
    with open(val_file, "w") as f:
        for record in tqdm(val_records, desc="Writing val file"):
            f.write(json.dumps(record) + "\n")
    print(f"Saved {len(val_records)} val records to: {val_file}")

    return train_file, val_file


def copy_existing_metadata_files():
    """Copy v3 metadata files to v5 (following Victor's pattern)"""
    print("Copying existing v3 metadata files to v5...")

    # Check source files exist
    if not os.path.exists(SOURCE_TRAIN_FILE):
        raise FileNotFoundError(f"Source train file not found: {SOURCE_TRAIN_FILE}")
    if not os.path.exists(SOURCE_VAL_FILE):
        raise FileNotFoundError(f"Source val file not found: {SOURCE_VAL_FILE}")

    # Copy files
    import shutil

    shutil.copy2(SOURCE_TRAIN_FILE, TARGET_TRAIN_FILE)
    shutil.copy2(SOURCE_VAL_FILE, TARGET_VAL_FILE)

    print(f"Copied {SOURCE_TRAIN_FILE} -> {TARGET_TRAIN_FILE}")
    print(f"Copied {SOURCE_VAL_FILE} -> {TARGET_VAL_FILE}")


def append_to_metadata_files(train_records, val_records):
    """Append new records to existing v5 metadata files"""
    print("Appending stems_captions records to v5 metadata files...")

    # Append to train file
    with open(TARGET_TRAIN_FILE, "a") as f:
        for record in tqdm(train_records, desc="Appending to train file"):
            f.write(json.dumps(record) + "\n")
    print(f"Appended {len(train_records)} records to: {TARGET_TRAIN_FILE}")

    # Append to val file
    with open(TARGET_VAL_FILE, "a") as f:
        for record in tqdm(val_records, desc="Appending to val file"):
            f.write(json.dumps(record) + "\n")
    print(f"Appended {len(val_records)} records to: {TARGET_VAL_FILE}")


def print_summary_stats(records, train_records, val_records):
    """Print summary statistics"""
    print("\n" + "=" * 60)
    print("PROCESSING SUMMARY")
    print("=" * 60)

    # Overall stats
    print(f"Total input records: {len(records)}")
    print(f"Train records: {len(train_records)}")
    print(f"Val records: {len(val_records)}")
    print(f"Train/Val ratio: {len(train_records)/(len(train_records)+len(val_records)):.3f}")

    # Stems and captions stats
    records_with_stems = sum(1 for r in records if r.get("stems"))
    records_with_captions = sum(1 for r in records if r.get("stems_captions"))
    records_with_both = sum(1 for r in records if r.get("stems") and r.get("stems_captions"))

    print(f"\nContent statistics:")
    print(f"Records with stems: {records_with_stems}")
    print(f"Records with stems_captions: {records_with_captions}")
    print(f"Records with both: {records_with_both}")

    # Sample record info
    if records:
        sample_record = records[0]
        print(f"\nSample record ID: {sample_record.get('id', 'N/A')}")
        if sample_record.get("stems"):
            print(f"Sample stems: {list(sample_record['stems'].keys())}")
        if sample_record.get("stems_captions"):
            print(f"Sample stems_captions keys: {list(sample_record['stems_captions'].keys())}")

    print("=" * 60)


def main():
    """Main processing function"""
    print("🎵 Stems Captions Processing Script")
    print("Following Victor's pattern from add_extreme.ipynb")
    print(f"Input: {INPUT_STEMS_CAPTIONS_FILE}")
    print(f"Output dir: {OUTPUT_DIR}")
    print(f"Metadata: v3 -> v5")
    print()

    try:
        # Step 1: Load stems_captions data
        records = load_stems_captions_data(INPUT_STEMS_CAPTIONS_FILE)

        if not records:
            print("❌ No records to process!")
            return

        # Step 2: Split into train/val
        train_records, val_records = split_train_val(records)

        # Step 3: Save split files
        save_split_files(train_records, val_records, OUTPUT_DIR)

        # Step 4: Copy existing metadata files (v3 -> v5)
        copy_existing_metadata_files()

        # Step 5: Append to metadata files
        append_to_metadata_files(train_records, val_records)

        # Step 6: Print summary
        print_summary_stats(records, train_records, val_records)

        print("\n✅ Processing completed successfully!")
        print(f"📁 Split files: {OUTPUT_DIR}/")
        print(f"📁 Updated metadata: {TARGET_TRAIN_FILE}, {TARGET_VAL_FILE}")

    except Exception as e:
        print(f"❌ Error during processing: {e}")
        raise


if __name__ == "__main__":
    main()