#!/usr/bin/env python3
"""
Test script to verify Snowflake and Feast connectivity.
Run this from the feast/ directory with: python test_connection.py
"""

import os
import sys
from pathlib import Path


def test_snowflake_connection():
    """Test 1: Direct Snowflake Connection"""
    print("=" * 60)
    print("TEST 1: Direct Snowflake Connection")
    print("=" * 60)

    try:
        import snowflake.connector
        from cryptography.hazmat.primitives import serialization
        
        # Load private key
        path_str = "/Users/$USER/documents/snowflake_secrets/rsa_key.p8"
        private_key_path = Path(os.path.expandvars(path_str))
        if not private_key_path.exists():
            print(f"❌ ERROR: Private key not found at {private_key_path}")
            return False
        # Set the environment variable to the private key path
        os.environ["SNOWFLAKE_PRIVATE_KEY"] = private_key_path.read_text()
        
        with open(private_key_path, "rb") as key_file:
            private_key = serialization.load_pem_private_key(
                key_file.read(),
                password=None,  # If your key has a passphrase, add it here
            )
        
        private_key_der = private_key.private_bytes(
            encoding=serialization.Encoding.DER,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        
        print("✅ Private key loaded successfully")
        
        # Connect to Snowflake
        conn = snowflake.connector.connect(
            account="ccntdmk-suno",
            user=os.getenv("SNOWFLAKE_USER"),
            private_key=private_key_der,
            role="ACCOUNTADMIN",
            warehouse="DBT_DEV_MEDIUM",
            database="SUNO_DEV_DELPHINE",
            schema="PROD",
        )
        
        print("✅ Connected to Snowflake successfully")
        
        # Test query
        cursor = conn.cursor()
        cursor.execute("SELECT CURRENT_USER(), CURRENT_ROLE(), CURRENT_WAREHOUSE()")
        result = cursor.fetchone()
        print(f"   User: {result[0]}")
        print(f"   Role: {result[1]}")
        print(f"   Warehouse: {result[2]}")
        
        # Check if table exists
        cursor.execute("""
            SELECT COUNT(*) 
            FROM INFORMATION_SCHEMA.TABLES 
            WHERE TABLE_SCHEMA = 'PROD' 
            AND TABLE_NAME = 'RECS_HOOKS_RANKER_USER_FEATURE_STORE'
        """)
        table_exists = cursor.fetchone()[0]
        
        if table_exists:
            print(f"✅ Table RECS_HOOKS_RANKER_USER_FEATURE_STORE exists")
            
            # Get table info
            cursor.execute("""
                SELECT COLUMN_NAME, DATA_TYPE 
                FROM INFORMATION_SCHEMA.COLUMNS 
                WHERE TABLE_SCHEMA = 'PROD' 
                AND TABLE_NAME = 'RECS_HOOKS_RANKER_USER_FEATURE_STORE'
                ORDER BY ORDINAL_POSITION
            """)
            print("\n   Table columns:")
            for col_name, col_type in cursor.fetchall():
                print(f"     - {col_name}: {col_type}")
            
            # Count rows
            cursor.execute("SELECT COUNT(*) FROM RECS_HOOKS_RANKER_USER_FEATURE_STORE")
            row_count = cursor.fetchone()[0]
            print(f"\n   Row count: {row_count:,}")
            
            # Sample data
            if row_count > 0:
                cursor.execute("""
                    SELECT * 
                    FROM RECS_HOOKS_RANKER_USER_FEATURE_STORE 
                    LIMIT 3
                """)
                print("\n   Sample rows:")
                for row in cursor.fetchall():
                    print(f"     {row}")
        else:
            print(f"❌ Table RECS_HOOKS_RANKER_USER_FEATURE_STORE does not exist")
            print("   Available tables in PROD schema:")
            cursor.execute("""
                SELECT TABLE_NAME 
                FROM INFORMATION_SCHEMA.TABLES 
                WHERE TABLE_SCHEMA = 'PROD'
                LIMIT 10
            """)
            for table in cursor.fetchall():
                print(f"     - {table[0]}")
        
        cursor.close()
        conn.close()
        print("\n✅ Snowflake connection test PASSED\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Snowflake connection test FAILED: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_feast_feature_store():
    """Test 2: Feast Feature Store"""
    print("=" * 60)
    print("TEST 2: Feast Feature Store")
    print("=" * 60)

    try:
        from feast import FeatureStore
        
        # Initialize feature store
        store = FeatureStore(repo_path=".")
        print("✅ Feast feature store initialized")
        
        # List feature views
        feature_views = store.list_feature_views()
        print(f"\n   Feature views ({len(feature_views)}):")
        for fv in feature_views:
            print(f"     - {fv.name}")
        
        # List entities
        entities = store.list_entities()
        print(f"\n   Entities ({len(entities)}):")
        for entity in entities:
            print(f"     - {entity.name}")
        
        print("\n✅ Feast feature store test PASSED\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Feast feature store test FAILED: {e}\n")
        print("\nTip: You may need to run 'feast apply' first to register your features.")
        import traceback
        traceback.print_exc()
        return False


def test_user_historical_features():
    """Test 3: Fetch Historical Features"""
    print("=" * 60)
    print("TEST 3: Fetch Historical Features")
    print("=" * 60)

    try:
        import pandas as pd
        from datetime import datetime, timedelta
        from feast import FeatureStore
        
        store = FeatureStore(repo_path=".")
        
        # Create entity dataframe (you'll need real user IDs from your table)
        entity_df = pd.DataFrame({
            "USER_ID": [4937689, 4937689],
            "event_timestamp": [
                datetime.now() - timedelta(days=10),
                datetime.now(),
            ],
        })
        
        print("   Fetching features for users:", entity_df["USER_ID"].tolist())
        
        training_df = store.get_historical_features(
            entity_df=entity_df,
            features=["user_features:AVERAGE_WATCH_TIME_30D"],
        ).to_df()
        
        print("\n   Retrieved features:")
        print(training_df)
        
        print("\n✅ Historical features fetch test PASSED\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Historical features fetch test FAILED: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_user_hook_historical_features():
    """Test 4: Fetch User Hook Historical Features"""
    print("=" * 60)
    print("TEST 4: Fetch User Hook Historical Features")
    print("=" * 60)
    
    try:
        import pandas as pd
        from datetime import datetime
        from feast import FeatureStore
        
        store = FeatureStore(repo_path=".")
        
        # Create entity dataframe (you'll need real user IDs from your table)
        entity_df = pd.DataFrame({
            "USER_ID": [4937689],
            "HOOK_ID": ["7bca153a-0454-4911-ae04-90255821e8e1"],
            "event_timestamp": [datetime.now()],
        })
        
        print("   Fetching features for users:", entity_df["USER_ID"].tolist())
        
        training_df = store.get_historical_features(
            entity_df=entity_df,
            features=[
                "user_hook_features:WATCH_TIME",
            ],
        ).to_df()
        
        print("\n   Retrieved features:")
        print(training_df)
        for col in training_df.columns:
            print(f"     - {col}: {training_df[col].iloc[0]}")
        
        print("\n✅ User hook historical features fetch test PASSED\n")
        return True
        
    except Exception as e:
        print(f"\n❌ User hook historical features fetch test FAILED: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_materialization():
    """Test 5: Materialize Features"""
    print("=" * 60)
    print("TEST 5: Materialize Features")
    print("=" * 60)
    try:
        from feast import FeatureStore
        from datetime import datetime, timedelta
        
        store = FeatureStore(repo_path=".")

        end_date = datetime.fromisoformat("2025-09-16T01:00:00")
        start_date = end_date - timedelta(days=1)

        print(f"Materializing features from {start_date} to {end_date}")

        # Materialize to online store (Redis)
        # in prod, use materialize_incremental instead of materialize
        store.materialize(
            start_date=start_date,
            end_date=end_date,
            feature_views=["hook_characteristics_features"]
        )
        print("✅ Materialization complete!")
        return True
        
    except Exception as e:
        print(f"\n❌ Materialization test FAILED: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_online_features():
    """Test 6: Fetch Online Features
    
    Note: Assumes features were already materialized in Test 5.
    If features are None, check that the hook's EVENT_TIMESTAMP 
    matches the materialized date range from Test 5.
    """
    print("=" * 60)
    print("TEST 6: Fetch Online Features")
    print("=" * 60)
    
    try:
        from feast import FeatureStore
        from datetime import datetime
        import pandas as pd
        
        store = FeatureStore(repo_path=".")

        hook_id = "7bf0047e-7ea9-4fd3-9b86-1c22eef4f14c"
        # Create entity dataframe (you'll need real user IDs from your table)
        entity_df = pd.DataFrame({
            "HOOK_ID": [hook_id],
        })

        online_df = store.get_online_features(
            features=["hook_characteristics_features:GEMINI_RATING", "hook_characteristics_features:IS_AI_GENERATED"],
            entity_rows=entity_df.to_dict(orient="records"),
            full_feature_names=True,
        ).to_df()
        
        print("\n   Retrieved features from Redis:")
        print(online_df)
    
    except Exception as e:
        print(f"\n❌ Online features fetch test FAILED: {e}\n")
        import traceback
        traceback.print_exc()
        return False

    print("\n✅ Online features fetch test PASSED\n")
    return True

def main():
    """Run all tests"""
    results = []
    
    # Run tests
    results.append(("Snowflake Connection", test_snowflake_connection()))
    results.append(("Feast Feature Store", test_feast_feature_store()))
    results.append(("Historical Features", test_user_historical_features()))
    results.append(("User Hook Historical Features", test_user_hook_historical_features()))

    # Print summary
    print("=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)
    for test_name, passed in results:
        status = "✅ PASSED" if passed else "❌ FAILED"
        print(f"{test_name}: {status}")
    
    all_passed = all(passed for _, passed in results)
    if all_passed:
        print("\n🎉 ALL TESTS PASSED!")
    else:
        print("\n⚠️  Some tests failed. Check the output above for details.")
        sys.exit(1)


if __name__ == "__main__":
    main()