#!/usr/bin/env python3

import os
import logging
from peewee import PostgresqlDatabase

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('postgres_optimization')

def optimize_postgresql():
    """Apply PostgreSQL optimizations to the migrated database"""
    
    # Connect to PostgreSQL
    db = PostgresqlDatabase(
        'sba_master',
        user='sba_admin', 
        password='sba_dev_password_2024',
        host='localhost',
        port=5432
    )
    
    try:
        db.connect()
        logger.info("✓ Connected to PostgreSQL database")
        
        # Read and execute optimization SQL
        with open('optimize_postgres.sql', 'r') as f:
            sql_commands = f.read()
        
        # Split into individual commands and execute
        commands = [cmd.strip() for cmd in sql_commands.split(';') if cmd.strip() and not cmd.strip().startswith('--')]
        
        successful_commands = 0
        total_commands = len(commands)
        
        for i, command in enumerate(commands, 1):
            try:
                if command.lower().startswith(('create index', 'analyze', 'commit')):
                    logger.info(f"Executing command {i}/{total_commands}: {command[:50]}...")
                    db.execute_sql(command)
                    successful_commands += 1
                    logger.info(f"✓ Command {i} completed successfully")
                else:
                    logger.info(f"Skipping command {i}: {command[:50]}...")
                    
            except Exception as e:
                if "already exists" in str(e).lower():
                    logger.info(f"⚠ Command {i} - Index already exists, skipping")
                    successful_commands += 1
                else:
                    logger.error(f"✗ Command {i} failed: {e}")
        
        logger.info(f"Optimization completed: {successful_commands}/{total_commands} commands successful")
        
        # Final statistics update
        try:
            db.execute_sql("ANALYZE;")
            logger.info("✓ Database statistics updated")
        except Exception as e:
            logger.error(f"✗ Failed to update statistics: {e}")
        
        db.close()
        return True
        
    except Exception as e:
        logger.error(f"✗ Database optimization failed: {e}")
        try:
            db.close()
        except:
            pass
        return False

if __name__ == "__main__":
    logger.info("=== PostgreSQL Database Optimization ===")
    success = optimize_postgresql()
    
    if success:
        logger.info("🚀 Database optimization completed successfully")
        logger.info("Recommended next steps:")
        logger.info("  1. Test query performance with sample API calls")
        logger.info("  2. Monitor query execution plans with EXPLAIN ANALYZE") 
        logger.info("  3. Adjust work_mem and other settings based on usage patterns")
    else:
        logger.error("❌ Database optimization failed")
    
    exit(0 if success else 1)