#!/usr/bin/env python
"""
Script to fix database connection issues in production
This script directly modifies the SQLAlchemy database URI at runtime
"""
import os
import sys
from sqlalchemy import create_engine, text

# Set the correct database URL directly in the environment
# PyMySQL doesn't support ssl_mode parameter directly in the URL
# Instead, we need to use ssl=true and specify SSL options in the connect_args
os.environ['DATABASE_URL'] = 'mysql+pymysql://avnadmin:AVNS_fCi3sA9701sLkgs-0F-@mysql-35c14f29-creations-3584.i.aivencloud.com:23574/sqlgenai?ssl=true'

# Import Flask app after setting environment variable
from app import create_app
from app.utils.db_init import db

def test_connection():
    """Test the database connection with the new URL"""
    print(f"Testing connection to: {os.environ['DATABASE_URL']}")
    try:
        # Create engine with SSL configuration
        # For Aiven MySQL, we need to set ssl=True but don't need to specify cert paths
        engine = create_engine(
            os.environ['DATABASE_URL'],
            connect_args={
                'ssl': {
                    'ca': None,  # No CA certificate needed
                    'check_hostname': False,  # Skip hostname verification
                    'verify_mode': None  # Skip certificate verification
                }
            }
        )
        
        # Try to connect
        with engine.connect() as conn:
            result = conn.execute(text("SELECT 1"))
            print("✅ Connection successful!")
            print(f"Result: {result.fetchone()}")
        return True
    except Exception as e:
        print(f"❌ Connection failed: {str(e)}")
        return False

def update_env_file():
    """Update the .env file with the correct database URL"""
    try:
        # Read the current .env file
        with open('.env', 'r') as f:
            lines = f.readlines()
        
        # Update the DATABASE_URL line
        new_lines = []
        db_url_updated = False
        db_config_added = False
        
        for line in lines:
            if line.startswith('DATABASE_URL='):
                new_lines.append(f"DATABASE_URL={os.environ['DATABASE_URL']}\n")
                # Add a comment explaining SSL configuration
                new_lines.append("# Note: SSL configuration is handled in the code via connect_args\n")
                new_lines.append("# For SQLAlchemy with PyMySQL, we need to set ssl=true in the URL\n")
                new_lines.append("# and configure SSL options in the engine creation\n")
                db_url_updated = True
            else:
                new_lines.append(line)
        
        # Add the DATABASE_URL if it doesn't exist
        if not db_url_updated:
            new_lines.append(f"\n# Database configuration - Aiven Cloud MySQL\nDATABASE_URL={os.environ['DATABASE_URL']}\n")
        
        # Write the updated .env file
        with open('.env', 'w') as f:
            f.writelines(new_lines)
        
        print("✅ .env file updated with correct database URL")
        return True
    except Exception as e:
        print(f"❌ Failed to update .env file: {str(e)}")
        return False

def create_db_config_file():
    """Create a database configuration file with proper SSL settings"""
    try:
        config_content = """
# SQLGenAI Database Configuration
# This file is automatically generated to configure database connections

from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker

# Database URL
DATABASE_URL = '{db_url}'

# Create engine with proper SSL configuration for Aiven MySQL
engine = create_engine(
    DATABASE_URL,
    connect_args={{
        'ssl': {{
            'ca': None,
            'check_hostname': False,
            'verify_mode': None
        }}
    }}
)

# Create session factory
db_session = scoped_session(
    sessionmaker(autocommit=False, autoflush=False, bind=engine)
)

# Function to initialize database
def init_db():
    from app.utils.db_models import Base
    Base.metadata.create_all(bind=engine)

# Function to get database session
def get_db():
    return db_session
""".format(db_url=os.environ['DATABASE_URL'])

        # Write the configuration file
        with open('app/utils/db_config.py', 'w') as f:
            f.write(config_content)
        
        print("✅ Created database configuration file with proper SSL settings")
        return True
    except Exception as e:
        print(f"❌ Failed to create database configuration file: {str(e)}")
        return False

def main():
    """Main function to fix database connection"""
    print("SQLGenAI Database Connection Fix")
    print("================================")
    
    # Test the connection
    if test_connection():
        # Update the .env file
        env_updated = update_env_file()
        
        # Create database configuration file
        config_created = create_db_config_file()
        
        if env_updated and config_created:
            print("\n✅ Database connection fixed!")
            print("\nTo use this configuration in your Flask app:")
            print("1. Import the database session from the new config file:")
            print("   from app.utils.db_config import db_session, init_db")
            print("\n2. Update your app/__init__.py to use this configuration")
            print("\n3. Restart your application:")
            print("   sudo systemctl restart apache2")
        else:
            print("\n⚠️ Connection works but some files could not be updated.")
            print("Please manually update your configuration with:")
            print(f"DATABASE_URL={os.environ['DATABASE_URL']}")
            print("And configure SSL options in your SQLAlchemy engine creation")
    else:
        print("\n❌ Connection failed even with the new URL.")
        print("Please check:")
        print("1. Network connectivity to the database server")
        print("2. Firewall settings")
        print("3. Database credentials")
        print("4. SSL requirements")

if __name__ == "__main__":
    main()