hyper-connector-mullvad 0.0.2

use std::{fmt, str::FromStr};

use crate::Error;

/// A geographic restriction applied to discovered Mullvad relays.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum LocationFilter {
    /// An ISO 3166-1 alpha-2 country code.
    Country(String),
    /// A continent containing one or more relay countries.
    Continent(Continent),
}

/// A supported continent/region.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Continent {
    Africa,
    Asia,
    Europe,
    NorthAmerica,
    SouthAmerica,
    Oceania,
}

impl LocationFilter {
    /// Parse a country code, English country name, or continent name.
    pub fn parse(value: &str) -> Result<Self, Error> {
        value.parse()
    }

    /// Construct a country filter from a code or English country name.
    pub fn country(value: &str) -> Result<Self, Error> {
        country_code(value)
            .map(Self::Country)
            .ok_or_else(|| Error::InvalidLocationFilter(value.trim().to_owned()))
    }

    /// Construct a continent filter.
    pub fn continent(value: &str) -> Result<Self, Error> {
        continent(value)
            .map(Self::Continent)
            .ok_or_else(|| Error::InvalidLocationFilter(value.trim().to_owned()))
    }

    pub(crate) fn matches(&self, code: &str) -> bool {
        match self {
            Self::Country(wanted) => wanted == code,
            Self::Continent(c) => country_continent(code) == Some(*c),
        }
    }
}

impl FromStr for LocationFilter {
    type Err = Error;
    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match (continent(value), country_code(value)) {
            (Some(_), Some(_)) => Err(Error::InvalidLocationFilter(value.trim().to_owned())),
            (Some(c), None) => Ok(Self::Continent(c)),
            (None, Some(code)) => Ok(Self::Country(code)),
            (None, None) => Err(Error::InvalidLocationFilter(value.trim().to_owned())),
        }
    }
}

impl fmt::Display for Continent {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

fn normalized(value: &str) -> String {
    value
        .trim()
        .to_ascii_lowercase()
        .replace([' ', '-', '_'], "")
}

fn continent(value: &str) -> Option<Continent> {
    match normalized(value).as_str() {
        "africa" => Some(Continent::Africa),
        "asia" => Some(Continent::Asia),
        "europe" => Some(Continent::Europe),
        "northamerica" => Some(Continent::NorthAmerica),
        "southamerica" => Some(Continent::SouthAmerica),
        "oceania" | "australia" => Some(Continent::Oceania),
        _ => None,
    }
}

fn country_code(value: &str) -> Option<String> {
    let n = normalized(value);
    if n.len() == 2 && n.bytes().all(|b| b.is_ascii_alphabetic()) {
        return Some(n);
    }
    let code = match n.as_str() {
        "sweden" => "se",
        "unitedstates" | "usa" => "us",
        "unitedkingdom" | "greatbritain" => "gb",
        "australia" => "au",
        "austria" => "at",
        "belgium" => "be",
        "brazil" => "br",
        "bulgaria" => "bg",
        "canada" => "ca",
        "chile" => "cl",
        "colombia" => "co",
        "croatia" => "hr",
        "czechia" | "czechrepublic" => "cz",
        "denmark" => "dk",
        "estonia" => "ee",
        "finland" => "fi",
        "france" => "fr",
        "germany" => "de",
        "greece" => "gr",
        "hongkong" => "hk",
        "hungary" => "hu",
        "ireland" => "ie",
        "israel" => "il",
        "italy" => "it",
        "japan" => "jp",
        "latvia" => "lv",
        "luxembourg" => "lu",
        "mexico" => "mx",
        "netherlands" => "nl",
        "newzealand" => "nz",
        "norway" => "no",
        "poland" => "pl",
        "portugal" => "pt",
        "romania" => "ro",
        "serbia" => "rs",
        "singapore" => "sg",
        "slovakia" => "sk",
        "slovenia" => "si",
        "southafrica" => "za",
        "southkorea" | "korea" => "kr",
        "spain" => "es",
        "switzerland" => "ch",
        "taiwan" => "tw",
        "turkey" | "turkiye" => "tr",
        "unitedarabemirates" => "ae",
        _ => return None,
    };
    Some(code.into())
}

fn country_continent(code: &str) -> Option<Continent> {
    const EUROPE: &[&str] = &[
        "al", "at", "be", "bg", "ch", "cz", "de", "dk", "ee", "es", "fi", "fr", "gb", "gr", "hr",
        "hu", "ie", "is", "it", "lt", "lu", "lv", "md", "mk", "mt", "nl", "no", "pl", "pt", "ro",
        "rs", "se", "si", "sk", "ua",
    ];
    const ASIA: &[&str] = &[
        "ae", "hk", "id", "il", "in", "jp", "kr", "my", "ph", "sg", "th", "tr", "tw", "vn",
    ];
    const NORTH_AMERICA: &[&str] = &["ca", "cr", "gt", "mx", "pa", "us"];
    const SOUTH_AMERICA: &[&str] = &["ar", "bo", "br", "cl", "co", "ec", "pe", "py", "uy", "ve"];
    const OCEANIA: &[&str] = &["au", "nz"];
    const AFRICA: &[&str] = &["dz", "eg", "gh", "ke", "ma", "ng", "tn", "za"];
    [
        (EUROPE, Continent::Europe),
        (ASIA, Continent::Asia),
        (NORTH_AMERICA, Continent::NorthAmerica),
        (SOUTH_AMERICA, Continent::SouthAmerica),
        (OCEANIA, Continent::Oceania),
        (AFRICA, Continent::Africa),
    ]
    .into_iter()
    .find_map(|(codes, c)| codes.contains(&code).then_some(c))
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn parses_locations() {
        assert_eq!(
            LocationFilter::parse(" Sweden ").unwrap(),
            LocationFilter::Country("se".into())
        );
        assert!(LocationFilter::parse("Europe").unwrap().matches("se"));
        assert!(
            LocationFilter::parse("North America")
                .unwrap()
                .matches("us")
        );
        assert!(LocationFilter::parse("Oceania").unwrap().matches("au"));
        assert!(LocationFilter::parse("Australia").is_err());
        assert!(LocationFilter::country("Australia").unwrap().matches("au"));
        assert!(LocationFilter::parse("Atlantis").is_err());
    }
}