# Copyright 2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at:
#
#    http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS
# OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the
# License.
from enum import IntEnum
from typing import NamedTuple, Optional, Any, Union, Sequence, Coroutine

from amazon.ion.symbols import SymbolToken

# in Python 3.10, abstract collections have moved into their own module
# for compatibility with 3.10+, first try imports from the new location
# if that fails, try from the pre-3.10 location
try:
    from collections.abc import MutableMapping, MutableSequence
    from collections import OrderedDict
except:
    from collections import MutableMapping, MutableSequence, OrderedDict

from datetime import datetime, timedelta, tzinfo
from decimal import Decimal, Context, Inexact
from math import isnan


class IonType(IntEnum):
    """Enumeration of the Ion data types."""
    NULL = 0
    BOOL = 1
    INT = 2
    FLOAT = 3
    DECIMAL = 4
    TIMESTAMP = 5
    SYMBOL = 6
    STRING = 7
    CLOB = 8
    BLOB = 9
    LIST = 10
    SEXP = 11
    STRUCT = 12

    @property
    def is_numeric(self):
        return IonType.INT <= self <= IonType.TIMESTAMP

    @property
    def is_text(self):
        """Returns whether the type is a Unicode textual type."""
        return self is IonType.SYMBOL or self is IonType.STRING

    @property
    def is_lob(self):
        """Returns whether the type is a LOB."""
        return self is IonType.CLOB or self is IonType.BLOB

    @property
    def is_container(self):
        """Returns whether the type is a container."""
        return self >= IonType.LIST


# TODO At some point we can add SCALAR_START/SCALAR_END for streaming large values.
class IonEventType(IntEnum):
    """Enumeration of Ion parser or serializer events.

    These types do not correspond directly to the Ion type system, but they are related.
    In particular, ``null.*`` will surface as a ``SCALAR`` even though they are containers.

    Attributes:
        INCOMPLETE: Indicates that parsing cannot be completed due to lack of input.
        STREAM_END: Indicates that the logical stream has terminated.
        VERSION_MARKER: Indicates that the **Ion Version Marker** has been encountered.
        SCALAR: Indicates an *atomic* value has been encountered.
        CONTAINER_START: Indicates that the start of a container has been reached.
        CONTAINER_END: Indicates that the end of a container has been reached.
    """
    INCOMPLETE = -2
    STREAM_END = -1
    VERSION_MARKER = 0
    SCALAR = 1
    CONTAINER_START = 2
    CONTAINER_END = 3

    @property
    def begins_value(self):
        """Indicates if the event type is a start of a value."""
        return self is IonEventType.SCALAR or self is IonEventType.CONTAINER_START

    @property
    def ends_container(self):
        """Indicates if the event type terminates a container or stream."""
        return self is IonEventType.STREAM_END or self is IonEventType.CONTAINER_END

    @property
    def is_stream_signal(self):
        """Indicates that the event type corresponds to a stream signal."""
        return self < 0


SymbolTokenOrStr = Union[SymbolToken, str]


