# 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
from core.Constants import Constants

# 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

"""
Purpose:
    This is a wrapper for other SQL functions to catch errors and do solutions/fixes.

Used/Called by:
    insert_into_model

Returns:
    DatabaseOperation Objects
"""

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


"""
If a wrapped function needs to pass up an error to this wrapper,
so that the wrapper can roll back sessions,

Then it uses this error exception and passes the DatabaseOperation object.
"""
class WrapperRollback(Exception):
    def __init__(self, database_operation: DatabaseOperation):
        self.database_operation = database_operation