1import datetime
  2from calendar import monthrange
  3from enum import Enum
  4
  5from plain.utils import timezone
  6
  7
  8class DatetimeRangeAliases(Enum):
  9    TODAY = "Today"
 10    THIS_WEEK = "This Week"
 11    THIS_WEEK_TO_DATE = "This Week-to-date"
 12    THIS_MONTH = "This Month"
 13    THIS_MONTH_TO_DATE = "This Month-to-date"
 14    THIS_QUARTER = "This Quarter"
 15    THIS_QUARTER_TO_DATE = "This Quarter-to-date"
 16    THIS_YEAR = "This Year"
 17    THIS_YEAR_TO_DATE = "This Year-to-date"
 18    LAST_WEEK = "Last Week"
 19    LAST_WEEK_TO_DATE = "Last Week-to-date"
 20    LAST_MONTH = "Last Month"
 21    LAST_MONTH_TO_DATE = "Last Month-to-date"
 22    LAST_QUARTER = "Last Quarter"
 23    LAST_QUARTER_TO_DATE = "Last Quarter-to-date"
 24    LAST_YEAR = "Last Year"
 25    LAST_YEAR_TO_DATE = "Last Year-to-date"
 26    SINCE_30_DAYS_AGO = "Since 30 Days Ago"
 27    SINCE_60_DAYS_AGO = "Since 60 Days Ago"
 28    SINCE_90_DAYS_AGO = "Since 90 Days Ago"
 29    SINCE_365_DAYS_AGO = "Since 365 Days Ago"
 30    NEXT_WEEK = "Next Week"
 31    NEXT_4_WEEKS = "Next 4 Weeks"
 32    NEXT_MONTH = "Next Month"
 33    NEXT_QUARTER = "Next Quarter"
 34    NEXT_YEAR = "Next Year"
 35
 36    # TODO doesn't include anything less than a day...
 37    # ex. SINCE_1_HOUR_AGO = "Since 1 Hour Ago"
 38
 39    def __str__(self):
 40        return self.value
 41
 42    @classmethod
 43    def from_value(cls, value):
 44        for member in cls:
 45            if member.value == value:
 46                return member
 47        raise ValueError(f"{value} is not a valid value for {cls.__name__}")
 48
 49    @classmethod
 50    def to_range(cls, value: str) -> tuple[datetime.datetime, datetime.datetime]:
 51        now = timezone.localtime()
 52        start_of_today = now.replace(hour=0, minute=0, second=0, microsecond=0)
 53        start_of_week = start_of_today - datetime.timedelta(
 54            days=start_of_today.weekday()
 55        )
 56        start_of_month = start_of_today.replace(day=1)
 57        start_of_quarter = start_of_today.replace(
 58            month=((start_of_today.month - 1) // 3) * 3 + 1, day=1
 59        )
 60        start_of_year = start_of_today.replace(month=1, day=1)
 61
 62        def end_of_day(dt):
 63            return dt.replace(hour=23, minute=59, second=59, microsecond=999999)
 64
 65        def end_of_month(dt):
 66            last_day = monthrange(dt.year, dt.month)[1]
 67            return end_of_day(dt.replace(day=last_day))
 68
 69        def end_of_quarter(dt):
 70            end_month = ((dt.month - 1) // 3 + 1) * 3
 71            return end_of_month(dt.replace(month=end_month))
 72
 73        def end_of_year(dt):
 74            return end_of_month(dt.replace(month=12))
 75
 76        if value == cls.TODAY:
 77            return DatetimeRange(start_of_today, end_of_day(now))
 78        if value == cls.THIS_WEEK:
 79            return DatetimeRange(
 80                start_of_week, end_of_day(start_of_week + datetime.timedelta(days=6))
 81            )
 82        if value == cls.THIS_WEEK_TO_DATE:
 83            return DatetimeRange(start_of_week, now)
 84        if value == cls.THIS_MONTH:
 85            return DatetimeRange(start_of_month, end_of_month(start_of_month))
 86        if value == cls.THIS_MONTH_TO_DATE:
 87            return DatetimeRange(start_of_month, now)
 88        if value == cls.THIS_QUARTER:
 89            return DatetimeRange(start_of_quarter, end_of_quarter(start_of_quarter))
 90        if value == cls.THIS_QUARTER_TO_DATE:
 91            return DatetimeRange(start_of_quarter, now)
 92        if value == cls.THIS_YEAR:
 93            return DatetimeRange(start_of_year, end_of_year(start_of_year))
 94        if value == cls.THIS_YEAR_TO_DATE:
 95            return DatetimeRange(start_of_year, now)
 96        if value == cls.LAST_WEEK:
 97            last_week_start = start_of_week - datetime.timedelta(days=7)
 98            return DatetimeRange(
 99                last_week_start,
100                end_of_day(last_week_start + datetime.timedelta(days=6)),
101            )
102        if value == cls.LAST_WEEK_TO_DATE:
103            return DatetimeRange(start_of_week - datetime.timedelta(days=7), now)
104        if value == cls.LAST_MONTH:
105            last_month = (start_of_month - datetime.timedelta(days=1)).replace(day=1)
106            return DatetimeRange(last_month, end_of_month(last_month))
107        if value == cls.LAST_MONTH_TO_DATE:
108            last_month = (start_of_month - datetime.timedelta(days=1)).replace(day=1)
109            return DatetimeRange(last_month, now)
110        if value == cls.LAST_QUARTER:
111            last_quarter = (start_of_quarter - datetime.timedelta(days=1)).replace(
112                day=1
113            )
114            return DatetimeRange(last_quarter, end_of_quarter(last_quarter))
115        if value == cls.LAST_QUARTER_TO_DATE:
116            last_quarter = (start_of_quarter - datetime.timedelta(days=1)).replace(
117                day=1
118            )
119            return DatetimeRange(last_quarter, now)
120        if value == cls.LAST_YEAR:
121            last_year = start_of_year.replace(year=start_of_year.year - 1)
122            return DatetimeRange(last_year, end_of_year(last_year))
123        if value == cls.LAST_YEAR_TO_DATE:
124            last_year = start_of_year.replace(year=start_of_year.year - 1)
125            return DatetimeRange(last_year, now)
126        if value == cls.SINCE_30_DAYS_AGO:
127            return DatetimeRange(now - datetime.timedelta(days=30), now)
128        if value == cls.SINCE_60_DAYS_AGO:
129            return DatetimeRange(now - datetime.timedelta(days=60), now)
130        if value == cls.SINCE_90_DAYS_AGO:
131            return DatetimeRange(now - datetime.timedelta(days=90), now)
132        if value == cls.SINCE_365_DAYS_AGO:
133            return DatetimeRange(now - datetime.timedelta(days=365), now)
134        if value == cls.NEXT_WEEK:
135            next_week_start = start_of_week + datetime.timedelta(days=7)
136            return DatetimeRange(
137                next_week_start,
138                end_of_day(next_week_start + datetime.timedelta(days=6)),
139            )
140        if value == cls.NEXT_4_WEEKS:
141            return DatetimeRange(now, end_of_day(now + datetime.timedelta(days=28)))
142        if value == cls.NEXT_MONTH:
143            next_month = (start_of_month + datetime.timedelta(days=31)).replace(day=1)
144            return DatetimeRange(next_month, end_of_month(next_month))
145        if value == cls.NEXT_QUARTER:
146            next_quarter = (start_of_quarter + datetime.timedelta(days=90)).replace(
147                day=1
148            )
149            return DatetimeRange(next_quarter, end_of_quarter(next_quarter))
150        if value == cls.NEXT_YEAR:
151            next_year = start_of_year.replace(year=start_of_year.year + 1)
152            return DatetimeRange(next_year, end_of_year(next_year))
153        raise ValueError(f"Invalid range: {value}")
154
155
156class DatetimeRange:
157    def __init__(self, start, end):
158        self.start = start
159        self.end = end
160
161        if isinstance(self.start, str) and self.start:
162            self.start = datetime.datetime.fromisoformat(self.start)
163
164        if isinstance(self.end, str) and self.end:
165            self.end = datetime.datetime.fromisoformat(self.end)
166
167        if isinstance(self.start, datetime.date) and not isinstance(
168            self.start, datetime.datetime
169        ):
170            self.start = timezone.localtime().replace(
171                year=self.start.year, month=self.start.month, day=self.start.day
172            )
173
174        if isinstance(self.end, datetime.date) and not isinstance(
175            self.start, datetime.datetime
176        ):
177            self.end = timezone.localtime().replace(
178                year=self.end.year, month=self.end.month, day=self.end.day
179            )
180
181    def as_tuple(self):
182        return (self.start, self.end)
183
184    def total_days(self):
185        return (self.end - self.start).days
186
187    def iter_days(self):
188        """Yields each day in the range."""
189        return iter(
190            self.start.date() + datetime.timedelta(days=i)
191            for i in range(0, self.total_days())
192        )
193
194    def iter_weeks(self):
195        """Yields the start of each week in the range."""
196        current = self.start - datetime.timedelta(days=self.start.weekday())
197        current = current.replace(hour=0, minute=0, second=0, microsecond=0)
198        while current <= self.end:
199            next_week = current + datetime.timedelta(weeks=1)
200            yield current
201            current = next_week
202
203    def iter_months(self):
204        """Yields the start of each month in the range."""
205        current = self.start.replace(day=1)
206        current = current.replace(hour=0, minute=0, second=0, microsecond=0)
207        while current <= self.end:
208            if current.month == 12:
209                next_month = current.replace(year=current.year + 1, month=1)
210            else:
211                next_month = current.replace(month=current.month + 1)
212            yield current
213            current = next_month
214
215    def iter_quarters(self):
216        """Yields the start of each quarter in the range."""
217        current = self.start.replace(month=((self.start.month - 1) // 3) * 3 + 1, day=1)
218        current = current.replace(hour=0, minute=0, second=0, microsecond=0)
219        while current <= self.end:
220            next_quarter_month = ((current.month - 1) // 3 + 1) * 3 + 1
221            if next_quarter_month > 12:
222                next_quarter_month -= 12
223                next_year = current.year + 1
224            else:
225                next_year = current.year
226            next_quarter = datetime.datetime(
227                next_year, next_quarter_month, 1, tzinfo=current.tzinfo
228            )
229            yield current
230            current = next_quarter
231
232    def iter_years(self):
233        """Yields the start of each year in the range."""
234        current = self.start.replace(month=1, day=1)
235        current = current.replace(hour=0, minute=0, second=0, microsecond=0)
236        while current <= self.end:
237            next_year = current.replace(year=current.year + 1)
238            yield current
239            current = next_year
240
241    def __repr__(self):
242        return f"DatetimeRange({self.start}, {self.end})"
243
244    def __str__(self):
245        return f"{self.start} to {self.end}"
246
247    def __eq__(self, other):
248        return self.start == other.start and self.end == other.end
249
250    def __hash__(self):
251        return hash((self.start, self.end))
252
253    def __contains__(self, item):
254        return self.start <= item <= self.end