"""
Python 3.8 / 3.9+ 时区兼容层。

Python 3.9+ 内置 zoneinfo；Python 3.8 回退到 pytz。
统一暴露 ZoneInfo 和 make_aware() 两个接口，
其余模块只从此处导入，不直接写 from zoneinfo import ZoneInfo。
"""

from __future__ import annotations

from datetime import datetime

try:
    from zoneinfo import ZoneInfo                               # Python 3.9+

    def make_aware(dt: datetime, key: str) -> datetime:
        """将朴素 datetime 本地化为指定时区的 aware datetime。"""
        if dt.tzinfo is not None:
            return dt.astimezone(ZoneInfo(key))
        return dt.replace(tzinfo=ZoneInfo(key))

except ImportError:
    import pytz                                                 # Python 3.8 fallback
    from datetime import tzinfo as _tzinfo

    class ZoneInfo(_tzinfo):                                    # type: ignore[no-redef]
        """pytz 驱动的 ZoneInfo 最小兼容实现，继承 tzinfo 供 datetime.now(tz) 使用。"""

        _cache: dict = {}

        def __new__(cls, key: str):
            if key not in cls._cache:
                obj = _tzinfo.__new__(cls)
                obj._tz  = pytz.timezone(key)
                obj._key = key
                cls._cache[key] = obj
            return cls._cache[key]

        def __init__(self, key: str) -> None:
            pass

        @property
        def key(self) -> str:
            return self._key

        def utcoffset(self, dt):
            naive = dt.replace(tzinfo=None) if dt is not None else datetime.utcnow()
            return self._tz.utcoffset(naive)

        def tzname(self, dt):
            return self._tz.zone

        def dst(self, dt):
            naive = dt.replace(tzinfo=None) if dt is not None else datetime.utcnow()
            return self._tz.dst(naive)

        def fromutc(self, dt: datetime) -> datetime:
            # pytz.fromutc 要求 dt.tzinfo 是其自身，先替换回 pytz tz
            return self._tz.fromutc(dt.replace(tzinfo=self._tz))

    def make_aware(dt: datetime, key: str) -> datetime:         # type: ignore[misc]
        """使用 pytz.localize() 正确处理 DST，避免朴素 datetime 赋值偏移问题。"""
        tz = pytz.timezone(key)
        if dt.tzinfo is not None:
            return dt.astimezone(tz)
        return tz.localize(dt)
