# custom
from core.errors.logger import logger
from core.models.DatabaseOperation import DatabaseOperation, DBErrorType
from core.models.manage.migrations import migrate_sql
from core.models.manage.session_management import get_session

# generic
from sqlalchemy.orm import Session
from sqlalchemy.exc import (
    OperationalError,
    IntegrityError,
    SQLAlchemyError
)
from typing import Type, Dict, Any, Callable, TypeVar
from functools import wraps
from enum import Enum


def safe_db_operation(func):
    def wrapper(*args, **kwargs):

        # get the Session for the function it's wrapping,
        session = get_session()

        try:
            # passes the session to the function,
            result = func(*args, session=session, **kwargs)
            
            # If func already returns DatabaseOperation, pass it through
            if isinstance(result, DatabaseOperation):
                return result
            
            # Otherwise, wrap success
            return DatabaseOperation(valid=True, data=result)
        
        except OperationalError as e:
            logger.error(f"Schema mismatch in {func.__name__}: {e}")
            session.rollback()
            
            try:
                logger.info("Attempting migration...")
                migrate_sql()

                # try again:
                result = func(*args, session=session, **kwargs)
                return DatabaseOperation(
                    valid=True,
                    data=result,
                    tried_migration=True,
                    message="Recovered via migration"
                )
            except Exception as retry_error: # migration failed:
                logger.error(f"Migration failed with error: {str(retry_error)}", exc_info=True)
                session.rollback()
                return DatabaseOperation(
                    valid=False,
                    error_type=DBErrorType.MIGRATION_FAILED,
                    message=str(retry_error)
                )

        # This is really TypeErrors, passed from the insert function.
        except WrapperRollback as e:
            session.rollback()
            return e.database_operation

        except IntegrityError as e:
            error_msg = f"IntegrityError caught: {e.orig}"
            logger.error(error_msg)
            session.rollback()
            return DatabaseOperation(
                valid=False,
                error_type=DBErrorType.INTEGRITY_ERROR,
                message=error_msg
            )

        except SQLAlchemyError as e:
            error_msg = f"SQLAlchemyError caught: {e}"
            logger.error(error_msg)
            session.rollback()
            return DatabaseOperation(
                valid=False,
                error_type=DBErrorType.UNKNOWN,
                message=str(e)
            )

        except Exception as e:
            # Catch-all for unexpected errors
            session.rollback()
            return DatabaseOperation(
                valid=False,
                error_type=DBErrorType.UNKNOWN,
                message=str(e)
            )
    
    return wrapper


class WrapperRollback(Exception):
    """Raised by wrapped function to signal wrapper should rollback and return error."""
    def __init__(self, database_operation: DatabaseOperation):
        self.database_operation = database_operation