"""
Database migration script to add indicator preferences and AI plan generation tables
Run this to add the new tables to your existing database
"""

import sys
import os
from sqlalchemy import create_engine, text

# Add parent directory to path to import app modules
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.config import settings
from app.models.models import Base, UserIndicatorPreferences, AIPlanGeneration


def run_migration():
    """Create the new tables in the database"""
    engine = create_engine(settings.DATABASE_URL)
    
    print("🔄 Starting database migration...")
    print(f"📊 Database URL: {settings.DATABASE_URL}")
    
    try:
        # Create only the new tables
        print("\n📝 Creating new tables...")
        UserIndicatorPreferences.__table__.create(engine, checkfirst=True)
        print("✅ Created table: user_indicator_preferences")
        
        AIPlanGeneration.__table__.create(engine, checkfirst=True)
        print("✅ Created table: ai_plan_generations")
        
        print("\n✨ Migration completed successfully!")
        print("\n📋 New tables created:")
        print("  - user_indicator_preferences: Store user's preferred technical indicators")
        print("  - ai_plan_generations: Store AI-generated daily trading plans")
        
        # Test connection
        with engine.connect() as conn:
            # Check if tables exist
            result = conn.execute(text("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public' 
                AND table_name IN ('user_indicator_preferences', 'ai_plan_generations')
            """))
            tables = [row[0] for row in result]
            
            print(f"\n✓ Verified tables in database: {', '.join(tables)}")
        
    except Exception as e:
        print(f"\n❌ Migration failed: {str(e)}")
        print("\nPlease check:")
        print("  1. Database is running")
        print("  2. Database credentials are correct in .env")
        print("  3. Database user has CREATE TABLE permissions")
        raise


def rollback_migration():
    """Drop the new tables (use with caution!)"""
    engine = create_engine(settings.DATABASE_URL)
    
    print("⚠️  ROLLBACK: Dropping new tables...")
    
    try:
        UserIndicatorPreferences.__table__.drop(engine, checkfirst=True)
        print("✅ Dropped table: user_indicator_preferences")
        
        AIPlanGeneration.__table__.drop(engine, checkfirst=True)
        print("✅ Dropped table: ai_plan_generations")
        
        print("\n✨ Rollback completed successfully!")
        
    except Exception as e:
        print(f"\n❌ Rollback failed: {str(e)}")
        raise


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Migrate database for indicator preferences and AI plans')
    parser.add_argument('--rollback', action='store_true', help='Rollback migration (drop tables)')
    
    args = parser.parse_args()
    
    if args.rollback:
        confirm = input("⚠️  Are you sure you want to rollback? This will DELETE data! (yes/no): ")
        if confirm.lower() == 'yes':
            rollback_migration()
        else:
            print("Rollback cancelled.")
    else:
        run_migration()