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