"""
Test Script for User Clustering Pipeline
=======================================
Creates synthetic test data and validates the clustering pipeline end-to-end.
"""

import pandas as pd
import numpy as np
import json
import os
import shutil
from datetime import datetime, timedelta
from pathlib import Path
import subprocess
import sys


def create_test_data():
    """Create synthetic test data representing different user types."""
    print("Creating synthetic test data...")

    # Set random seed for reproducibility
    np.random.seed(42)

    # Initialize lists to store data
    total_clips = []
    reactions = []
    boosts = []
    playlists = []

    # Current timestamp
    base_time = datetime(2024, 1, 1)

    # User ID counter
    user_id_counter = 1000

    # 1. Create BOT users (20 users)
    print("Creating bot users...")
    for i in range(20):
        user_id = f"bot_{user_id_counter + i}"

        # Bots create many clips rapidly
        creation_time = base_time
        for j in range(100):  # Each bot creates 100 clips
            clip_id = f"clip_{user_id}_{j}"

            # Rapid creation (every 30 seconds)
            creation_time += timedelta(seconds=30)

            clip = {
                "id": clip_id,
                "user_id": user_id,
                "created_at": creation_time,
                "updated_at": creation_time,
                "metadata": json.dumps({}),  # No lyrics generation
                "model_name": "model_v1",  # Always same model
                "task": "generate",
                "source": "web",
                "creation_source": "api",
                "duration": 30,
                "is_public": True,
                "is_deleted": False,
                "is_pro_user": False,
                "play_count": np.random.randint(0, 10),
                "upvote_count": 0,
                "prompt_text": f"prompt {j % 5}",  # Reuses prompts
                "continued_parent": None,
            }
            total_clips.append(clip)

            # Minimal boosts data
            boost = {
                "clip_id": clip_id,
                "created_at": creation_time,
                "download_audio_count": 0,
                "download_video_count": 0,
                "share_count": 0,
                "reuse_prompt_count": j % 5,  # High reuse
            }
            boosts.append(boost)

        # Bots don't react to anything (no consumption)

    user_id_counter += 20

    # 2. Create CASUAL users (30 users)
    print("Creating casual users...")
    for i in range(30):
        user_id = f"casual_{user_id_counter + i}"

        # Casual users create moderate number of clips
        creation_time = base_time
        num_clips = np.random.randint(10, 30)

        for j in range(num_clips):
            clip_id = f"clip_{user_id}_{j}"

            # Random gaps between creations (hours to days)
            creation_time += timedelta(hours=np.random.randint(1, 48))

            # 70% of clips have lyrics generation
            metadata = {}
            if np.random.random() < 0.7:
                metadata["gpt_description_prompt"] = (
                    f"Write lyrics about {np.random.choice(['love', 'life', 'happiness'])}"
                )

            clip = {
                "id": clip_id,
                "user_id": user_id,
                "created_at": creation_time,
                "updated_at": creation_time,
                "metadata": json.dumps(metadata),
                "model_name": np.random.choice(["model_v1", "model_v2"]),
                "task": np.random.choice(["generate", "extend"]),
                "source": np.random.choice(["web", "ios", "android"]),
                "creation_source": np.random.choice(["web", "mobile"]),
                "duration": np.random.randint(20, 40),
                "is_public": np.random.random() < 0.8,
                "is_deleted": False,
                "is_pro_user": False,  # Not pro
                "play_count": np.random.randint(10, 100),
                "upvote_count": np.random.randint(0, 5),
                "prompt_text": f"casual prompt {j}",
                "continued_parent": None,
            }
            total_clips.append(clip)

            # Moderate engagement
            boost = {
                "clip_id": clip_id,
                "created_at": creation_time,
                "download_audio_count": np.random.randint(0, 5),
                "download_video_count": np.random.randint(0, 2),
                "share_count": np.random.randint(0, 3),
                "reuse_prompt_count": 0,
            }
            boosts.append(boost)

            # Add to playlist sometimes
            if np.random.random() < 0.3:
                playlist = {
                    "id": f"pl_{len(playlists)}",
                    "clip_id": clip_id,
                    "playlist_id": f"playlist_{i % 5}",
                    "relative_index": j,
                    "updated_at": creation_time,
                }
                playlists.append(playlist)

        # Casual users have some reactions
        for _ in range(np.random.randint(5, 20)):
            reaction = {
                "id": f"reaction_{len(reactions)}",
                "user_id": user_id,
                "clip_id": np.random.choice([c["id"] for c in total_clips]),
                "reaction_type": "play",
                "play_count": 1,
                "updated_at": creation_time,
            }
            reactions.append(reaction)

    user_id_counter += 30

    # 3. Create PRO SERIOUS users (20 users)
    print("Creating pro serious users...")
    for i in range(20):
        user_id = f"pro_{user_id_counter + i}"

        # Pro users create quality content regularly
        creation_time = base_time
        num_clips = np.random.randint(40, 80)

        for j in range(num_clips):
            clip_id = f"clip_{user_id}_{j}"

            # Regular creation pattern
            creation_time += timedelta(hours=np.random.randint(12, 36))

            # Rarely use lyrics generation (they write their own)
            metadata = {}
            if np.random.random() < 0.1:
                metadata["gpt_description_prompt"] = "Professional prompt"

            clip = {
                "id": clip_id,
                "user_id": user_id,
                "created_at": creation_time,
                "updated_at": creation_time,
                "metadata": json.dumps(metadata),
                "model_name": np.random.choice(
                    ["model_v2", "model_v3", "model_v4"]
                ),  # Uses latest models
                "task": np.random.choice(
                    ["generate", "extend", "remix", "cover"]
                ),  # Diverse tasks
                "source": np.random.choice(["web", "ios"]),  # Professional platforms
                "creation_source": np.random.choice(["studio", "web", "pro_tools"]),
                "duration": np.random.randint(30, 60),  # Longer clips
                "is_public": True,
                "is_deleted": np.random.random() < 0.05,  # Low deletion rate
                "is_pro_user": True,  # Pro user
                "play_count": np.random.randint(100, 1000),  # High engagement
                "upvote_count": np.random.randint(10, 50),
                "prompt_text": f"professional prompt version {j}",
                "continued_parent": f"clip_{user_id}_{j-1}"
                if j > 0 and np.random.random() < 0.3
                else None,
            }
            total_clips.append(clip)

            # High engagement
            boost = {
                "clip_id": clip_id,
                "created_at": creation_time,
                "download_audio_count": np.random.randint(5, 20),
                "download_video_count": np.random.randint(2, 10),
                "share_count": np.random.randint(5, 15),
                "reuse_prompt_count": np.random.randint(0, 5),
            }
            boosts.append(boost)

            # Often in playlists
            if np.random.random() < 0.6:
                playlist = {
                    "id": f"pl_{len(playlists)}",
                    "clip_id": clip_id,
                    "playlist_id": f"playlist_pro_{i}",
                    "relative_index": j % 10,
                    "updated_at": creation_time,
                }
                playlists.append(playlist)

        # Pro users are active consumers too
        for _ in range(np.random.randint(50, 100)):
            reaction = {
                "id": f"reaction_{len(reactions)}",
                "user_id": user_id,
                "clip_id": np.random.choice([c["id"] for c in total_clips]),
                "reaction_type": "play",
                "play_count": 1,
                "updated_at": creation_time,
            }
            reactions.append(reaction)

    user_id_counter += 20

    # 4. Create OTHER/NORMAL users (30 users)
    print("Creating other/normal users...")
    for i in range(30):
        user_id = f"normal_{user_id_counter + i}"

        # Variable behavior
        creation_time = base_time
        num_clips = np.random.randint(5, 25)

        for j in range(num_clips):
            clip_id = f"clip_{user_id}_{j}"

            # Irregular creation pattern
            creation_time += timedelta(days=np.random.randint(1, 7))

            # Sometimes use lyrics
            metadata = {}
            if np.random.random() < 0.3:
                metadata["gpt_description_prompt"] = "Some lyrics"

            clip = {
                "id": clip_id,
                "user_id": user_id,
                "created_at": creation_time,
                "updated_at": creation_time,
                "metadata": json.dumps(metadata),
                "model_name": np.random.choice(["model_v1", "model_v2"]),
                "task": "generate",
                "source": np.random.choice(["web", "ios", "android"]),
                "creation_source": "web",
                "duration": 30,
                "is_public": np.random.random() < 0.6,
                "is_deleted": np.random.random() < 0.1,
                "is_pro_user": np.random.random() < 0.2,  # Some are pro
                "play_count": np.random.randint(5, 50),
                "upvote_count": np.random.randint(0, 10),
                "prompt_text": f"normal prompt {j}",
                "continued_parent": None,
            }
            total_clips.append(clip)

            # Variable engagement
            boost = {
                "clip_id": clip_id,
                "created_at": creation_time,
                "download_audio_count": np.random.randint(0, 10),
                "download_video_count": np.random.randint(0, 5),
                "share_count": np.random.randint(0, 5),
                "reuse_prompt_count": 0,
            }
            boosts.append(boost)

        # Some reactions
        for _ in range(np.random.randint(10, 30)):
            reaction = {
                "id": f"reaction_{len(reactions)}",
                "user_id": user_id,
                "clip_id": np.random.choice([c["id"] for c in total_clips]),
                "reaction_type": "play",
                "play_count": 1,
                "updated_at": creation_time,
            }
            reactions.append(reaction)

    # Create DataFrames
    total_clip_df = pd.DataFrame(total_clips)
    reaction_df = pd.DataFrame(reactions)
    boosts_action_df = pd.DataFrame(boosts)
    playlist_clip_df = pd.DataFrame(playlists)

    # Add missing columns with default values
    reaction_df["skip_count"] = 0
    reaction_df["flagged"] = False
    reaction_df["is_pro_user"] = reaction_df["user_id"].str.contains("pro_")

    boosts_action_df["updated_at"] = boosts_action_df["created_at"]
    boosts_action_df["first_published_at"] = boosts_action_df["created_at"]
    boosts_action_df["is_public_approved"] = True

    # Add remaining columns with defaults
    for col in [
        "time_used",
        "status",
        "discord_message_id",
        "prompt_id",
        "request_id",
        "is_generated",
        "s3_id",
        "batch_index",
        "daily_theme_id",
        "image_s3_id",
        "dislike_count",
        "flag_count",
        "skip_count",
        "title",
        "slug",
        "allow_comments",
        "is_hidden",
        "display_tags",
        "clip_type",
        "edited_clip_id",
        "date",
        "hour",
    ]:
        if col not in total_clip_df.columns:
            total_clip_df[col] = None

    total_clip_df["hour"] = total_clip_df["created_at"].dt.hour
    total_clip_df["date"] = total_clip_df["created_at"].dt.date

    print(f"Created test data:")
    print(f"  - Total clips: {len(total_clip_df)}")
    print(f"  - Unique users: {total_clip_df['user_id'].nunique()}")
    print(f"  - Reactions: {len(reaction_df)}")
    print(f"  - Boosts: {len(boosts_action_df)}")
    print(f"  - Playlist entries: {len(playlist_clip_df)}")

    return total_clip_df, reaction_df, boosts_action_df, playlist_clip_df


