from mongoengine.fields import DateField
from datetime import datetime


DEFAULT_DATE_FORMAT = '%b %d, %Y'


class QuarterDateField(DateField):
    def __init__(self, *args, date_format='%b %d, %Y', **kwargs):
        self.date_format = date_format
        super().__init__(*args, **kwargs)

    def validate(self, value):
        if isinstance(value, str):
            try:
                datetime.strptime(value, self.date_format)
            except Exception as ex:
                self.error(f'Date time is not the correct format: (expected {self.date_format}). Error: {str(ex)}')

    def to_mongo(self, value):
        value = super().to_mongo(value)
        if isinstance(value, datetime):
            value = value.strftime(self.date_format)
        return value

    def to_python(self, value):
        value = super().to_python(value)
        if isinstance(value, datetime):
            value = value.strftime(self.date_format)
        return value

    def get_dt_format(self):
        return self.date_format

    @staticmethod
    def make_datetime(value, dt_format='%b %d, %Y'):
        if isinstance(value, str):
            return datetime.strptime(value, dt_format)
        return value


class QuarterYearField(QuarterDateField):
    def __init__(self, *args, **kwargs):
        kwargs['date_format'] = '%Y'
        super().__init__(*args, **kwargs)

    def validate(self, value):
        if isinstance(value, int):
            value = str(value)
        return super().validate(value)

    def to_mongo(self, value):
        if isinstance(value, int):
            value = str(value)
        return super().to_mongo(value)

    def to_python(self, value):
        if isinstance(value, int):
            value = str(value)
        return super().to_python(value)