class IonEvent(NamedTuple):
    """A parse or serialization event.

    Args:
        event_type (IonEventType): The type of event.
        ion_type (Optional(amazon.ion.core.IonType)): The Ion data model type
            associated with the event.
        value (Optional[any]): The data value associated with the event.
        field_name (Optional[SymbolTokenOrStr]): The field name
            associated with the event.
        annotations (Sequence[SymbolTokenOrStr]): The annotations
            associated with the event.
        depth (Optional[int]): The tree depth of the event if applicable.
    """
    event_type: IonEventType
    ion_type: Optional[IonType] = None
    value: Optional[Any] = None
    field_name: Optional[SymbolTokenOrStr] = None
    annotations: Sequence[SymbolTokenOrStr] = ()
    depth: Optional[int] = None

    def __eq__(self, other):
        if not isinstance(other, IonEvent):
            return False

        if isinstance(self.value, float):
            if not isinstance(other.value, float):
                return False

            # Need to deal with NaN appropriately.
            if self.value != other.value and not (isnan(self.value) and isnan(other.value)):
                return False
        else:
            if self.value != other.value:
                return False

            # Timestamp precision has additional requirements.
            if isinstance(self.value, Timestamp) or isinstance(other.value, Timestamp):
                # Special case for timestamps to capture equivalence over precision.
                self_precision = getattr(self.value, TIMESTAMP_PRECISION_FIELD, None)
                other_precision = getattr(other.value, TIMESTAMP_PRECISION_FIELD, None)
                if self_precision != other_precision \
                        and not ((self_precision is None and other_precision is TimestampPrecision.SECOND) or
                                 (self_precision is TimestampPrecision.SECOND and other_precision is None)):
                    # The absence of precision indicates a naive datetime, which always has SECOND precision.
                    return False

            if isinstance(self.value, Timestamp) and isinstance(other.value, Timestamp):
                self_fractional_seconds = getattr(self.value, TIMESTAMP_FRACTIONAL_SECONDS_FIELD, None)
                other_fractional_seconds = getattr(other.value, TIMESTAMP_FRACTIONAL_SECONDS_FIELD, None)
                if self_fractional_seconds != other_fractional_seconds:
                    return False

            if isinstance(self.value, datetime):
                if self.value.utcoffset() != other.value.utcoffset():
                    return False

        return (self.event_type == other.event_type
            and self.ion_type == other.ion_type
            and self.field_name == other.field_name
            and self.annotations == other.annotations
            and self.depth == other.depth
        )

    def derive_field_name(self, field_name):
        """Derives a new event from this one setting the ``field_name`` attribute.

        Args:
            field_name (Union[amazon.ion.symbols.SymbolToken, unicode]): The field name to set.
        Returns:
            IonEvent: The newly generated event.
        """
        cls = type(self)
        # We use ordinals to avoid thunk materialization.
        return cls(
            self[0],
            self[1],
            self[2],
            field_name,
            self[4],
            self[5]
        )

    def derive_annotations(self, annotations):
        """Derives a new event from this one setting the ``annotations`` attribute.

        Args:
            annotations: (Sequence[Union[amazon.ion.symbols.SymbolToken, unicode]]):
                The annotations associated with the derived event.

        Returns:
            IonEvent: The newly generated event.
        """
        cls = type(self)
        # We use ordinals to avoid thunk materialization.
        return cls(
            self[0],
            self[1],
            self[2],
            self[3],
            annotations,
            self[5]
        )

    def derive_value(self, value):
        """Derives a new event from this one setting the ``value`` attribute.

        Args:
            value: (any):
                The value associated with the derived event.

        Returns:
            IonEvent: The newly generated non-thunk event.
        """
        return IonEvent(
            self.event_type,
            self.ion_type,
            value,
            self.field_name,
            self.annotations,
            self.depth
        )

    def derive_depth(self, depth):
        """Derives a new event from this one setting the ``depth`` attribute.

        Args:
            depth: (int):
                The annotations associated with the derived event.

        Returns:
            IonEvent: The newly generated event.
        """
        cls = type(self)
        # We use ordinals to avoid thunk materialization.
        return cls(
            self[0],
            self[1],
            self[2],
            self[3],
            self[4],
            depth
        )


class MemoizingThunk(object):
    """A :class:`callable` that invokes a ``delegate`` and caches and returns the result."""
    def __init__(self, delegate):
        self.delegate = delegate

    def __call__(self):
        if hasattr(self, 'value'):
            return self.value
        self.value = self.delegate()
        return self.value

    def __str__(self):
        return str(self())

    def __repr__(self):
        return repr(self())


