"""
Async database layer using aiomysql.

Provides connection pooling, query helpers, and proper resource management.
"""

import asyncio
import aiomysql
from typing import Optional, Any, Tuple, List
import logging
from config import config

logger = logging.getLogger(__name__)

# Global connection pool
_pool: Optional[aiomysql.Pool] = None
_pool_lock = asyncio.Lock()


async def initialize_pool(
    pool_size: int = None,
    max_overflow: int = 10
) -> aiomysql.Pool:
    """
    Initialize the async database connection pool.

    Args:
        pool_size: Maximum pool size (default from config.MAX_CONNECTIONS)
        max_overflow: Additional connections beyond pool_size

    Returns:
        The connection pool
    """
    global _pool

    async with _pool_lock:
        if _pool is not None:
            logger.warning("Pool already initialized")
            return _pool

        pool_size = pool_size or config.MAX_CONNECTIONS

        try:
            _pool = await aiomysql.create_pool(
                host=config.DB_HOST,
                port=config.DB_PORT,
                user=config.DB_USER,
                password=config.DB_PASSWORD,
                db=config.DB_NAME,
                minsize=1,  # Start with minimal connections to avoid exhaustion
                maxsize=pool_size,
                autocommit=True,  # Auto-commit by default
                pool_recycle=3600,  # Recycle connections after 1 hour
                echo=config.DEBUG,
                connect_timeout=10,  # 10 second connection timeout
            )

            logger.info(f"Database pool initialized: size={pool_size}, host={config.DB_HOST}")
            return _pool

        except Exception as e:
            logger.error(f"Failed to initialize database pool: {e}")
            raise


async def close_pool() -> None:
    """Close the database connection pool"""
    global _pool

    async with _pool_lock:
        if _pool is not None:
            _pool.close()
            await _pool.wait_closed()
            _pool = None
            logger.info("Database pool closed")


async def get_connection() -> aiomysql.Connection:
    """
    Get a connection from the pool.

    Usage:
        async with get_connection() as conn:
            async with conn.cursor() as cursor:
                await cursor.execute("SELECT ...")

    Returns:
        Database connection (context manager)
    """
    if _pool is None:
        raise RuntimeError("Database pool not initialized. Call initialize_pool() first.")

    return _pool.acquire()


async def execute_query(
    query: str,
    params: Optional[Tuple] = None,
    commit: bool = True
) -> int:
    """
    Execute a query (INSERT, UPDATE, DELETE).

    Args:
        query: SQL query with %s placeholders
        params: Query parameters
        commit: Whether to commit (default True due to autocommit)

    Returns:
        Number of affected rows
    """
    if _pool is None:
        raise RuntimeError("Database pool not initialized")

    async with _pool.acquire() as conn:
        async with conn.cursor() as cursor:
            await cursor.execute(query, params)
            return cursor.rowcount


async def execute_many(
    query: str,
    params_list: List[Tuple]
) -> int:
    """
    Execute a query multiple times with different parameters (batch insert).

    Args:
        query: SQL query with %s placeholders
        params_list: List of parameter tuples

    Returns:
        Total number of affected rows
    """
    if _pool is None:
        raise RuntimeError("Database pool not initialized")

    async with _pool.acquire() as conn:
        async with conn.cursor() as cursor:
            await cursor.executemany(query, params_list)
            return cursor.rowcount


async def fetch_one(
    query: str,
    params: Optional[Tuple] = None
) -> Optional[Tuple]:
    """
    Fetch a single row.

    Args:
        query: SQL query with %s placeholders
        params: Query parameters

    Returns:
        Single row as tuple, or None if not found
    """
    if _pool is None:
        raise RuntimeError("Database pool not initialized")

    async with _pool.acquire() as conn:
        async with conn.cursor() as cursor:
            await cursor.execute(query, params)
            return await cursor.fetchone()


async def fetch_all(
    query: str,
    params: Optional[Tuple] = None
) -> List[Tuple]:
    """
    Fetch all rows.

    Args:
        query: SQL query with %s placeholders
        params: Query parameters

    Returns:
        List of rows as tuples
    """
    if _pool is None:
        raise RuntimeError("Database pool not initialized")

    async with _pool.acquire() as conn:
        async with conn.cursor() as cursor:
            await cursor.execute(query, params)
            return await cursor.fetchall()


async def fetch_dict_one(
    query: str,
    params: Optional[Tuple] = None
) -> Optional[dict]:
    """
    Fetch a single row as dictionary.

    Args:
        query: SQL query with %s placeholders
        params: Query parameters

    Returns:
        Single row as dict, or None if not found
    """
    if _pool is None:
        raise RuntimeError("Database pool not initialized")

    async with _pool.acquire() as conn:
        async with conn.cursor(aiomysql.DictCursor) as cursor:
            await cursor.execute(query, params)
            return await cursor.fetchone()


async def fetch_dict_all(
    query: str,
    params: Optional[Tuple] = None
) -> List[dict]:
    """
    Fetch all rows as dictionaries.

    Args:
        query: SQL query with %s placeholders
        params: Query parameters

    Returns:
        List of rows as dicts
    """
    if _pool is None:
        raise RuntimeError("Database pool not initialized")

    async with _pool.acquire() as conn:
        async with conn.cursor(aiomysql.DictCursor) as cursor:
            await cursor.execute(query, params)
            return await cursor.fetchall()


class Transaction:
    """
    Context manager for database transactions.

    Usage:
        async with Transaction() as conn:
            async with conn.cursor() as cursor:
                await cursor.execute("UPDATE ...")
                await cursor.execute("INSERT ...")
            # Auto-commit on success, rollback on exception
    """

    def __init__(self):
        self.conn = None

    async def __aenter__(self) -> aiomysql.Connection:
        if _pool is None:
            raise RuntimeError("Database pool not initialized")

        self.conn = await _pool.acquire()
        await self.conn.begin()
        return self.conn

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            # Exception occurred, rollback
            await self.conn.rollback()
            logger.warning(f"Transaction rolled back due to {exc_type.__name__}")
        else:
            # Success, commit
            await self.conn.commit()

        # Return connection to pool
        await _pool.release(self.conn)
        self.conn = None