"""Decorators that handle model-related errors."""

import functools

from oto import response
from sqlalchemy import exc

from sentry_sdk import capture_exception
from sales_goals.constants import error


def integrity_error_handler(error_code):
    """Decorator that handles IntegrityError from SQLAlchemy.

    Args:
        error_code (str): Error code to add to response.Response object.
    """
    def wrap(function):
        @functools.wraps(function)
        def wrapper(*args, **kwargs):
            try:
                return function(*args, **kwargs)
            except exc.IntegrityError as e:
                message = str(e)
                # here we check the first argument of origin args tuple
                if e.orig.args[0] == error.ERROR_MYSQL_DUPLICATE_ENTRY:
                    message = error.ERROR_MESSAGE_DUPLICATE_ENTRY
                return response.create_error_response(
                    code=error_code,
                    message=message)
        return wrapper
    return wrap


def sqlalchemy_error_handler(function):
    """Decorator that handles SQLAlchemyError from SQLAlchemy."""
    @functools.wraps(function)
    def wrapper(*args, **kwargs):
        try:
            return function(*args, **kwargs)
        except exc.SQLAlchemyError as e:
            capture_exception(e)
            return response.create_fatal_response(str(e))
    return wrapper