class IonThunkEvent(IonEvent):
    """
    A lazy `IonEvent` whose ``value`` field is a thunk.

    The `value` will be materialized on first access and cached.

    Accessing the value by its slot: ``event[2]`` will avoid materialization.
    """

    @property
    def value(self):
        if hasattr(self, 'cached_value'):
            return self.cached_value
        self.cached_value = self[2]()
        return self.cached_value

# Singletons for structural events
ION_STREAM_END_EVENT = IonEvent(IonEventType.STREAM_END)
ION_STREAM_INCOMPLETE_EVENT = IonEvent(IonEventType.INCOMPLETE)
ION_VERSION_MARKER_EVENT = IonEvent(
    IonEventType.VERSION_MARKER, ion_type=None, value=(1, 0), depth=0
)


class DataEvent(NamedTuple):
    """Event generated as a result of the writer or as input into the reader.

    Args:
        type (IntEnum): The type of event.
        data (bytes):  The serialized data returned.  If no data is to be serialized,
            this should be the empty byte string.
    """
    type: IntEnum
    data: bytes


class Transition(NamedTuple):
    """A pair of event and co-routine delegate.

    This is generally used as a result of a state-machine.

    Args:
        event (Union[DataEvent, IonEvent, None]): The event associated with the transition.
        delegate (Coroutine): The co-routine delegate which can be the same routine from
            whence this transition came.
    """
    event: Union[DataEvent, IonEvent, None]
    delegate: Coroutine


_MIN_OFFSET = timedelta(hours=-24)
_MAX_OFFSET = timedelta(hours=24)
_ZERO_DELTA = timedelta()


class OffsetTZInfo(tzinfo):
    """A trivial UTC offset :class:`tzinfo`.

    todo: as of Python 3 this can be dropped in favor of datetime.timezone(offset_delta)
    """
    def __init__(self, delta=_ZERO_DELTA):
        if delta <= _MIN_OFFSET or delta >= _MAX_OFFSET:
            raise ValueError('Invalid UTC offset: %s' % delta)
        self.delta = delta

    def dst(self, date_time):
        return timedelta()

    def tzname(self, date_time):
        return None

    def utcoffset(self, date_time):
        return self.delta

    def __repr__(self):
        sign = '+'
        delta = self.delta
        if delta < _ZERO_DELTA:
            sign = '-'
            delta = _ZERO_DELTA - delta
        return 'OffsetTZInfo(%s%s)' % (sign, delta)


class TimestampPrecision(IntEnum):
    """The different levels of precision supported in an Ion timestamp."""
    YEAR = 0
    MONTH = 1
    DAY = 2
    MINUTE = 3
    SECOND = 4

    @property
    def includes_month(self):
        """Precision has at least the ``month`` field."""
        return self >= TimestampPrecision.MONTH

    @property
    def includes_day(self):
        """Precision has at least the ``day`` field."""
        return self >= TimestampPrecision.DAY

    @property
    def includes_minute(self):
        """Precision has at least the ``minute`` field."""
        return self >= TimestampPrecision.MINUTE

    @property
    def includes_second(self):
        """Precision has at least the ``second`` field."""
        return self >= TimestampPrecision.SECOND


TIMESTAMP_PRECISION_FIELD = 'precision'
TIMESTAMP_FRACTION_PRECISION_FIELD = 'fractional_precision'
TIMESTAMP_FRACTIONAL_SECONDS_FIELD = 'fractional_seconds'
TIMESTAMP_MICROSECOND_FIELD = 'microsecond'
MICROSECOND_PRECISION = 6
BASE_TEN_MICROSECOND_PRECISION_EXPONENTIATION = 10 ** MICROSECOND_PRECISION
DECIMAL_ZERO = Decimal(0)
PRECISION_LIMIT_LOOKUP = (
    DECIMAL_ZERO,
    Decimal('0.1'),
    Decimal('0.01'),
    Decimal('0.001'),
    Decimal('0.0001'),
    Decimal('0.00001'),
    Decimal('0.000001')
)
DATETIME_CONSTRUCTOR_MICROSECOND_ARGUMENT_INDEX = 6


