from enum import StrEnum


class USState(StrEnum):
    AL = "AL"
    AK = "AK"
    AZ = "AZ"
    AR = "AR"
    CA = "CA"
    CO = "CO"
    CT = "CT"
    DE = "DE"
    FL = "FL"
    GA = "GA"
    HI = "HI"
    ID = "ID"
    IL = "IL"
    IN = "IN"
    IA = "IA"
    KS = "KS"
    KY = "KY"
    LA = "LA"
    ME = "ME"
    MD = "MD"
    MA = "MA"
    MI = "MI"
    MN = "MN"
    MS = "MS"
    MO = "MO"
    MT = "MT"
    NE = "NE"
    NV = "NV"
    NH = "NH"
    NJ = "NJ"
    NM = "NM"
    NY = "NY"
    NC = "NC"
    ND = "ND"
    OH = "OH"
    OK = "OK"
    OR = "OR"
    PA = "PA"
    RI = "RI"
    SC = "SC"
    SD = "SD"
    TN = "TN"
    TX = "TX"
    UT = "UT"
    VT = "VT"
    VA = "VA"
    WA = "WA"
    WV = "WV"
    WI = "WI"
    WY = "WY"


class USTerritory(StrEnum):
    DC = "DC"
    AS = "AS"
    GU = "GU"
    MP = "MP"
    PR = "PR"
    UM = "UM"
    VI = "VI"


USStateOrTerritory = USState | USTerritory


STATE_NAME_TO_CODE: dict[str, str] = {
    "alabama": "AL",
    "alaska": "AK",
    "arizona": "AZ",
    "arkansas": "AR",
    "california": "CA",
    "colorado": "CO",
    "connecticut": "CT",
    "delaware": "DE",
    "florida": "FL",
    "georgia": "GA",
    "hawaii": "HI",
    "idaho": "ID",
    "illinois": "IL",
    "indiana": "IN",
    "iowa": "IA",
    "kansas": "KS",
    "kentucky": "KY",
    "louisiana": "LA",
    "maine": "ME",
    "maryland": "MD",
    "massachusetts": "MA",
    "michigan": "MI",
    "minnesota": "MN",
    "mississippi": "MS",
    "missouri": "MO",
    "montana": "MT",
    "nebraska": "NE",
    "nevada": "NV",
    "new hampshire": "NH",
    "new jersey": "NJ",
    "new mexico": "NM",
    "new york": "NY",
    "north carolina": "NC",
    "north dakota": "ND",
    "ohio": "OH",
    "oklahoma": "OK",
    "oregon": "OR",
    "pennsylvania": "PA",
    "rhode island": "RI",
    "south carolina": "SC",
    "south dakota": "SD",
    "tennessee": "TN",
    "texas": "TX",
    "utah": "UT",
    "vermont": "VT",
    "virginia": "VA",
    "washington": "WA",
    "west virginia": "WV",
    "wisconsin": "WI",
    "wyoming": "WY",
    "district of columbia": "DC",
    "american samoa": "AS",
    "guam": "GU",
    "northern mariana islands": "MP",
    "puerto rico": "PR",
    "united states minor outlying islands": "UM",
    "u.s. virgin islands": "VI",
}


def parse_us_state_or_territory(value: str | None) -> USStateOrTerritory | None:
    """
    Parse a US state or territory from either a full name or 2-letter code. Case insensitive.

    Args:
        value: A state/territory name (e.g., "New Mexico", "Puerto Rico") or code (e.g., "NM", "PR").

    Returns:
        The corresponding USState or USTerritory enum value, or None if not found.
    """
    if not value:
        return None

    normalized = value.strip().lower()
    if len(normalized) == 2:
        code = normalized.upper()
        try:
            return USState(code)
        except ValueError:
            pass
        try:
            return USTerritory(code)
        except ValueError:
            return None

    code = STATE_NAME_TO_CODE.get(normalized)
    if code:
        try:
            return USState(code)
        except ValueError:
            return USTerritory(code)

    return None


def parse_us_state(value: str | None) -> USState | None:
    """
    Parse a US state (not territory) from either a full name or 2-letter code. Case insensitive.

    Args:
        value: A state name (e.g., "New Mexico") or code (e.g., "NM").

    Returns:
        The corresponding USState enum value, or None if not found.
    """
    result = parse_us_state_or_territory(value)
    if isinstance(result, USState):
        return result
    return None

