from __future__ import annotations

from typing import TYPE_CHECKING
from zoneinfo import ZoneInfo

if TYPE_CHECKING:
    from skyvern.schemas.runs import ProxyLocation


# Mapping derived from the IANA tz database (zone.tab). For multi-timezone
# countries we choose the primary commercial hub; Kosovo falls back to Europe/Belgrade.
RESIDENTIAL_COUNTRY_TIMEZONES: dict[str, str] = {
    "AD": "Europe/Andorra",
    "AE": "Asia/Dubai",
    "AF": "Asia/Kabul",
    "AG": "America/Antigua",
    "AI": "America/Anguilla",
    "AL": "Europe/Tirane",
    "AM": "Asia/Yerevan",
    "AO": "Africa/Luanda",
    "AR": "America/Argentina/Buenos_Aires",
    "AS": "Pacific/Pago_Pago",
    "AT": "Europe/Vienna",
    "AU": "Australia/Sydney",
    "AW": "America/Aruba",
    "AX": "Europe/Mariehamn",
    "AZ": "Asia/Baku",
    "BA": "Europe/Sarajevo",
    "BB": "America/Barbados",
    "BD": "Asia/Dhaka",
    "BE": "Europe/Brussels",
    "BF": "Africa/Ouagadougou",
    "BG": "Europe/Sofia",
    "BH": "Asia/Bahrain",
    "BI": "Africa/Bujumbura",
    "BJ": "Africa/Porto-Novo",
    "BM": "Atlantic/Bermuda",
    "BN": "Asia/Brunei",
    "BO": "America/La_Paz",
    "BQ": "America/Kralendijk",
    "BR": "America/Sao_Paulo",
    "BS": "America/Nassau",
    "BT": "Asia/Thimphu",
    "BW": "Africa/Gaborone",
    "BY": "Europe/Minsk",
    "BZ": "America/Belize",
    "CA": "America/Toronto",
    "CD": "Africa/Kinshasa",
    "CF": "Africa/Bangui",
    "CG": "Africa/Brazzaville",
    "CH": "Europe/Zurich",
    "CI": "Africa/Abidjan",
    "CK": "Pacific/Rarotonga",
    "CL": "America/Santiago",
    "CM": "Africa/Douala",
    "CN": "Asia/Shanghai",
    "CO": "America/Bogota",
    "CR": "America/Costa_Rica",
    "CU": "America/Havana",
    "CV": "Atlantic/Cape_Verde",
    "CW": "America/Curacao",
    "CY": "Asia/Nicosia",
    "CZ": "Europe/Prague",
    "DE": "Europe/Berlin",
    "DJ": "Africa/Djibouti",
    "DK": "Europe/Copenhagen",
    "DM": "America/Dominica",
    "DO": "America/Santo_Domingo",
    "DZ": "Africa/Algiers",
    "EC": "America/Guayaquil",
    "EE": "Europe/Tallinn",
    "EG": "Africa/Cairo",
    "ES": "Europe/Madrid",
    "ET": "Africa/Addis_Ababa",
    "FI": "Europe/Helsinki",
    "FJ": "Pacific/Fiji",
    "FK": "Atlantic/Stanley",
    "FM": "Pacific/Pohnpei",
    "FO": "Atlantic/Faroe",
    "FR": "Europe/Paris",
    "GA": "Africa/Libreville",
    "GB": "Europe/London",
    "GD": "America/Grenada",
    "GE": "Asia/Tbilisi",
    "GF": "America/Cayenne",
    "GG": "Europe/Guernsey",
    "GH": "Africa/Accra",
    "GI": "Europe/Gibraltar",
    "GL": "America/Nuuk",
    "GM": "Africa/Banjul",
    "GN": "Africa/Conakry",
    "GP": "America/Guadeloupe",
    "GQ": "Africa/Malabo",
    "GR": "Europe/Athens",
    "GT": "America/Guatemala",
    "GU": "Pacific/Guam",
    "GW": "Africa/Bissau",
    "GY": "America/Guyana",
    "HK": "Asia/Hong_Kong",
    "HN": "America/Tegucigalpa",
    "HR": "Europe/Zagreb",
    "HT": "America/Port-au-Prince",
    "HU": "Europe/Budapest",
    "ID": "Asia/Jakarta",
    "IE": "Europe/Dublin",
    "IL": "Asia/Jerusalem",
    "IM": "Europe/Isle_of_Man",
    "IN": "Asia/Kolkata",
    "IO": "Indian/Chagos",
    "IQ": "Asia/Baghdad",
    "IR": "Asia/Tehran",
    "IS": "Atlantic/Reykjavik",
    "IT": "Europe/Rome",
    "JE": "Europe/Jersey",
    "JM": "America/Jamaica",
    "JO": "Asia/Amman",
    "JP": "Asia/Tokyo",
    "KE": "Africa/Nairobi",
    "KG": "Asia/Bishkek",
    "KH": "Asia/Phnom_Penh",
    "KI": "Pacific/Tarawa",
    "KM": "Indian/Comoro",
    "KN": "America/St_Kitts",
    "KR": "Asia/Seoul",
    "KW": "Asia/Kuwait",
    "KY": "America/Cayman",
    "KZ": "Asia/Almaty",
    "LA": "Asia/Vientiane",
    "LB": "Asia/Beirut",
    "LC": "America/St_Lucia",
    "LI": "Europe/Vaduz",
    "LK": "Asia/Colombo",
    "LR": "Africa/Monrovia",
    "LS": "Africa/Maseru",
    "LT": "Europe/Vilnius",
    "LU": "Europe/Luxembourg",
    "LV": "Europe/Riga",
    "LY": "Africa/Tripoli",
    "MA": "Africa/Casablanca",
    "MC": "Europe/Monaco",
    "MD": "Europe/Chisinau",
    "ME": "Europe/Podgorica",
    "MF": "America/Marigot",
    "MG": "Indian/Antananarivo",
    "MK": "Europe/Skopje",
    "ML": "Africa/Bamako",
    "MM": "Asia/Yangon",
    "MN": "Asia/Ulaanbaatar",
    "MO": "Asia/Macau",
    "MP": "Pacific/Saipan",
    "MQ": "America/Martinique",
    "MR": "Africa/Nouakchott",
    "MS": "America/Montserrat",
    "MT": "Europe/Malta",
    "MU": "Indian/Mauritius",
    "MV": "Indian/Maldives",
    "MW": "Africa/Blantyre",
    "MX": "America/Mexico_City",
    "MY": "Asia/Kuala_Lumpur",
    "MZ": "Africa/Maputo",
    "NA": "Africa/Windhoek",
    "NC": "Pacific/Noumea",
    "NE": "Africa/Niamey",
    "NG": "Africa/Lagos",
    "NI": "America/Managua",
    "NL": "Europe/Amsterdam",
    "NO": "Europe/Oslo",
    "NP": "Asia/Kathmandu",
    "NZ": "Pacific/Auckland",
    "OM": "Asia/Muscat",
    "PA": "America/Panama",
    "PE": "America/Lima",
    "PF": "Pacific/Tahiti",
    "PG": "Pacific/Port_Moresby",
    "PH": "Asia/Manila",
    "PK": "Asia/Karachi",
    "PL": "Europe/Warsaw",
    "PM": "America/Miquelon",
    "PR": "America/Puerto_Rico",
    "PS": "Asia/Hebron",
    "PT": "Europe/Lisbon",
    "PW": "Pacific/Palau",
    "PY": "America/Asuncion",
    "QA": "Asia/Qatar",
    "RE": "Indian/Reunion",
    "RO": "Europe/Bucharest",
    "RS": "Europe/Belgrade",
    "RU": "Europe/Moscow",
    "RW": "Africa/Kigali",
    "SA": "Asia/Riyadh",
    "SB": "Pacific/Guadalcanal",
    "SC": "Indian/Mahe",
    "SD": "Africa/Khartoum",
    "SE": "Europe/Stockholm",
    "SG": "Asia/Singapore",
    "SI": "Europe/Ljubljana",
    "SK": "Europe/Bratislava",
    "SL": "Africa/Freetown",
    "SM": "Europe/San_Marino",
    "SN": "Africa/Dakar",
    "SO": "Africa/Mogadishu",
    "SR": "America/Paramaribo",
    "SS": "Africa/Juba",
    "ST": "Africa/Sao_Tome",
    "SV": "America/El_Salvador",
    "SX": "America/Lower_Princes",
    "SY": "Asia/Damascus",
    "SZ": "Africa/Mbabane",
    "TC": "America/Grand_Turk",
    "TD": "Africa/Ndjamena",
    "TG": "Africa/Lome",
    "TH": "Asia/Bangkok",
    "TJ": "Asia/Dushanbe",
    "TL": "Asia/Dili",
    "TM": "Asia/Ashgabat",
    "TN": "Africa/Tunis",
    "TO": "Pacific/Tongatapu",
    "TR": "Europe/Istanbul",
    "TT": "America/Port_of_Spain",
    "TW": "Asia/Taipei",
    "TZ": "Africa/Dar_es_Salaam",
    "UA": "Europe/Kyiv",
    "UG": "Africa/Kampala",
    "US": "America/New_York",
    "UY": "America/Montevideo",
    "UZ": "Asia/Tashkent",
    "VA": "Europe/Vatican",
    "VC": "America/St_Vincent",
    "VE": "America/Caracas",
    "VG": "America/Tortola",
    "VI": "America/St_Thomas",
    "VN": "Asia/Ho_Chi_Minh",
    "VU": "Pacific/Efate",
    "WS": "Pacific/Apia",
    "XK": "Europe/Belgrade",
    "YE": "Asia/Aden",
    "YT": "Indian/Mayotte",
    "ZA": "Africa/Johannesburg",
    "ZM": "Africa/Lusaka",
    "ZW": "Africa/Harare",
}


def build_residential_locations_by_iso(enum_cls: type[ProxyLocation]) -> dict[str, ProxyLocation]:
    mapping: dict[str, ProxyLocation] = {"US": enum_cls.RESIDENTIAL}
    for iso_code in sorted(RESIDENTIAL_COUNTRY_TIMEZONES):
        if iso_code == "US":
            continue
        mapping[iso_code] = getattr(enum_cls, f"RESIDENTIAL_{iso_code}")
    return mapping


def build_iso_by_residential_location(
    residential_locations_by_iso: dict[str, ProxyLocation],
) -> dict[ProxyLocation, str]:
    return {location: iso for iso, location in residential_locations_by_iso.items()}


def build_tzinfo_by_residential_location(
    residential_locations_by_iso: dict[str, ProxyLocation],
) -> dict[ProxyLocation, ZoneInfo]:
    return {
        location: ZoneInfo(RESIDENTIAL_COUNTRY_TIMEZONES[iso]) for iso, location in residential_locations_by_iso.items()
    }