class Timestamp(datetime):
    """Subclass of ``datetime`` with extended attributes to model Ion Timestamps.

    Extended Attributes:
        precision (TimestampPrecision):
            The year -> second granularity of the timestamp.
            Note that there is no invariant checking of "zero-ness" past the defined
            precision.
        fractional_precision (int):
            The number of decimal digits of ``microseconds`` that are significant.
            Must be >= 0 and <= 6.
        fractional_seconds (Decimal):
            The arbitrarily precise fractional seconds of the timestamp.
            Must be >= 0 and < 1.

    The ``microseconds``, ``fractional_precision``, and ``fractional_seconds``
        will all be present and normalized in a ``Timestamp`` instance.

    If there are more than 6 digits of precision in the ``fractional_seconds``
        then the ``microseconds`` field holds the most significant 6 digits
        and ``fractional_precision`` is ``6``.

    Consider some examples:
        * `2019-10-01T12:45:01Z` would have the following fields set:
          * ``microseconds == 0``, ``fractional_precision == 0``, ``fractional_seconds == Decimal('0')``
        * `2019-10-01T12:45:01.100Z` would have the following fields set:
          * ``microseconds == 100000``, ``fractional_precision == 3``, ``fractional_seconds == Decimal('0.100')``
        * `2019-10-01T12:45:01.123456789Z` would have the following fields set:
          * ``microseconds == 123456``, ``fractional_precision == 6``, ``fractional_seconds ==
          Decimal('0.123456789')``
        """
    __slots__ = [TIMESTAMP_PRECISION_FIELD, TIMESTAMP_FRACTION_PRECISION_FIELD, TIMESTAMP_FRACTIONAL_SECONDS_FIELD]

    def __new__(cls, year, month=1, day=1, hour=0, minute=0, second=0, microsecond=None, tzinfo=None,
                precision=None, fractional_precision=None, fractional_seconds=None):
        """__new__ does the minimum amount of stuff that must be done before the immutable datetime
           super instance is created. This gives internal code, which we assume to know and preserve
           the invariants, fast object construction.

           If you're wondering if you "could" or "should" call this, the answer is no: construct
           Timestamps with the normal constructor.

           If you're working on ion-python itself, and you know what the invariants are and how
           to preserve them, and have good tests, then have at it!
        """
        # We _could_ also do the normalization from microseconds and (optional) fractional_precision
        # here as well. That would be easier to reason about, but that is a more costly operation
        # and is not strictly necessary.
        if microsecond is None:
            if fractional_seconds:
                microsecond = int(fractional_seconds * BASE_TEN_MICROSECOND_PRECISION_EXPONENTIATION)
            else:
                microsecond = 0

        instance = super(Timestamp, cls).__new__(cls, year, month, day, hour, minute, second, microsecond, tzinfo)
        setattr(instance, TIMESTAMP_PRECISION_FIELD, precision)
        setattr(instance, TIMESTAMP_FRACTION_PRECISION_FIELD, fractional_precision)
        setattr(instance, TIMESTAMP_FRACTIONAL_SECONDS_FIELD, fractional_seconds)

        return instance

    def __init__(self, year, month=1, day=1, hour=0, minute=0, second=0, microsecond=None, tzinfo=None,
                precision=None, fractional_precision=None, fractional_seconds=None):
        """Construct a Timestamp with optional attributes defined in the Ion datamodel but not datetimes.

        If None, ``fractional_seconds`` will be set from ``microsecond`` and an optional ``fractional_precision``.
        If None, ``microsecond`` and ``fractional_precision`` will be set from ``fractional_seconds``.
            (though as an implementation detail, this occurs in __new__)
        If both ``fractional_seconds`` and ``microsecond`` are None, then both are set to 0 along with
            ``fractional_precision``.

        The invariant that ``microsecond`` == first ``fractional_precision`` digits of ``fractional_seconds``
        is checked.
        """
        if fractional_precision is not None:
            if not (0 <= fractional_precision <= MICROSECOND_PRECISION):
                raise ValueError('Cannot construct a Timestamp with fractional precision of %d digits, '
                                 'which is out of the supported range of [0, %d].'
                                 % (fractional_precision, MICROSECOND_PRECISION))
            if microsecond is None:
                raise ValueError('Cannot construct a Timestamp with fractional precision and no microseconds')
        if fractional_seconds is not None:
            if not (0 <= fractional_seconds < 1):
                raise ValueError('Cannot construct a Timestamp with fractional seconds of %s, '
                                 'which is out of the supported range of [0, 1).'
                                 % str(fractional_seconds))

            fractional_seconds_exponent = fractional_seconds.as_tuple().exponent
            if fractional_seconds == DECIMAL_ZERO and fractional_seconds_exponent > 0:
                # Zero with a positive exponent is just zero. Set the exponent to zero so fractional_precision is
                # calculated correctly.
                fractional_seconds_exponent = 0
                fractional_seconds = DECIMAL_ZERO
            exp_frac_precision = min(-fractional_seconds_exponent, MICROSECOND_PRECISION)
            if fractional_precision is None:
                fractional_precision = exp_frac_precision
            elif fractional_precision != exp_frac_precision:
                raise ValueError('Cannot construct a Timestamp with fractional precision of %d digits, '
                                 'which is different than the precision of %s.'
                                 % (fractional_precision, str(fractional_seconds)))
            if self.microsecond != int(fractional_seconds * BASE_TEN_MICROSECOND_PRECISION_EXPONENTIATION):
                raise ValueError('Cannot construct a Timestamp with fractional seconds of %s, '
                                 'which is different than the microseconds of %d.'
                                 % (str(fractional_seconds), self.microsecond))
        elif microsecond is not None:
            if fractional_precision is None:
                fractional_precision = MICROSECOND_PRECISION
            if fractional_precision == 0:
                if microsecond > 0:
                    raise ValueError('Cannot construct a Timestamp with fractional precision of 0 digits, '
                                     'but microseconds of %d.' % microsecond)
                fractional_seconds = DECIMAL_ZERO
            else:
                try:
                    fractional_seconds = Decimal(microsecond).scaleb(-MICROSECOND_PRECISION)\
                        .quantize(PRECISION_LIMIT_LOOKUP[fractional_precision], context=Context(traps=[Inexact]))
                except Inexact:
                    raise ValueError('microsecond value %d cannot be expressed exactly in %d digits.'
                                     % (microsecond, fractional_precision))
        else:
            fractional_precision = 0
            fractional_seconds = DECIMAL_ZERO

        setattr(self, TIMESTAMP_PRECISION_FIELD, precision)
        setattr(self, TIMESTAMP_FRACTION_PRECISION_FIELD, fractional_precision)
        setattr(self, TIMESTAMP_FRACTIONAL_SECONDS_FIELD, fractional_seconds)

    def __repr__(self):
        return 'Timestamp(%04d-%02d-%02dT%02d:%02d:%02d.%06d, %r, %r, %s=%s, %s=%s)' % \
               (self.year, self.month, self.day,
                self.hour, self.minute, self.second, self.microsecond,
                self.tzinfo, self.precision,
                TIMESTAMP_FRACTION_PRECISION_FIELD, self.fractional_precision,
                TIMESTAMP_FRACTIONAL_SECONDS_FIELD, self.fractional_seconds)

    @staticmethod
    def adjust_from_utc_fields(*args, **kwargs):
        """Constructs a timestamp from UTC fields adjusted to the local offset if given."""
        raw_ts = Timestamp(*args, **kwargs)
        offset = raw_ts.utcoffset()
        if offset is None or offset == timedelta():
            return raw_ts

        # XXX This returns a datetime, not a Timestamp (which has our precision if defined)
        adjusted = raw_ts + offset
        if raw_ts.precision is None:
            # No precision means we can just return a regular datetime
            return adjusted

        return Timestamp(
            adjusted.year,
            adjusted.month,
            adjusted.day,
            adjusted.hour,
            adjusted.minute,
            adjusted.second,
            None,
            raw_ts.tzinfo,
            precision=raw_ts.precision,
            fractional_precision=None,
            fractional_seconds=raw_ts.fractional_seconds
        )


