"""MySQL connection adapter.

Implements the DBConnection protocol for persistent storage operations (Aurora).
"""

from __future__ import annotations

from collections.abc import Callable, Generator
from contextlib import closing, contextmanager
from enum import IntEnum
from functools import wraps
from typing import ContextManager, ParamSpec, TypeVar

import pymysql
from pydantic import BaseModel, ConfigDict, Field
from pymysql.connections import Connection as PyMySQLConnection
from pymysql.cursors import Cursor as PyMySQLCursor
from pymysql.err import OperationalError

from src.connectors.db import Connection
from src.errors import TransientError


class MySQLConfig(BaseModel):
    """Configuration model for a MySQL database connection."""

    model_config = ConfigDict(frozen=True)

    # Connection
    host: str = Field(..., description='Hostname or IP address of the MySQL server.')
    port: int = Field(
        default=3306,
        description='The connection port (e.g., 3306).',
    )
    user: str = Field(..., description='Username for authentication.')
    password: str = Field(..., description='Password for authentication.')
    database: str = Field(..., description='Target database name.')
    connect_timeout: int = Field(
        default=10, description='time to wait for db connection (seconds)'
    )
    # Settings
    autocommit: bool = Field(
        default=False, description='Whether transactions should commit automatically.'
    )
    local_infile: bool = Field(
        default=False, description='Enable LOAD DATA LOCAL INFILE capability.'
    )


class MySQLConnection(Connection):
    """MySQL database connection adapter."""

    def __init__(self, conn: PyMySQLConnection):
        """Initialize the adapter with a raw pymysql connection.

        Args:
            conn: An active pymysql connection.
        """
        self._conn = conn

    def close(self) -> None:
        """Close the connection."""
        self._conn.close()

    def commit(self) -> None:
        """Commit the current transaction to the database."""
        return self._conn.commit()

    def cursor(self) -> ContextManager[PyMySQLCursor]:  # type: ignore
        """Return a raw pymysql cursor for low-level use."""
        return closing(self._conn.cursor())

    # pd.read_sql(sql, self._conn, params=params, dtype_backend='pyarrow')

    def rollback(self) -> None:
        """Roll back the current transaction, discarding changes."""
        return self._conn.rollback()


class MySQLConnectionFactory:
    """Factory for creating MySQL database connections with configuration management."""

    def __init__(self, config: MySQLConfig) -> None:
        """Initialize the connection factory with MySQL configuration.

        Args:
            config: MySQLConfig instance containing database connection params
        """
        self._config = config

    @contextmanager
    def connection(self) -> Generator[MySQLConnection, None, None]:
        """Create and yield a MySQL database connection as a context manager.

        The connection is automatically closed when the context exits.

        Yields:
            MySQLConnection: MySQL database adapter.

        Raises:
            TransientError: If connection fails due to transient database errors.
            OperationalError: If connection fails due to non-transient errors.
        """
        with mysql_connection(self._config) as conn:
            yield MySQLConnection(conn)


class MySQLErrorCode(IntEnum):
    """Known MySQL error codes, mostly related to transient errors."""

    # Lock wait timeout exceeded; try restarting transaction
    ER_LOCK_WAIT_TIMEOUT = 1205
    # Deadlock found when trying to get lock; try restarting transaction
    ER_LOCK_DEADLOCK = 1213
    # Can't connect to MySQL server
    CR_CONN_HOST_ERROR = 2003
    # MySQL server has gone away
    CR_SERVER_GONE_ERROR = 2006
    # Lost connection to MySQL server during query
    CR_SERVER_LOST = 2013
    # Lost connection to MySQL server
    CR_SERVER_LOST_EXTENDED = 2055


P = ParamSpec('P')
T = TypeVar('T')


TRANSIENT_ERROR_CODES: frozenset[int] = frozenset(
    {
        MySQLErrorCode.ER_LOCK_WAIT_TIMEOUT,
        MySQLErrorCode.ER_LOCK_DEADLOCK,
        MySQLErrorCode.CR_CONN_HOST_ERROR,
        MySQLErrorCode.CR_SERVER_GONE_ERROR,
        MySQLErrorCode.CR_SERVER_LOST,
        MySQLErrorCode.CR_SERVER_LOST_EXTENDED,
    }
)


def _try_transient_error(e: OperationalError) -> None:
    """Check if an OperationalError is caused by a transient condition.

    If transient, a custom TransientError is raised.

    Args:
        e: The caught pymysql.err.OperationalError exception.

    Raises:
        TransientError: If the error code matches one in TRANSIENT_ERROR_CODES.
    """
    if e.args and e.args[0] in TRANSIENT_ERROR_CODES:
        raise TransientError(f'Database operation failed: {e.args[1]}') from e


def handle_mysql_errors(func: Callable[P, T]) -> Callable[P, T]:
    """Decorate database operations and handle transient errors.

    If a transient error occurs, it is converted into a TransientError to
    allow for retry logic. Non-transient errors are re-raised as is.
    """

    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        try:
            return func(*args, **kwargs)
        except OperationalError as e:
            _try_transient_error(e)
            raise  # Re-raise if it wasn't transient

    return wrapper


@contextmanager
def mysql_connection(config: MySQLConfig) -> Generator[PyMySQLConnection, None, None]:
    """Context manager for connection objects.

    Args:
        config: MySQL configuration object.

    Yields:
        Connection: A raw active pymysql connection.

    Raises:
        TransientError: If the connection attempt fails due to a transient error.
        OperationalError: If the connection attempt fails due to a non-transient error.
    """
    connect_timeout = config.connect_timeout
    if connect_timeout is None:
        connect_timeout = 10

    port = config.port if config.port is not None else 3306
    autocommit = config.autocommit if config.autocommit is not None else False
    local_infile = config.local_infile if config.local_infile is not None else False

    try:
        conn = pymysql.connect(
            host=config.host,
            user=config.user,
            password=config.password,
            database=config.database,
            connect_timeout=connect_timeout,
            port=port,
            cursorclass=PyMySQLCursor,
            autocommit=autocommit,
            local_infile=local_infile,
        )
    except OperationalError as e:
        _try_transient_error(e)
        raise

    try:
        yield conn
    finally:
        conn.close()