from sqlalchemy import and_
import pendulum, pytz


# Custom functions
def get_by_partnership_account(cls, paid, count_only=False, from_date=None, to_date=None):
    """
    Get the result_objects in a particular partnership account of id: `paid`.
    if count_only = True, return the count instead of objects.
    """
    if from_date and to_date:
        results = cls.query\
            .filter(cls.partnership_account_id == paid)\
            .filter(and_(cls.created_on >= from_date, cls.created_on <= to_date ))
    elif from_date:
        results = cls.query\
            .filter(cls.partnership_account_id == paid)\
            .filter(cls.created_on >= from_date)
    elif to_date:
        results = cls.query\
            .filter(cls.partnership_account_id == paid)\
            .filter(cls.created_on <= to_date)
    else:
        results = cls.query\
            .filter(cls.partnership_account_id == paid)
            
    return results.count() if count_only else results.all() 


def get_by_partnership_account_expr(cls, paid, count_only=False, date_range_expr=None):
    """
    :param date_range_expr: tpye:string, possible values: this_month, last_month
    """

    if date_range_expr:
        if date_range_expr == 'this_month':
            today = pendulum.today(tz=pytz.utc)
            first_day = today.replace(day=1)
            return get_by_partnership_account(cls, paid, count_only, first_day, today)
        
        elif date_range_expr == 'last_month':
            today = pendulum.today(tz=pytz.utc)
            pre_month = today.month - 1
            first_day_this_month = today.replace(day=1)
            first_day_pre_month = today.replace(month=pre_month, day=1)
            last_day_pre_month = first_day_this_month  # Time considers as 00:00:00
            return get_by_partnership_account(cls, paid, count_only, first_day_pre_month, last_day_pre_month)
        else:
            return get_by_partnership_account(cls, paid, count_only)
    else:
        return get_by_partnership_account(cls, paid, count_only)