1"""Provides consistent datetime handling across Airbyte with ISO8601/RFC3339 compliance.
  2
  3Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  4
  5This module provides a custom datetime class (AirbyteDateTime) and helper functions that ensure
  6consistent datetime handling across Airbyte. All datetime strings are formatted according to
  7ISO8601/RFC3339 standards with 'T' delimiter and '+00:00' for UTC timezone.
  8
  9Key Features:
 10    - Timezone-aware datetime objects (defaults to UTC)
 11    - ISO8601/RFC3339 compliant string formatting
 12    - Consistent parsing of various datetime formats
 13    - Support for Unix timestamps and milliseconds
 14    - Type-safe datetime arithmetic with timedelta
 15
 16## Basic Usage
 17
 18```python
 19from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now, ab_datetime_parse
 20from datetime import timedelta, timezone
 21
 22## Current time in UTC
 23now = ab_datetime_now()
 24print(now)  # 2023-03-14T15:09:26.535897Z
 25
 26# Parse various datetime formats
 27dt = ab_datetime_parse("2023-03-14T15:09:26Z")  # ISO8601/RFC3339
 28dt = ab_datetime_parse("2023-03-14")  # Date only (assumes midnight UTC)
 29dt = ab_datetime_parse(1678806566)  # Unix timestamp
 30
 31## Create with explicit timezone
 32dt = AirbyteDateTime(2023, 3, 14, 15, 9, 26, tzinfo=timezone.utc)
 33print(dt)  # 2023-03-14T15:09:26+00:00
 34
 35# Datetime arithmetic with timedelta
 36tomorrow = dt + timedelta(days=1)
 37yesterday = dt - timedelta(days=1)
 38time_diff = tomorrow - yesterday  # timedelta object
 39```
 40
 41## Millisecond Timestamp Handling
 42
 43```python
 44# Convert to millisecond timestamp
 45dt = ab_datetime_parse("2023-03-14T15:09:26Z")
 46ms = dt.to_epoch_millis()  # 1678806566000
 47
 48# Create from millisecond timestamp
 49dt = AirbyteDateTime.from_epoch_millis(1678806566000)
 50print(dt)  # 2023-03-14T15:09:26Z
 51```
 52
 53## Timezone Handling
 54
 55```python
 56# Create with non-UTC timezone
 57tz = timezone(timedelta(hours=-4))  # EDT
 58dt = AirbyteDateTime(2023, 3, 14, 15, 9, 26, tzinfo=tz)
 59print(dt)  # 2023-03-14T15:09:26-04:00
 60
 61## Parse with timezone
 62dt = ab_datetime_parse("2023-03-14T15:09:26-04:00")
 63print(dt)  # 2023-03-14T15:09:26-04:00
 64
 65## Naive datetimes are automatically converted to UTC
 66dt = ab_datetime_parse("2023-03-14T15:09:26")
 67print(dt)  # 2023-03-14T15:09:26Z
 68```
 69
 70# Format Validation
 71
 72```python
 73from airbyte_cdk.utils.datetime_helpers import ab_datetime_try_parse
 74
 75# Validate ISO8601/RFC3339 format
 76assert ab_datetime_try_parse("2023-03-14T15:09:26Z")       # Basic UTC format
 77assert ab_datetime_try_parse("2023-03-14T15:09:26-04:00")  # With timezone offset
 78assert ab_datetime_try_parse("2023-03-14T15:09:26+00:00")  # With explicit UTC offset
 79assert ab_datetime_try_parse("2023-03-14 15:09:26Z")       # Missing T delimiter but still parsable
 80assert not ab_datetime_try_parse("foo")                    # Invalid: not parsable, returns `None`
 81```
 82"""
 83
 84from datetime import datetime, timedelta, timezone
 85from typing import Any, Union, overload
 86
 87from dateutil import parser
 88from whenever import Instant
 89
 90
 91class AirbyteDateTime(datetime):
 92    """A timezone-aware datetime class with ISO8601/RFC3339 string representation and operator overloading.
 93
 94    This class extends the standard datetime class to provide consistent timezone handling
 95    (defaulting to UTC) and ISO8601/RFC3339 compliant string formatting. It also supports
 96    operator overloading for datetime arithmetic with timedelta objects.
 97
 98    Example:
 99        >>> dt = AirbyteDateTime(2023, 3, 14, 15, 9, 26, tzinfo=timezone.utc)
100        >>> str(dt)
101        '2023-03-14T15:09:26+00:00'
102        >>> dt + timedelta(hours=1)
103        '2023-03-14T16:09:26+00:00'
104    """
105
106    def __new__(cls, *args: Any, **kwargs: Any) -> "AirbyteDateTime":
107        """Creates a new timezone-aware AirbyteDateTime instance.
108
109        Ensures all instances are timezone-aware by defaulting to UTC if no timezone is provided.
110
111        Returns:
112            AirbyteDateTime: A new timezone-aware datetime instance.
113        """
114        self = super().__new__(cls, *args, **kwargs)
115        if self.tzinfo is None:
116            return self.replace(tzinfo=timezone.utc)
117        return self
118
119    @classmethod
120    def from_datetime(cls, dt: datetime) -> "AirbyteDateTime":
121        """Converts a standard datetime to AirbyteDateTime.
122
123        Args:
124            dt: A standard datetime object to convert.
125
126        Returns:
127            AirbyteDateTime: A new timezone-aware AirbyteDateTime instance.
128        """
129        return cls(
130            dt.year,
131            dt.month,
132            dt.day,
133            dt.hour,
134            dt.minute,
135            dt.second,
136            dt.microsecond,
137            dt.tzinfo or timezone.utc,
138        )
139
140    def to_datetime(self) -> datetime:
141        """Converts this AirbyteDateTime to a standard datetime object.
142
143        Today, this just returns `self` because AirbyteDateTime is a subclass of `datetime`.
144        In the future, we may modify our internal representation to use a different base class.
145        """
146        return self
147
148    def __str__(self) -> str:
149        """Returns the datetime in ISO8601/RFC3339 format with 'T' delimiter.
150
151        Ensures consistent string representation with timezone, using '+00:00' for UTC.
152        Preserves full microsecond precision when present, omits when zero.
153
154        Returns:
155            str: ISO8601/RFC3339 formatted string.
156        """
157        aware_self = self if self.tzinfo else self.replace(tzinfo=timezone.utc)
158        return aware_self.isoformat(sep="T", timespec="auto")
159
160    def __repr__(self) -> str:
161        """Returns the same string representation as __str__ for consistency.
162
163        Returns:
164            str: ISO8601/RFC3339 formatted string.
165        """
166        return self.__str__()
167
168    def add(self, delta: timedelta) -> "AirbyteDateTime":
169        """Add a timedelta interval to this datetime.
170
171        This method provides a more explicit alternative to the + operator
172        for adding time intervals to datetimes.
173
174        Args:
175            delta: The timedelta interval to add.
176
177        Returns:
178            AirbyteDateTime: A new datetime with the interval added.
179
180        Example:
181            >>> dt = AirbyteDateTime(2023, 3, 14, tzinfo=timezone.utc)
182            >>> dt.add(timedelta(hours=1))
183            '2023-03-14T01:00:00Z'
184        """
185        return self + delta
186
187    def subtract(self, delta: timedelta) -> "AirbyteDateTime":
188        """Subtract a timedelta interval from this datetime.
189
190        This method provides a more explicit alternative to the - operator
191        for subtracting time intervals from datetimes.
192
193        Args:
194            delta: The timedelta interval to subtract.
195
196        Returns:
197            AirbyteDateTime: A new datetime with the interval subtracted.
198
199        Example:
200            >>> dt = AirbyteDateTime(2023, 3, 14, tzinfo=timezone.utc)
201            >>> dt.subtract(timedelta(hours=1))
202            '2023-03-13T23:00:00Z'
203        """
204        result = super().__sub__(delta)
205        if isinstance(result, datetime):
206            return AirbyteDateTime.from_datetime(result)
207        raise TypeError("Invalid operation")
208
209    def __add__(self, other: timedelta) -> "AirbyteDateTime":
210        """Adds a timedelta to this datetime.
211
212        Args:
213            other: A timedelta object to add.
214
215        Returns:
216            AirbyteDateTime: A new datetime with the timedelta added.
217
218        Raises:
219            TypeError: If other is not a timedelta.
220        """
221        result = super().__add__(other)
222        if isinstance(result, datetime):
223            return AirbyteDateTime.from_datetime(result)
224        raise TypeError("Invalid operation")
225
226    def __radd__(self, other: timedelta) -> "AirbyteDateTime":
227        """Supports timedelta + AirbyteDateTime operation.
228
229        Args:
230            other: A timedelta object to add.
231
232        Returns:
233            AirbyteDateTime: A new datetime with the timedelta added.
234
235        Raises:
236            TypeError: If other is not a timedelta.
237        """
238        return self.__add__(other)
239
240    @overload  # type: ignore[override]
241    def __sub__(self, other: timedelta) -> "AirbyteDateTime": ...
242
243    @overload  # type: ignore[override]
244    def __sub__(self, other: Union[datetime, "AirbyteDateTime"]) -> timedelta: ...
245
246    def __sub__(
247        self, other: Union[datetime, "AirbyteDateTime", timedelta]
248    ) -> Union[timedelta, "AirbyteDateTime"]:  # type: ignore[override]
249        """Subtracts a datetime, AirbyteDateTime, or timedelta from this datetime.
250
251        Args:
252            other: A datetime, AirbyteDateTime, or timedelta object to subtract.
253
254        Returns:
255            Union[timedelta, AirbyteDateTime]: A timedelta if subtracting datetime/AirbyteDateTime,
256                or a new datetime if subtracting timedelta.
257
258        Raises:
259            TypeError: If other is not a datetime, AirbyteDateTime, or timedelta.
260        """
261        if isinstance(other, timedelta):
262            result = super().__sub__(other)  # type: ignore[call-overload]
263            if isinstance(result, datetime):
264                return AirbyteDateTime.from_datetime(result)
265        elif isinstance(other, (datetime, AirbyteDateTime)):
266            result = super().__sub__(other)  # type: ignore[call-overload]
267            if isinstance(result, timedelta):
268                return result
269        raise TypeError(
270            f"unsupported operand type(s) for -: '{type(self).__name__}' and '{type(other).__name__}'"
271        )
272
273    def __rsub__(self, other: datetime) -> timedelta:
274        """Supports datetime - AirbyteDateTime operation.
275
276        Args:
277            other: A datetime object.
278
279        Returns:
280            timedelta: The time difference between the datetimes.
281
282        Raises:
283            TypeError: If other is not a datetime.
284        """
285        if not isinstance(other, datetime):
286            return NotImplemented
287        result = other - datetime(
288            self.year,
289            self.month,
290            self.day,
291            self.hour,
292            self.minute,
293            self.second,
294            self.microsecond,
295            self.tzinfo,
296        )
297        if isinstance(result, timedelta):
298            return result
299        raise TypeError("Invalid operation")
300
301    def to_epoch_millis(self) -> int:
302        """Return the Unix timestamp in milliseconds for this datetime.
303
304        Returns:
305            int: Number of milliseconds since Unix epoch (January 1, 1970).
306
307        Example:
308            >>> dt = AirbyteDateTime(2023, 3, 14, 15, 9, 26, tzinfo=timezone.utc)
309            >>> dt.to_epoch_millis()
310            1678806566000
311        """
312        return int(self.timestamp() * 1000)
313
314    @classmethod
315    def from_epoch_millis(cls, milliseconds: int) -> "AirbyteDateTime":
316        """Create an AirbyteDateTime from Unix timestamp in milliseconds.
317
318        Args:
319            milliseconds: Number of milliseconds since Unix epoch (January 1, 1970).
320
321        Returns:
322            AirbyteDateTime: A new timezone-aware datetime instance (UTC).
323
324        Example:
325            >>> dt = AirbyteDateTime.from_epoch_millis(1678806566000)
326            >>> str(dt)
327            '2023-03-14T15:09:26+00:00'
328        """
329        return cls.fromtimestamp(milliseconds / 1000.0, timezone.utc)
330
331    @classmethod
332    def from_str(cls, dt_str: str) -> "AirbyteDateTime":
333        """Thin convenience wrapper around `ab_datetime_parse()`.
334
335        This method attempts to create a new `AirbyteDateTime` using all available parsing
336        strategies.
337
338        Raises:
339            ValueError: If the value cannot be parsed into a valid datetime object.
340        """
341        return ab_datetime_parse(dt_str)
342
343
344def ab_datetime_now() -> AirbyteDateTime:
345    """Returns the current time as an AirbyteDateTime in UTC timezone.
346
347    Previously named: now()
348
349    Returns:
350        AirbyteDateTime: Current UTC time.
351
352    Example:
353        >>> dt = ab_datetime_now()
354        >>> str(dt)  # Returns current time in ISO8601/RFC3339
355        '2023-03-14T15:09:26.535897Z'
356    """
357    return AirbyteDateTime.from_datetime(datetime.now(timezone.utc))
358
359
360def ab_datetime_parse(dt_str: str | int) -> AirbyteDateTime:
361    """Parses a datetime string or timestamp into an AirbyteDateTime with timezone awareness.
362
363    This implementation is as flexible as possible to handle various datetime formats.
364    Always returns a timezone-aware datetime (defaults to UTC if no timezone specified).
365
366    Handles:
367        - ISO8601/RFC3339 format strings (with ' ' or 'T' delimiter)
368        - Unix timestamps (as integers or strings)
369        - Date-only strings (YYYY-MM-DD)
370        - Timezone-aware formats (+00:00 for UTC, or ±HH:MM offset)
371        - Anything that can be parsed by `dateutil.parser.parse()`
372
373    Args:
374        dt_str: A datetime string in ISO8601/RFC3339 format, Unix timestamp (int/str),
375            or other recognizable datetime format.
376
377    Returns:
378        AirbyteDateTime: A timezone-aware datetime object.
379
380    Raises:
381        ValueError: If the input cannot be parsed as a valid datetime.
382
383    Example:
384        >>> ab_datetime_parse("2023-03-14T15:09:26+00:00")
385        '2023-03-14T15:09:26+00:00'
386        >>> ab_datetime_parse(1678806000)  # Unix timestamp
387        '2023-03-14T15:00:00+00:00'
388        >>> ab_datetime_parse("2023-03-14")  # Date-only
389        '2023-03-14T00:00:00+00:00'
390    """
391    try:
392        # Handle numeric values as Unix timestamps (UTC)
393        if isinstance(dt_str, int) or (
394            isinstance(dt_str, str)
395            and (dt_str.isdigit() or (dt_str.startswith("-") and dt_str[1:].isdigit()))
396        ):
397            timestamp = int(dt_str)
398            if timestamp < 0:
399                raise ValueError("Timestamp cannot be negative")
400            if len(str(abs(timestamp))) > 10:
401                raise ValueError("Timestamp value too large")
402            instant = Instant.from_timestamp(timestamp)
403            return AirbyteDateTime.from_datetime(instant.py_datetime())
404
405        if not isinstance(dt_str, str):
406            raise ValueError(
407                f"Could not parse datetime string: expected string or integer, got {type(dt_str)}"
408            )
409
410        # Handle date-only format first
411        if ":" not in dt_str and dt_str.count("-") == 2 and "/" not in dt_str:
412            try:
413                year, month, day = map(int, dt_str.split("-"))
414                if not (1 <= month <= 12 and 1 <= day <= 31):
415                    raise ValueError(f"Invalid date format: {dt_str}")
416                instant = Instant.from_utc(year, month, day, 0, 0, 0)
417                return AirbyteDateTime.from_datetime(instant.py_datetime())
418            except (ValueError, TypeError):
419                raise ValueError(f"Invalid date format: {dt_str}")
420
421        # Reject time-only strings without date
422        if ":" in dt_str and dt_str.count("-") < 2 and dt_str.count("/") < 2:
423            raise ValueError(f"Missing date part in datetime string: {dt_str}")
424
425        # Try parsing with dateutil for timezone handling
426        try:
427            parsed = parser.parse(dt_str)
428            if parsed.tzinfo is None:
429                parsed = parsed.replace(tzinfo=timezone.utc)
430
431            return AirbyteDateTime.from_datetime(parsed)
432        except (ValueError, TypeError):
433            raise ValueError(f"Could not parse datetime string: {dt_str}")
434    except ValueError as e:
435        if "Invalid date format:" in str(e):
436            raise
437        if "Timestamp cannot be negative" in str(e):
438            raise
439        if "Timestamp value too large" in str(e):
440            raise
441        raise ValueError(f"Could not parse datetime string: {dt_str}")
442
443
444def ab_datetime_try_parse(dt_str: str) -> AirbyteDateTime | None:
445    """Try to parse the input as a datetime, failing gracefully instead of raising an exception.
446
447    This is a thin wrapper around `ab_datetime_parse()` that catches parsing errors and
448    returns `None` instead of raising an exception.
449    The implementation is as flexible as possible to handle various datetime formats.
450    Always returns a timezone-aware datetime (defaults to `UTC` if no timezone specified).
451
452    Example:
453        >>> ab_datetime_try_parse("2023-03-14T15:09:26Z")  # Returns AirbyteDateTime
454        >>> ab_datetime_try_parse("2023-03-14 15:09:26Z")  # Missing 'T' delimiter still parsable
455        >>> ab_datetime_try_parse("2023-03-14")            # Returns midnight UTC time
456    """
457    try:
458        return ab_datetime_parse(dt_str)
459    except (ValueError, TypeError):
460        return None
461
462
463def ab_datetime_format(
464    dt: Union[datetime, AirbyteDateTime],
465    format: str | None = None,
466) -> str:
467    """Formats a datetime object as an ISO8601/RFC3339 string with 'T' delimiter and timezone.
468
469    Previously named: format()
470
471    Converts any datetime object to a string with 'T' delimiter and proper timezone.
472    If the datetime is naive (no timezone), UTC is assumed.
473    Uses '+00:00' for UTC timezone, otherwise keeps the original timezone offset.
474
475    Args:
476        dt: Any datetime object to format.
477        format: Optional format string. If provided, calls `strftime()` with this format.
478            Otherwise, uses the default ISO8601/RFC3339 format, adapted for available precision.
479
480    Returns:
481        str: ISO8601/RFC3339 formatted datetime string.
482
483    Example:
484        >>> dt = datetime(2023, 3, 14, 15, 9, 26, tzinfo=timezone.utc)
485        >>> ab_datetime_format(dt)
486        '2023-03-14T15:09:26+00:00'
487    """
488    if isinstance(dt, AirbyteDateTime):
489        return str(dt)
490
491    if dt.tzinfo is None:
492        dt = dt.replace(tzinfo=timezone.utc)
493
494    if format:
495        return dt.strftime(format)
496
497    # Format with consistent timezone representation and "T" delimiter
498    return dt.isoformat(sep="T", timespec="auto")