# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Utility data models and classes for common tools and mappings.
"""

class CountryCodes:
    """
    Centralized country code mapping for different services.
    
    Maps ISO-2 country codes to service-specific country names.
    Each service may have different requirements for country names.
    """
    
    # Mapping: ISO-2 code -> (Tavily country name, Perplexity country name)
    _COUNTRY_MAP = {
        "AF": ("afghanistan", "Afghanistan"),
        "AL": ("albania", "Albania"),
        "DZ": ("algeria", "Algeria"),
        "AD": ("andorra", "Andorra"),
        "AO": ("angola", "Angola"),
        "AR": ("argentina", "Argentina"),
        "AM": ("armenia", "Armenia"),
        "AU": ("australia", "Australia"),
        "AT": ("austria", "Austria"),
        "AZ": ("azerbaijan", "Azerbaijan"),
        "BS": ("bahamas", "Bahamas"),
        "BH": ("bahrain", "Bahrain"),
        "BD": ("bangladesh", "Bangladesh"),
        "BB": ("barbados", "Barbados"),
        "BY": ("belarus", "Belarus"),
        "BE": ("belgium", "Belgium"),
        "BZ": ("belize", "Belize"),
        "BJ": ("benin", "Benin"),
        "BT": ("bhutan", "Bhutan"),
        "BO": ("bolivia", "Bolivia"),
        "BA": ("bosnia and herzegovina", "Bosnia and Herzegovina"),
        "BW": ("botswana", "Botswana"),
        "BR": ("brazil", "Brazil"),
        "BN": ("brunei", "Brunei"),
        "BG": ("bulgaria", "Bulgaria"),
        "BF": ("burkina faso", "Burkina Faso"),
        "BI": ("burundi", "Burundi"),
        "KH": ("cambodia", "Cambodia"),
        "CM": ("cameroon", "Cameroon"),
        "CA": ("canada", "Canada"),
        "CV": ("cape verde", "Cape Verde"),
        "CF": ("central african republic", "Central African Republic"),
        "TD": ("chad", "Chad"),
        "CL": ("chile", "Chile"),
        "CN": ("china", "China"),
        "CO": ("colombia", "Colombia"),
        "KM": ("comoros", "Comoros"),
        "CG": ("congo", "Congo"),
        "CR": ("costa rica", "Costa Rica"),
        "HR": ("croatia", "Croatia"),
        "CU": ("cuba", "Cuba"),
        "CY": ("cyprus", "Cyprus"),
        "CZ": ("czech republic", "Czech Republic"),
        "DK": ("denmark", "Denmark"),
        "DJ": ("djibouti", "Djibouti"),
        "DO": ("dominican republic", "Dominican Republic"),
        "EC": ("ecuador", "Ecuador"),
        "EG": ("egypt", "Egypt"),
        "SV": ("el salvador", "El Salvador"),
        "GQ": ("equatorial guinea", "Equatorial Guinea"),
        "ER": ("eritrea", "Eritrea"),
        "EE": ("estonia", "Estonia"),
        "ET": ("ethiopia", "Ethiopia"),
        "FJ": ("fiji", "Fiji"),
        "FI": ("finland", "Finland"),
        "FR": ("france", "France"),
        "GA": ("gabon", "Gabon"),
        "GM": ("gambia", "Gambia"),
        "GE": ("georgia", "Georgia"),
        "DE": ("germany", "Germany"),
        "GH": ("ghana", "Ghana"),
        "GR": ("greece", "Greece"),
        "GT": ("guatemala", "Guatemala"),
        "GN": ("guinea", "Guinea"),
        "HT": ("haiti", "Haiti"),
        "HN": ("honduras", "Honduras"),
        "HU": ("hungary", "Hungary"),
        "IS": ("iceland", "Iceland"),
        "IN": ("india", "India"),
        "ID": ("indonesia", "Indonesia"),
        "IR": ("iran", "Iran"),
        "IQ": ("iraq", "Iraq"),
        "IE": ("ireland", "Ireland"),
        "IL": ("israel", "Israel"),
        "IT": ("italy", "Italy"),
        "JM": ("jamaica", "Jamaica"),
        "JP": ("japan", "Japan"),
        "JO": ("jordan", "Jordan"),
        "KZ": ("kazakhstan", "Kazakhstan"),
        "KE": ("kenya", "Kenya"),
        "KW": ("kuwait", "Kuwait"),
        "KG": ("kyrgyzstan", "Kyrgyzstan"),
        "LV": ("latvia", "Latvia"),
        "LB": ("lebanon", "Lebanon"),
        "LS": ("lesotho", "Lesotho"),
        "LR": ("liberia", "Liberia"),
        "LY": ("libya", "Libya"),
        "LI": ("liechtenstein", "Liechtenstein"),
        "LT": ("lithuania", "Lithuania"),
        "LU": ("luxembourg", "Luxembourg"),
        "MG": ("madagascar", "Madagascar"),
        "MW": ("malawi", "Malawi"),
        "MY": ("malaysia", "Malaysia"),
        "MV": ("maldives", "Maldives"),
        "ML": ("mali", "Mali"),
        "MT": ("malta", "Malta"),
        "MR": ("mauritania", "Mauritania"),
        "MU": ("mauritius", "Mauritius"),
        "MX": ("mexico", "Mexico"),
        "MD": ("moldova", "Moldova"),
        "MC": ("monaco", "Monaco"),
        "MN": ("mongolia", "Mongolia"),
        "ME": ("montenegro", "Montenegro"),
        "MA": ("morocco", "Morocco"),
        "MZ": ("mozambique", "Mozambique"),
        "MM": ("myanmar", "Myanmar"),
        "NA": ("namibia", "Namibia"),
        "NP": ("nepal", "Nepal"),
        "NL": ("netherlands", "Netherlands"),
        "NZ": ("new zealand", "New Zealand"),
        "NI": ("nicaragua", "Nicaragua"),
        "NE": ("niger", "Niger"),
        "NG": ("nigeria", "Nigeria"),
        "KP": ("north korea", "North Korea"),
        "MK": ("north macedonia", "North Macedonia"),
        "NO": ("norway", "Norway"),
        "OM": ("oman", "Oman"),
        "PK": ("pakistan", "Pakistan"),
        "PA": ("panama", "Panama"),
        "PG": ("papua new guinea", "Papua New Guinea"),
        "PY": ("paraguay", "Paraguay"),
        "PE": ("peru", "Peru"),
        "PH": ("philippines", "Philippines"),
        "PL": ("poland", "Poland"),
        "PT": ("portugal", "Portugal"),
        "QA": ("qatar", "Qatar"),
        "RO": ("romania", "Romania"),
        "RU": ("russia", "Russia"),
        "RW": ("rwanda", "Rwanda"),
        "SA": ("saudi arabia", "Saudi Arabia"),
        "SN": ("senegal", "Senegal"),
        "RS": ("serbia", "Serbia"),
        "SG": ("singapore", "Singapore"),
        "SK": ("slovakia", "Slovakia"),
        "SI": ("slovenia", "Slovenia"),
        "SO": ("somalia", "Somalia"),
        "ZA": ("south africa", "South Africa"),
        "KR": ("south korea", "South Korea"),
        "SS": ("south sudan", "South Sudan"),
        "ES": ("spain", "Spain"),
        "LK": ("sri lanka", "Sri Lanka"),
        "SD": ("sudan", "Sudan"),
        "SE": ("sweden", "Sweden"),
        "CH": ("switzerland", "Switzerland"),
        "SY": ("syria", "Syria"),
        "TW": ("taiwan", "Taiwan"),
        "TJ": ("tajikistan", "Tajikistan"),
        "TZ": ("tanzania", "Tanzania"),
        "TH": ("thailand", "Thailand"),
        "TG": ("togo", "Togo"),
        "TT": ("trinidad and tobago", "Trinidad and Tobago"),
        "TN": ("tunisia", "Tunisia"),
        "TR": ("turkey", "Turkey"),
        "TM": ("turkmenistan", "Turkmenistan"),
        "UG": ("uganda", "Uganda"),
        "UA": ("ukraine", "Ukraine"),
        "AE": ("united arab emirates", "United Arab Emirates"),
        "GB": ("united kingdom", "United Kingdom"),
        "US": ("united states", "United States"),
        "UY": ("uruguay", "Uruguay"),
        "UZ": ("uzbekistan", "Uzbekistan"),
        "VE": ("venezuela", "Venezuela"),
        "VN": ("vietnam", "Vietnam"),
        "YE": ("yemen", "Yemen"),
        "ZM": ("zambia", "Zambia"),
        "ZW": ("zimbabwe", "Zimbabwe"),
    }
    
    @classmethod
    def getForTavily(cls, isoCode: str) -> str:
        """
        Get Tavily-compatible country name from ISO-2 code.
        
        Args:
            isoCode: ISO-2 country code (e.g., "CH", "ch", "US", "us")
            
        Returns:
            Country name in lowercase as required by Tavily (e.g., "switzerland", "united states")
        """
        # Convert to uppercase for lookup
        isoCodeUpper = isoCode.upper() if isoCode else ""
        mapping = cls._COUNTRY_MAP.get(isoCodeUpper)
        return mapping[0] if mapping else isoCode
    
    @classmethod
    def getForPerplexity(cls, isoCode: str) -> str:
        """
        Get Perplexity-compatible country name from ISO-2 code.
        
        Args:
            isoCode: ISO-2 country code (e.g., "CH", "US")
            
        Returns:
            Full country name as required by Perplexity (e.g., "Switzerland", "United States")
        """
        isoCodeUpper = isoCode.upper()
        mapping = cls._COUNTRY_MAP.get(isoCodeUpper)
        return mapping[1] if mapping else isoCode
    
    @classmethod
    def isValid(cls, isoCode: str) -> bool:
        """
        Check if ISO-2 code is valid.
        
        Args:
            isoCode: ISO-2 country code to check
            
        Returns:
            True if valid, False otherwise
        """
        return isoCode.upper() in cls._COUNTRY_MAP

