"""Helper functions for multiprocessing."""
import os
import logging
from sqlalchemy import event
from sqlalchemy import exc
import pymysql

logger = logging.getLogger(__name__)


def add_engine_pidguard(engine):
    """Add multiprocessing guards.

    Forces a connection to be reconnected if it is detected
    as having been shared to a sub-process.
    """
    @event.listens_for(engine, 'connect')
    def connect(dbapi_connection, connection_record):
        connection_record.info['pid'] = os.getpid()

    @event.listens_for(engine, 'checkout')
    def checkout(dbapi_connection, connection_record, connection_proxy):
        pid = os.getpid()
        conn_pid = connection_record.info.get('pid')
        if conn_pid != pid:
            logger.warning(
                f'Parent process {conn_pid} forked ({pid}) with an open '
                'database connection, which is being discarded and recreated.'
            )

            try:
                # Check if the connection is already closed
                if not dbapi_connection.open:
                    logger.info(
                        "Connection already closed, skipping close operation")
                else:
                    # Close the connection to ensure resources are freed
                    dbapi_connection.close()
            except (pymysql.err.Error, AttributeError) as e:
                logger.warning(f"Error closing connection: {str(e)}")
            except Exception as e:
                logger.exception(
                    f"Unexpected error closing connection: {str(e)}")

            try:
                # Use SQLAlchemy's built-in invalidation method
                connection_record.invalidate()
            except Exception as e:
                logger.exception(f"Error invalidating connection: {str(e)}")

            # Raise DisconnectionError to trigger a reconnection
            raise exc.DisconnectionError(
                f'Connection record belongs to pid {conn_pid}, '
                f'attempting to check out in pid {pid}'
            )

        logger.debug(f'Successfully checked out connection for pid {pid}')
        return dbapi_connection

    @event.listens_for(engine, 'checkin')
    def checkin(dbapi_connection, connection_record):
        # Update the PID on checkin, in case the connection is reused
        new_pid = os.getpid()
        old_pid = connection_record.info.get('pid', 'Unknown')
        connection_record.info['pid'] = new_pid
        logger.debug(
            f'Checked in connection. Old pid: {old_pid}, New pid: {new_pid}')