"""MySQL DB Connector."""

from __future__ import annotations

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

import pymysql
from pymysql.connections import Connection
from pymysql.cursors import DictCursor
from pymysql.err import OperationalError

from src.errors import TransientError

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


class MySQLErrorCode(IntEnum):
    """MySQL error codes."""

    # 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


TRANSIENT_ERROR_CODES: set[MySQLErrorCode] = {
    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 handle_mysql_errors(func: Callable[P, T]) -> Callable[P, T]:
    """Decorate to handle transient database errors.

    Catches OperationalError for connection issues
    and converts them to TransientError.
    """

    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        try:
            return func(*args, **kwargs)
        except OperationalError as e:
            # Retriable error codes
            if e.args[0] in TRANSIENT_ERROR_CODES:
                raise TransientError(f'Database connection failed: {e.args[1]}') from e
            # Other errors
            raise

    return wrapper


@contextmanager
def mysql_connection(
    host: str,
    user: str,
    password: str,
    database: str,
    connect_timeout: int = 5,
    port: int = 3306,
    autocommit: bool = False,
) -> Generator[Connection[DictCursor], None, None]:
    """Context manager for mysql connection objects.

    Args:
        host (str): hostname of database server
        user (str): user name
        password (str): password
        database (str): database name
        connect_timeout (int): time to wait for db connection
        port (int): the port of the sql connection
        autocommit (bool): indicates whether transactions
            should be committed automatically

    Yields:
        Connection: connection to the specified database
    """
    conn = None
    try:
        conn = pymysql.connect(
            host=host,
            user=user,
            password=password,
            database=database,
            connect_timeout=connect_timeout,
            port=port,
            cursorclass=DictCursor,
            autocommit=autocommit,
        )

        yield conn
    finally:
        if conn:
            conn.close()