def timestamp(year, month=1, day=1,
              hour=0, minute=0, second=0, microsecond=None,
              off_hours=None, off_minutes=None,
              precision=None, fractional_precision=None, fractional_seconds=None):
    """Shorthand for the :class:`Timestamp` constructor.

    Specifically, converts ``off_hours`` and ``off_minutes`` parameters to a suitable
    :class:`OffsetTZInfo` instance.
    """
    delta = None
    if off_hours is not None:
        if off_hours < -23 or off_hours > 23:
            raise ValueError('Hour offset %d is out of required range -23..23.' % (off_hours,))
        delta = timedelta(hours=off_hours)
    if off_minutes is not None:
        if off_minutes < -59 or off_minutes > 59:
            raise ValueError('Minute offset %d is out of required range -59..59.' % (off_minutes,))
        minutes_delta = timedelta(minutes=off_minutes)
        if delta is None:
            delta = minutes_delta
        else:
            delta += minutes_delta

    tz = None
    if delta is not None:
        tz = OffsetTZInfo(delta)

    return Timestamp(
        year, month, day,
        hour, minute, second, microsecond,
        tz, precision=precision, fractional_precision=fractional_precision, fractional_seconds=fractional_seconds
    )


class Multimap(MutableMapping):
    """
    Dictionary that can hold multiple values for the same key

    In order not to break existing customers, getting and inserting elements with ``[]`` keeps the same behaviour
    as the built-in dict. If multiple elements are already mapped to the key, ``[]`  will return
    the newest one.

    To map multiple elements to a key, use the ``add_item`` operation.
    To retrieve all the values map to a key, use ``get_all_values``.
    """

    def __init__(self, *args, **kwargs):
        super(Multimap, self).__init__()
        self.__store = OrderedDict()
        if args is not None and len(args) > 0:
            for key, value in iter(args[0].items()):
                self.__store[key] = MultimapValue(value)

    def __getitem__(self, key):
        return self.__store[key][len(self.__store[key]) - 1]  # Return only one in order not to break clients

    def __delitem__(self, key):
        del self.__store[key]

    def __setitem__(self, key, value):
        self.__store[key] = MultimapValue(value)

    def __len__(self):
        return sum([len(values) for values in iter(self.__store.values())])

    def __iter__(self):
        for key in iter(self.__store.keys()):
            yield key

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return '{%s}' % ', '.join(['%r: %r' % (k, v) for k, v in self.items()])

    def add_item(self, key, value):
        if key in self.__store:
            self.__store[key].append(value)
        else:
            self.__setitem__(key, value)

    def get_all_values(self, key):
        return self.__store[key]

    def iteritems(self):
        for key in self.__store:
            for value in self.__store[key]:
                yield (key, value)

    def items(self):
        output = []
        for k, v in self.iteritems():
            output.append((k, v))
        return output


class MultimapValue(MutableSequence):

    def __init__(self, *args):
        if args is not None:
            self.__store = [x for x in args]
        else:
            self.__store = []

    def insert(self, index, value):
        self.__setitem__(index, value)

    def __len__(self):
        return len(self.__store)

    def __getitem__(self, index):
        return self.__store[index]

    def __setitem__(self, index, value):
        self.__store.insert(index, value)

    def __delitem__(self, index):
        del self.__store[index]

    def __iter__(self):
        for x in self.__store:
            yield x