def save_test_data(output_dir: str):
    """Save test data to files."""
    print(f"\nSaving test data to {output_dir}...")

    # Create output directory
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    # Generate test data
    total_clip_df, reaction_df, boosts_action_df, playlist_clip_df = create_test_data()

    # Save as CSV files
    total_clip_df.to_csv(f"{output_dir}/total_clip_df.csv", index=False)
    reaction_df.to_csv(f"{output_dir}/reaction_df.csv", index=False)
    boosts_action_df.to_csv(f"{output_dir}/boosts_action_df.csv", index=False)
    playlist_clip_df.to_csv(f"{output_dir}/playlist_clip_df.csv", index=False)

    # Also save one as pickle to test pickle loading
    total_clip_df.to_pickle(f"{output_dir}/total_clip_df.pkl")

    print("Test data saved successfully!")

    return total_clip_df


def run_clustering_pipeline(input_dir: str, output_dir: str):
    """Run the clustering pipeline."""
    print(f"\nRunning clustering pipeline...")
    print(f"  Input: {input_dir}")
    print(f"  Output: {output_dir}")

    # Run the pipeline
    cmd = [
        sys.executable,
        "user_clustering.py",
        "--input-dir",
        input_dir,
        "--output-dir",
        output_dir,
        "--sample-size",
        "1000",  # Use sampling for faster test
    ]

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print("Pipeline failed!")
        print("STDOUT:", result.stdout)
        print("STDERR:", result.stderr)
        raise RuntimeError("Pipeline execution failed")

    print("Pipeline completed successfully!")


