"""Use zoneinfo timezones"""
from __future__ import annotations
try:
    import zoneinfo
except:
    from backports import zoneinfo
from icalendar import prop
from dateutil.rrule import rrule, rruleset
from datetime import datetime, tzinfo
from typing import Optional
from .provider import TZProvider
from .. import cal
from io import StringIO
from dateutil.tz import tzical
from dateutil.tz.tz import _tzicalvtz
import copyreg
import functools
import copy


class ZONEINFO(TZProvider):
    """Provide icalendar with timezones from zoneinfo."""

    name = "zoneinfo"
    utc = zoneinfo.ZoneInfo("UTC")
    _available_timezones = zoneinfo.available_timezones()

    def localize(self, dt: datetime, tz: zoneinfo.ZoneInfo) -> datetime:
        """Localize a datetime to a timezone."""
        return dt.replace(tzinfo=tz)

    def localize_utc(self, dt: datetime) -> datetime:
        """Return the datetime in UTC."""
        if getattr(dt, 'tzinfo', False) and dt.tzinfo is not None:
            return dt.astimezone(self.utc)
        return self.localize(dt, self.utc)

    def timezone(self, name: str) -> Optional[tzinfo]:
        """Return a timezone with a name or None if we cannot find it."""
        try:
            return zoneinfo.ZoneInfo(name)
        except zoneinfo.ZoneInfoNotFoundError:
            pass
        except ValueError:
            # ValueError: ZoneInfo keys may not be absolute paths, got: /Europe/CUSTOM
            pass

    def knows_timezone_id(self, id: str) -> bool:
        """Whether the timezone is already cached by the implementation."""
        return id in self._available_timezones

    def fix_rrule_until(self, rrule:rrule, ical_rrule:prop.vRecur) -> None:
        """Make sure the until value works for the rrule generated from the ical_rrule."""
        if not {'UNTIL', 'COUNT'}.intersection(ical_rrule.keys()):
            # zoninfo does not know any transition dates after 2038
            rrule._until = datetime(2038, 12, 31, tzinfo=self.utc)

    def create_timezone(self, tz: cal.Timezone) -> tzinfo:
        """Create a timezone from the given information."""
        try:
            return self._create_timezone(tz)
        except ValueError:
            # We might have a custom component in there.
            # see https://github.com/python/cpython/issues/120217
            tz = copy.deepcopy(tz)
            for sub in tz.walk():
                for attr in list(sub.keys()):
                    if attr.lower().startswith("x-"):
                        sub.pop(attr)
            return self._create_timezone(tz)

    def _create_timezone(self, tz: cal.Timezone) -> tzinfo:
        """Create a timezone and maybe fail"""
        file = StringIO(tz.to_ical().decode("UTF-8", "replace"))
        return tzical(file).get()

    def uses_pytz(self) -> bool:
        """Whether we use pytz."""
        return False

    def uses_zoneinfo(self) -> bool:
        """Whether we use zoneinfo."""
        return True


def pickle_tzicalvtz(tzicalvtz:tz._tzicalvtz):
    """Because we use dateutil.tzical, we need to make it pickle-able."""
    return _tzicalvtz, (tzicalvtz._tzid, tzicalvtz._comps)

copyreg.pickle(_tzicalvtz, pickle_tzicalvtz)


def pickle_rrule_with_cache(self: rrule):
    """Make sure we can also pickle rrules that cache.

    This is mainly copied from rrule.replace.
    """
    new_kwargs = {"interval": self._interval,
                  "count": self._count,
                  "dtstart": self._dtstart,
                  "freq": self._freq,
                  "until": self._until,
                  "wkst": self._wkst,
                  "cache": False if self._cache is None else True }
    new_kwargs.update(self._original_rule)
    # from https://stackoverflow.com/a/64915638/1320237
    return functools.partial(rrule, new_kwargs.pop("freq"), **new_kwargs), ()

copyreg.pickle(rrule, pickle_rrule_with_cache)

def pickle_rruleset_with_cache(rs: rruleset):
    """Pickle an rruleset."""
    # self._rrule = []
    # self._rdate = []
    # self._exrule = []
    # self._exdate = []
    return unpickle_rruleset_with_cache, (
        rs._rrule, rs._rdate, rs._exrule,
        rs._exdate, False if rs._cache is None else True
    )

def unpickle_rruleset_with_cache(rrule, rdate, exrule, exdate, cache):
    """unpickling the rruleset."""
    rs = rruleset(cache)
    for o in rrule: rs.rrule(o)
    for o in rdate: rs.rdate(o)
    for o in exrule: rs.exrule(o)
    for o in exdate: rs.exdate(o)
    return rs

copyreg.pickle(rruleset, pickle_rruleset_with_cache)

__all__ = ["ZONEINFO"]
