thc-webhook/thctime.py

from datetime import datetime
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

import pytz


TZDB_CACHE: dict | None = None


def get_tzdb_cache() -> dict | None:
    return TZDB_CACHE


def init_tzdb_cache() -> dict:
    """Initialize a cached lookup structure for tzdb data."""
    global TZDB_CACHE
    if TZDB_CACHE is not None:
        return TZDB_CACHE

    timezones = load_timezones()
    countries = load_countries()

    tz_to_country_code: dict[str, str] = {}
    tz_meta: dict[str, dict] = {}
    for tz in timezones:
        zone_name = tz.get("zone_name")
        country_code = tz.get("country_code")
        if not isinstance(zone_name, str) or not zone_name:
            continue
        if not isinstance(country_code, str) or not country_code:
            continue

        tz_to_country_code[zone_name] = country_code
        region, city = split_tz_name(zone_name)
        tz_meta[zone_name] = {
            "zone_name": zone_name,
            "country_code": country_code,
            "region": region,
            "city": city,
        }

    country_code_to_name: dict[str, str] = {}
    for c in countries:
        code = c.get("country_code")
        name = c.get("country_name")
        if code and name:
            country_code_to_name[code] = str(name).strip().strip('"')

    for zone_name, meta in tz_meta.items():
        code = meta.get("country_code")
        if isinstance(code, str):
            meta["country_name"] = country_code_to_name.get(code)

    tz_names: list[str] = []
    for zone_name in tz_to_country_code.keys():
        try:
            ZoneInfo(zone_name)
        except ZoneInfoNotFoundError:
            continue
        tz_names.append(zone_name)

    TZDB_CACHE = {
        "tz_to_country_code": tz_to_country_code,
        "country_code_to_name": country_code_to_name,
        "tz_names": tz_names,
        "tz_meta": tz_meta,
    }
    return TZDB_CACHE


def split_tz_name(zone_name: str) -> tuple[str, str]:
    """Split an IANA timezone name into (region, city)."""
    if "/" not in zone_name:
        return zone_name, ""
    region, rest = zone_name.split("/", 1)
    return region, rest


def load_timezones() -> list[dict]:
    """Load timezones from csv file."""
    with open("tzdb/TimeZoneDB.csv/time_zone.csv", "r", encoding="utf-8") as f:
        lines = f.readlines()

    timezones = []
    for line in lines:
        fields = line.strip().split(",")
        if len(fields) >= 5:
            timezones.append({
                "zone_name": fields[0],
                "country_code": fields[1],
                "abbreviation": fields[2],
                "time_start": fields[3],
                "gmt_offset": int(fields[4]),
                "dst": fields[5] == "1",
            })
    return timezones


def load_countries() -> list[dict]:
    """Load countries from csv file."""
    with open("tzdb/TimeZoneDB.csv/country.csv", "r", encoding="utf-8") as f:
        lines = f.readlines()

    countries = []
    for line in lines:
        fields = line.strip().split(",")
        if len(fields) >= 2:
            countries.append({
                "country_code": fields[0],
                "country_name": fields[1],
            })
    return countries


def get_tz_info(tz_name: str, timezones: list[dict]) -> dict | None:
    """Get timezone info by name."""
    return next((tz for tz in timezones if tz["zone_name"] == tz_name), None)


def get_country_info(country_code: str, countries: list[dict]) -> dict | None:
    """Get country info by country code."""
    return next((c for c in countries if c["country_code"] == country_code), None)


def where_is_it_420(
    timezones: list[dict],
    countries: list[dict],
    tz_names: list[str] | None = None,
    tz_to_country_code: dict[str, str] | None = None,
    country_code_to_name: dict[str, str] | None = None,
) -> list[str]:
    """Get timezones where the current hour is 4 or 16."""
    if tz_to_country_code is None:
        tz_to_country_code = {}
        for tz in timezones:
            zone_name = tz.get("zone_name")
            country_code = tz.get("country_code")
            if isinstance(zone_name, str) and isinstance(country_code, str):
                tz_to_country_code[zone_name] = country_code

    if country_code_to_name is None:
        country_code_to_name = {}
        for c in countries:
            code = c.get("country_code")
            name = c.get("country_name")
            if isinstance(code, str) and name is not None:
                country_code_to_name[code] = str(name).strip().strip('"')

    names_to_check = tz_names if tz_names is not None else pytz.all_timezones
    results: list[str] = []
    seen: set[str] = set()

    for tz_name in names_to_check:
        try:
            tz_obj = pytz.timezone(tz_name)
        except Exception:
            continue

        now = datetime.now(tz_obj)
        if now.hour != 4 and now.hour != 16:
            continue

        country_code = tz_to_country_code.get(tz_name)
        if not country_code:
            continue
        country_name = country_code_to_name.get(country_code)
        if not country_name:
            continue
        if country_name in seen:
            continue

        seen.add(country_name)
        results.append(country_name)

    return results