def validate_results(output_dir: str, total_clip_df: pd.DataFrame):
    """Validate the clustering results."""
    print(f"\nValidating results in {output_dir}...")

    # Load results
    assignments = pd.read_csv(f"{output_dir}/user_cluster_assignments.csv")
    profiles = pd.read_csv(f"{output_dir}/cluster_profiles.csv", index_col=0)

    with open(f"{output_dir}/cluster_type_mapping.json", "r") as f:
        cluster_mapping = json.load(f)

    print(f"\nCluster assignments:")
    print(f"  Total users clustered: {len(assignments)}")
    print(f"  Number of clusters: {assignments['cluster_label'].nunique()}")

    # Check cluster types
    print(f"\nCluster type mapping:")
    for cluster_id, cluster_type in cluster_mapping.items():
        count = len(assignments[assignments["cluster_label"] == int(cluster_id)])
        print(f"  Cluster {cluster_id} ({cluster_type}): {count} users")

    # Validate expected patterns
    print(f"\nValidating expected patterns...")

    # Check bot detection
    bot_clusters = [k for k, v in cluster_mapping.items() if v == "bot"]
    if bot_clusters:
        bot_cluster_id = int(bot_clusters[0])
        bot_users = assignments[assignments["cluster_label"] == bot_cluster_id][
            "user_id"
        ]
        bot_user_types = bot_users.str.split("_").str[0].value_counts()
        print(f"\nBot cluster contains:")
        print(bot_user_types)

        # Check bot characteristics
        bot_profile = profiles.loc[bot_cluster_id]
        assert (
            bot_profile["consumption_ratio"] < 0.1
        ), "Bot cluster should have low consumption"
        assert (
            bot_profile["creation_burst_score"] > 50
        ), "Bot cluster should have high burst score"
        print("✓ Bot cluster characteristics validated")

    # Check casual users
    casual_clusters = [k for k, v in cluster_mapping.items() if v == "casual"]
    if casual_clusters:
        casual_cluster_id = int(casual_clusters[0])
        casual_profile = profiles.loc[casual_cluster_id]
        assert (
            casual_profile["has_lyrics_generation"] > 30
        ), "Casual cluster should use lyrics generation"
        print("✓ Casual cluster characteristics validated")

    # Check pro users
    pro_clusters = [k for k, v in cluster_mapping.items() if v == "pro_serious"]
    if pro_clusters:
        pro_cluster_id = int(pro_clusters[0])
        pro_profile = profiles.loc[pro_cluster_id]
        assert (
            pro_profile["pro_user_ratio"] > 0.5
        ), "Pro cluster should have high pro ratio"
        print("✓ Pro cluster characteristics validated")

    # Load and display summary report
    print(f"\n{'='*60}")
    print("CLUSTERING SUMMARY REPORT:")
    print("=" * 60)
    with open(f"{output_dir}/clustering_summary_report.txt", "r") as f:
        print(f.read())

    print(f"\n{'='*60}")
    print("✓ All validations passed!")
    print("✓ Test completed successfully!")


def main():
    """Main test execution."""
    print("=" * 60)
    print("USER CLUSTERING PIPELINE - END-TO-END TEST")
    print("=" * 60)

    # Setup directories
    test_dir = "test_data"
    output_dir = "test_output"

    try:
        # Clean up previous test runs
        for dir_path in [test_dir, output_dir]:
            if os.path.exists(dir_path):
                shutil.rmtree(dir_path)

        # Step 1: Create and save test data
        total_clip_df = save_test_data(test_dir)

        # Step 2: Run clustering pipeline
        run_clustering_pipeline(test_dir, output_dir)

        # Step 3: Validate results
        validate_results(output_dir, total_clip_df)

    except Exception as e:
        print(f"\n❌ Test failed: {str(e)}")
        raise
    finally:
        # Cleanup (optional - comment out to inspect results)
        # for dir_path in [test_dir, output_dir]:
        #     if os.path.exists(dir_path):
        #         shutil.rmtree(dir_path)
        pass


if __name__ == "__main__":
    main()