import pandas as pd
from pandas.api.types import is_string_dtype

from upgini.utils.base_search_key_detector import BaseSearchKeyDetector


class CountrySearchKeyDetector(BaseSearchKeyDetector):
    def _is_search_key_by_name(self, column_name: str) -> bool:
        return "country" in column_name.lower()

    def _is_search_key_by_values(self, column: pd.Series) -> bool:
        if not is_string_dtype(column):
            return False

        all_count = len(column)
        is_countries_count = len(column[column.isin(self.COUNTRY_CODES)])
        return is_countries_count / all_count > 0.1

    COUNTRY_CODES = [
        "AF",
        "AL",
        "DZ",
        "AD",
        "AO",
        "AG",
        "AR",
        "AM",
        "AU",
        "AT",
        "AZ",
        "BS",
        "BH",
        "BD",
        "BB",
        "BY",
        "BE",
        "BZ",
        "BJ",
        "BT",
        "BO",
        "BA",
        "BW",
        "BR",
        "BN",
        "BG",
        "BF",
        "MM",
        "BI",
        "CV",
        "KH",
        "CM",
        "CA",
        "CF",
        "TD",
        "CL",
        "CN",
        "CO",
        "KM",
        "CG",
        "CD",
        "CR",
        "CI",
        "HR",
        "CU",
        "CY",
        "CZ",
        "DK",
        "Unknown code",
        "DJ",
        "DM",
        "DO",
        "EC",
        "EG",
        "SV",
        "GQ",
        "ER",
        "EE",
        "SZ",
        "ET",
        "FJ",
        "FI",
        "FR",
        "GA",
        "GM",
        "GE",
        "DE",
        "GH",
        "GR",
        "GD",
        "GT",
        "GN",
        "GW",
        "GY",
        "HT",
        "VA",
        "HN",
        "HU",
        "IS",
        "IN",
        "ID",
        "IR",
        "IQ",
        "IE",
        "IL",
        "IT",
        "JM",
        "JP",
        "JO",
        "KZ",
        "KE",
        "KR",
        "RS",
        "KW",
        "KG",
        "LA",
        "LV",
        "LB",
        "LR",
        "LY",
        "LI",
        "LT",
        "LU",
        "MG",
        "MW",
        "MY",
        "MV",
        "ML",
        "MT",
        "MR",
        "MU",
        "MX",
        "MD",
        "MC",
        "MN",
        "ME",
        "MA",
        "MZ",
        "NA",
        "NP",
        "NL",
        "NZ",
        "NI",
        "NE",
        "NG",
        "MK",
        "NO",
        "OM",
        "PK",
        "PA",
        "PG",
        "PY",
        "PE",
        "PH",
        "PL",
        "PT",
        "QA",
        "RO",
        "RU",
        "RW",
        "KN",
        "LC",
        "VC",
        "SM",
        "ST",
        "SA",
        "SN",
        "SC",
        "SL",
        "SG",
        "SK",
        "SI",
        "SO",
        "ZA",
        "SS",
        "ES",
        "LK",
        "SD",
        "SR",
        "SE",
        "CH",
        "SY",
        "TJ",
        "TZ",
        "TH",
        "TL",
        "TG",
        "TT",
        "TN",
        "TR",
        "US",
        "UG",
        "UA",
        "AE",
        "GB",
        "UY",
        "UZ",
        "VE",
        "VN",
        "EH",
        "YE",
        "ZM",
        "ZW",
        "Afghanistan",
        "Albania",
        "Algeria",
        "Andorra",
        "Angola",
        "Antigua and Barbuda",
        "Argentina",
        "Armenia",
        "Australia",
        "Austria",
        "Azerbaijan",
        "Bahamas",
        "Bahrain",
        "Bangladesh",
        "Barbados",
        "Belarus",
        "Belgium",
        "Belize",
        "Benin",
        "Bhutan",
        "Bolivia",
        "Bosnia and Herzegovina",
        "Botswana",
        "Brazil",
        "Brunei",
        "Bulgaria",
        "Burkina Faso",
        "Burma",
        "Burundi",
        "Cabo Verde",
        "Cambodia",
        "Cameroon",
        "Canada",
        "Central African Republic",
        "Chad",
        "Chile",
        "China",
        "Colombia",
        "Comoros",
        "Congo (Brazzaville)",
        "Congo (Kinshasa)",
        "Costa Rica",
        "Cote d'Ivoire",
        "Croatia",
        "Cuba",
        "Cyprus",
        "Czechia",
        "Denmark",
        "Diamond Princess",
        "Djibouti",
        "Dominica",
        "Dominican Republic",
        "Ecuador",
        "Egypt",
        "El Salvador",
        "Equatorial Guinea",
        "Eritrea",
        "Estonia",
        "Eswatini",
        "Ethiopia",
        "Fiji",
        "Finland",
        "France",
        "Gabon",
        "Gambia",
        "Georgia",
        "Germany",
        "Ghana",
        "Greece",
        "Grenada",
        "Guatemala",
        "Guinea",
        "Guinea-Bissau",
        "Guyana",
        "Haiti",
        "Holy See",
        "Honduras",
        "Hungary",
        "Iceland",
        "India",
        "Indonesia",
        "Iran",
        "Iraq",
        "Ireland",
        "Israel",
        "Italy",
        "Jamaica",
        "Japan",
        "Jordan",
        "Kazakhstan",
        "Kenya",
        "Korea, South",
        "Kosovo",
        "Kuwait",
        "Kyrgyzstan",
        "Laos",
        "Latvia",
        "Lebanon",
        "Liberia",
        "Libya",
        "Liechtenstein",
        "Lithuania",
        "Luxembourg",
        "MS Zaandam",
        "Madagascar",
        "Malawi",
        "Malaysia",
        "Maldives",
        "Mali",
        "Malta",
        "Mauritania",
        "Mauritius",
        "Mexico",
        "Moldova",
        "Monaco",
        "Mongolia",
        "Montenegro",
        "Morocco",
        "Mozambique",
        "Namibia",
        "Nepal",
        "Netherlands",
        "New Zealand",
        "Nicaragua",
        "Niger",
        "Nigeria",
        "North Macedonia",
        "Norway",
        "Oman",
        "Pakistan",
        "Panama",
        "Papua New Guinea",
        "Paraguay",
        "Peru",
        "Philippines",
        "Poland",
        "Portugal",
        "Qatar",
        "Romania",
        "Russia",
        "Rwanda",
        "Saint Kitts and Nevis",
        "Saint Lucia",
        "Saint Vincent and the Grenadines",
        "San Marino",
        "Sao Tome and Principe",
        "Saudi Arabia",
        "Senegal",
        "Serbia",
        "Seychelles",
        "Sierra Leone",
        "Singapore",
        "Slovakia",
        "Slovenia",
        "Somalia",
        "South Africa",
        "South Sudan",
        "Spain",
        "Sri Lanka",
        "Sudan",
        "Suriname",
        "Sweden",
        "Switzerland",
        "Syria",
        "Taiwan*",
        "Tajikistan",
        "Tanzania",
        "Thailand",
        "Timor-Leste",
        "Togo",
        "Trinidad and Tobago",
        "Tunisia",
        "Turkey",
        "US",
        "Uganda",
        "Ukraine",
        "United Arab Emirates",
        "United Kingdom",
        "Uruguay",
        "Uzbekistan",
        "Venezuela",
        "Vietnam",
        "West Bank and Gaza",
        "Western Sahara",
        "Yemen",
        "Zambia",
        "Zimbabwe",
    ]
