#!/usr/bin/env python3
"""
Merge multiple interesting_clips_*.pkl files into a single fully_merged file.

This script merges all interesting_clips_*.pkl files in a specified folder,
handling duplicates by 'id' column (keeping the last occurrence) and managing
memory efficiently for large files.
"""

import argparse
import gc
import glob
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional

import pandas as pd
from tqdm import tqdm


def merge_preference_files(
    folder_path: str,
    output_name: Optional[str] = None,
    force_recreate: bool = False,
) -> None:
    """
    Merge all interesting_clips_*.pkl files in a folder into one file.

    Args:
        folder_path: Path to the folder containing pkl files
        output_name: Optional custom output filename (without .pkl extension)
        force_recreate: If True, recreate from scratch ignoring existing merged file

    Raises:
        ValueError: If folder doesn't exist or no pkl files found
    """
    folder = Path(folder_path)
    if not folder.exists():
        raise ValueError(f"Folder not found: {folder_path}")

    # Extract folder name for output file
    folder_name = folder.name
    today = datetime.now().strftime("%Y%m%d")
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    if output_name:
        output_file = folder / f"{output_name}.pkl"
        summary_file = folder / f"{output_name}.txt"
    else:
        output_file = folder / f"fully_merged_{folder_name}_{today}.pkl"
        summary_file = folder / f"fully_merged_{folder_name}_{today}.txt"

    print(f"\n{'='*80}")
    print(f"📁 Merging preference data in: {folder_name}")
    print(f"🎯 Output file: {output_file.name}")
    print(f"📝 Summary file: {summary_file.name}")
    print(f"{'='*80}\n")

    # Track statistics for summary
    file_stats: List[Dict[str, any]] = []
    initial_size = 0

    # Load base dataframe if exists and not forcing recreate
    base_ctime = 0
    if output_file.exists() and not force_recreate:
        print(f"📂 Loading existing merged file: {output_file.name}")
        df = pd.read_pickle(output_file)
        base_ctime = os.path.getctime(output_file)
        initial_size = len(df)
        print(f"   Current size: {len(df):,} rows")
        print(f"   Shape: {df.shape}")
    else:
        if force_recreate:
            print("🔄 Force recreate mode - starting from scratch")
        else:
            print("✨ No existing merged file found - creating new one")
        df = pd.DataFrame()
        initial_size = 0

    # Find all interesting_clips_*.pkl files
    pattern = str(folder / "interesting_clips_*.pkl")
    pkl_files = glob.glob(pattern)

    if not pkl_files:
        raise ValueError(f"No interesting_clips_*.pkl files found in {folder_path}")

    print(f"\n🔍 Found {len(pkl_files)} interesting_clips_*.pkl files")

    # Filter files newer than base file
    newer_files = []
    for f in pkl_files:
        f_ctime = os.path.getctime(f)
        if f_ctime > base_ctime:
            newer_files.append(f)

    # Sort by creation time
    newer_files.sort(key=lambda x: os.path.getctime(x))

    if not newer_files:
        print("\n✅ No new files to process - merged file is up to date!")
        return

    print(f"📥 Processing {len(newer_files)} newer files\n")

    # Process each file
    for pkl_file in tqdm(newer_files, desc="Merging files"):
        file_name = os.path.basename(pkl_file)
        print(f"\n  Processing: {file_name}")

        prev_size = len(df)

        # Load temp dataframe with exception handling
        try:
            temp_df = pd.read_pickle(pkl_file)
            new_size = len(temp_df)
        except Exception as e:
            print(f"    ⚠️  Error loading file: {e}")
            print(f"    Skipping {file_name}")
            continue

        # Convert datetime columns if they exist
        for col in ["created_at", "updated_at"]:
            if col in temp_df.columns:
                temp_df[col] = pd.to_datetime(temp_df[col], utc=True)

        # Merge and handle duplicates - OPTIMIZED VERSION
        # Instead of concat-then-drop_duplicates (slow for large data):
        # 1. Use set lookup O(1) to identify duplicates
        # 2. Compare by updated_at timestamp (keep most recent)
        # 3. Filter dataframes before concat (smaller operations)
        # This is ~10-100x faster for large dataframes!
        if "id" in temp_df.columns and "id" in df.columns:
            # Build set of existing IDs for O(1) lookup
            existing_ids = set(df["id"].values)

            # Find duplicates and new records
            is_duplicate = temp_df["id"].isin(existing_ids)
            duplicate_ids = temp_df[is_duplicate]["id"].values
            new_records = temp_df[~is_duplicate]

            # Handle duplicate records by comparing updated_at timestamps
            if (
                len(duplicate_ids) > 0
                and "updated_at" in df.columns
                and "updated_at" in temp_df.columns
            ):
                # VECTORIZED approach - much faster than iterrows()!
                # Get records with duplicate IDs from both dataframes
                existing_dupes = df[df["id"].isin(duplicate_ids)][
                    ["id", "updated_at"]
                ].copy()
                incoming_dupes = temp_df[is_duplicate][["id", "updated_at"]].copy()

                # Merge to compare timestamps side-by-side (vectorized operation)
                comparison = incoming_dupes.merge(
                    existing_dupes,
                    on="id",
                    how="left",
                    suffixes=("_incoming", "_existing"),
                )

                # Clean up intermediate dataframes immediately
                del existing_dupes, incoming_dupes

                # Vectorized comparison: keep incoming if it's more recent
                # Handle NaT (missing timestamps) - if existing is NaT, always keep incoming
                keep_incoming_mask = comparison["updated_at_existing"].isna() | (
                    comparison["updated_at_incoming"].notna()
                    & (
                        comparison["updated_at_incoming"]
                        > comparison["updated_at_existing"]
                    )
                )

                # Get IDs to keep from incoming (more recent)
                keep_incoming_ids = comparison.loc[keep_incoming_mask, "id"].values

                # Clean up comparison dataframe
                del comparison, keep_incoming_mask

                # Remove outdated records and add updated ones
                if len(keep_incoming_ids) > 0:
                    df = df[~df["id"].isin(keep_incoming_ids)]
                    updated_records = temp_df[temp_df["id"].isin(keep_incoming_ids)]
                    df = pd.concat([df, updated_records], ignore_index=True)
                    del updated_records

                duplicates = len(keep_incoming_ids)
                del keep_incoming_ids
            elif len(duplicate_ids) > 0:
                # No updated_at column - fall back to file order (newer file wins)
                df = df[~df["id"].isin(duplicate_ids)]
                updated_records = temp_df[is_duplicate]
                df = pd.concat([df, updated_records], ignore_index=True)
                duplicates = len(duplicate_ids)
            else:
                duplicates = 0

            # Add completely new records
            if len(new_records) > 0:
                df = pd.concat([df, new_records], ignore_index=True)
        elif "id" in temp_df.columns:
            # First file with IDs
            df = temp_df.copy()
            duplicates = 0
        else:
            # No ID column - simple concat
            df = pd.concat([df, temp_df], ignore_index=True)
            duplicates = 0

        current_size = len(df)
        net_increase = current_size - prev_size

        # Store statistics
        file_stats.append(
            {
                "filename": file_name,
                "input_size": new_size,
                "prev_total": prev_size,
                "current_total": current_size,
                "net_increase": net_increase,
                "duplicates": duplicates,
            }
        )

        # Print statistics
        print(f"    Input size: {new_size:,}")
        print(f"    Previous total: {prev_size:,}")
        print(f"    Current total: {current_size:,}")
        print(f"    Net increase: {net_increase:,}")
        if duplicates > 0:
            print(f"    Duplicates removed: {duplicates:,}")

        # Aggressive memory management after each file
        del temp_df
        gc.collect()

        # Extra gc every 2 files to free up memory
        if len(file_stats) % 2 == 0:
            print(
                f"    🧹 Running extra garbage collection (processed {len(file_stats)} files)"
            )
            gc.collect()

    # Final statistics
    print(f"\n{'='*80}")
    print("📊 Final Statistics:")
    print(f"{'='*80}")
    print(f"Final shape: {df.shape}")
    print(f"Total rows: {len(df):,}")
    if "id" in df.columns:
        print(f"Unique IDs: {df['id'].nunique():,}")
    print(f"\n💾 Saving to: {output_file}")

    # Save merged dataframe
    df.to_pickle(output_file)

    # Get file size
    file_size = output_file.stat().st_size
    file_size_gb = file_size / (1024**3)
    print(f"✅ Saved successfully! File size: {file_size_gb:.2f} GB")

    # Calculate final statistics
    final_size = len(df)
    unique_ids = df["id"].nunique() if "id" in df.columns else None
    total_duplicates = sum(stat["duplicates"] for stat in file_stats)
    total_net_increase = final_size - initial_size

    # Write summary file
    print(f"\n📝 Writing summary to: {summary_file.name}")
    with open(summary_file, "w") as f:
        f.write("=" * 80 + "\n")
        f.write("PREFERENCE DATA MERGE SUMMARY\n")
        f.write("=" * 80 + "\n\n")

        f.write(f"Timestamp: {timestamp}\n")
        f.write(f"Folder: {folder_name}\n")
        f.write(f"Output File: {output_file.name}\n")
        f.write(f"Output Size: {file_size_gb:.2f} GB ({file_size:,} bytes)\n\n")

        f.write("=" * 80 + "\n")
        f.write("OVERALL STATISTICS\n")
        f.write("=" * 80 + "\n\n")

        f.write(f"Initial size (before merge): {initial_size:,} rows\n")
        f.write(f"Final size (after merge): {final_size:,} rows\n")
        f.write(f"Total net increase: {total_net_increase:,} rows\n")
        if unique_ids is not None:
            f.write(f"Unique IDs: {unique_ids:,}\n")
        f.write(f"Total files processed: {len(file_stats)}\n")
        f.write(f"Total duplicates removed: {total_duplicates:,}\n\n")

        if file_stats:
            f.write("=" * 80 + "\n")
            f.write("FILE-BY-FILE DETAILS\n")
            f.write("=" * 80 + "\n\n")

            for i, stat in enumerate(file_stats, 1):
                f.write(f"File {i}: {stat['filename']}\n")
                f.write(f"  Input size: {stat['input_size']:,} rows\n")
                f.write(f"  Previous total: {stat['prev_total']:,} rows\n")
                f.write(f"  Current total: {stat['current_total']:,} rows\n")
                f.write(f"  Net increase: {stat['net_increase']:,} rows\n")
                if stat["duplicates"] > 0:
                    f.write(f"  Duplicates removed: {stat['duplicates']:,}\n")
                f.write("\n")

        f.write("=" * 80 + "\n")
        f.write("END OF SUMMARY\n")
        f.write("=" * 80 + "\n")

    print("✅ Summary saved successfully!")

    # Clean up memory
    del df
    gc.collect()

    print(f"\n{'='*80}")
    print("✨ Merge complete!")
    print(f"{'='*80}\n")


def main() -> None:
    """Main entry point for the script."""
    parser = argparse.ArgumentParser(
        description="Merge multiple interesting_clips_*.pkl files into one file",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Merge files in auk_t0 folder
  python merge_preference_data.py /home/tony/Data/Preference/auk_t0
  
  # Merge with custom output name
  python merge_preference_data.py /home/tony/Data/Preference/crow_t1 --output-name custom_merged
  
  # Force recreate from scratch
  python merge_preference_data.py /home/tony/Data/Preference/auk_t1 --force-recreate
        """,
    )

    parser.add_argument(
        "folder",
        type=str,
        help="Path to folder containing interesting_clips_*.pkl files",
    )

    parser.add_argument(
        "-o",
        "--output-name",
        type=str,
        help="Custom output filename (without .pkl extension)",
        default=None,
    )

    parser.add_argument(
        "-f",
        "--force-recreate",
        action="store_true",
        help="Force recreate from scratch, ignoring existing merged file",
    )

    args = parser.parse_args()

    try:
        merge_preference_files(
            folder_path=args.folder,
            output_name=args.output_name,
            force_recreate=args.force_recreate,
        )
    except Exception as e:
        print(f"\n❌ Error: {e}")
        raise


if __name__ == "__main__":
    main()