import aiomysql
from typing import Optional
from common.config.app_settings import app_settings
from common.log.module_logger import ModuleLogger


class DatabaseConnectionPool:
    """Database connection pool manager for StarRocks"""
    
    def __init__(self):
        self.pool: Optional[aiomysql.Pool] = None
        self.module_logger = ModuleLogger(__file__)
    
    async def create_pool(self) -> bool:
        """Create database connection pool"""
        try:
            self.pool = await aiomysql.create_pool(
                host=app_settings.STARROCKS_HOST,
                port=app_settings.STARROCKS_PORT,
                user=app_settings.STARROCKS_USER,
                password=app_settings.STARROCKS_PASSWORD,
                db=app_settings.STARROCKS_DATABASE,
                charset='utf8mb4',
                autocommit=True,
                minsize=5,  # Minimum number of connections in the pool
                maxsize=20,  # Maximum number of connections in the pool
                pool_recycle=3600,  # Recycle connections after 1 hour
                echo=False  # Set to True for SQL query logging
            )
            await self.module_logger.log_info(
                f"Database connection pool created successfully for StarRocks at "
                f"{app_settings.STARROCKS_HOST}:{app_settings.STARROCKS_PORT}"
            )
            return True
        except Exception as e:
            await self.module_logger.log_error(f"Failed to create database connection pool: {e}")
            return False
    
    async def close_pool(self):
        """Close database connection pool"""
        if self.pool:
            self.pool.close()
            await self.pool.wait_closed()
            self.pool = None
            await self.module_logger.log_info("Database connection pool closed")
    
    def get_pool(self) -> aiomysql.Pool:
        """Get the database connection pool"""
        if not self.pool:
            raise Exception("Database connection pool not initialized")
        return self.pool


# Global database connection pool instance
db_pool = DatabaseConnectionPool()


def register(app):
    """Register database provider with FastAPI application"""
    
    @app.on_event("startup")
    async def startup():
        """Initialize database connection pool on application startup"""
        await db_pool.create_pool()
    
    @app.on_event("shutdown")
    async def shutdown():
        """Close database connection pool on application shutdown"""
        await db_pool.close_pool()


def boot(app):
    """Boot database provider (if needed for additional setup)"""
    pass