"""
Test script to verify PostgreSQL database migration is working properly.

Run this after setting up PostgreSQL to test the database integration.
"""
import asyncio
import os
import sys
from pathlib import Path

# Add src to path so we can import our modules
sys.path.insert(0, str(Path(__file__).parent / "src"))

async def test_database_migration():
    """Test the PostgreSQL database migration."""
    try:
        # Test basic imports
        print("✓ Testing imports...")
        from app.dashboard.data import (
            create_user, create_organization, get_user_by_email,
            create_api_key, get_dashboard_stats, init_demo_data
        )
        from app.dashboard.models import IndustryType, CompanySize
        from app.storage.database import get_connection
        print("✓ All imports successful")
        
        # Test database connection
        print("✓ Testing database connection...")
        async with get_connection() as conn:
            version = await conn.fetchval("SELECT version()")
            print(f"✓ Connected to: {version[:50]}...")
        
        # Test demo data creation
        print("✓ Testing demo data creation...")
        try:
            demo_user_id, demo_org_id = await init_demo_data()
            print(f"✓ Created demo user: {demo_user_id}")
            print(f"✓ Created demo organization: {demo_org_id}")
            
            # Test user retrieval
            user = await get_user_by_email("admin@salesai.com")
            if user:
                print(f"✓ Retrieved user: {user['first_name']} {user['last_name']}")
            
            # Test dashboard stats
            stats = await get_dashboard_stats(demo_org_id)
            print(f"✓ Dashboard stats - API usage: {stats.api_usage_current_month}/{stats.api_usage_limit}")
            
        except Exception as e:
            if "already exists" in str(e).lower():
                print("✓ Demo data already exists (this is fine)")
            else:
                raise e
        
        print("\n🎉 All database tests passed!")
        print("Your PostgreSQL migration is working correctly.")
        
    except Exception as e:
        print(f"\n❌ Database test failed: {e}")
        print("\nTroubleshooting steps:")
        print("1. Make sure PostgreSQL is running: docker-compose up -d")
        print("2. Check DATABASE_URL in .env file")
        print("3. Initialize schema: python init_db.py")
        return False
    
    return True

if __name__ == "__main__":
    # Load environment variables
    from dotenv import load_dotenv
    load_dotenv()
    
    # Check if DATABASE_URL is set
    if not os.getenv("DATABASE_URL"):
        print("❌ DATABASE_URL not found in environment")
        print("Please create .env file with DATABASE_URL=postgresql://salesbrain:salesbrain123@localhost:5432/salesbrain")
        sys.exit(1)
    
    print("PostgreSQL Database Migration Test")
    print("=" * 40)
    
    success = asyncio.run(test_database_migration())
    sys.exit(0 if success else 1)