"""Assign United States Census regions from state or territory codes."""

from __future__ import annotations

from typing import Any, Dict, Optional

NORTHEAST = {"CT", "ME", "MA", "NH", "RI", "VT", "NJ", "NY", "PA"}
MIDWEST = {"IL", "IN", "MI", "OH", "WI", "IA", "KS", "MN", "MO", "NE", "ND", "SD"}
SOUTH = {
    "DE",
    "FL",
    "GA",
    "MD",
    "NC",
    "SC",
    "VA",
    "WV",
    "DC",
    "AL",
    "KY",
    "MS",
    "TN",
    "AR",
    "LA",
    "OK",
    "TX",
}
WEST = {"AZ", "CO", "ID", "MT", "NV", "NM", "UT", "WY", "AK", "CA", "HI", "OR", "WA"}

REGION_BY_STATE = {code: "Northeast" for code in NORTHEAST}
REGION_BY_STATE.update({code: "Midwest" for code in MIDWEST})
REGION_BY_STATE.update({code: "South" for code in SOUTH})
REGION_BY_STATE.update({code: "West" for code in WEST})

STATE_NAME_TO_CODE = {
    "ALABAMA": "AL",
    "ALASKA": "AK",
    "ARIZONA": "AZ",
    "ARKANSAS": "AR",
    "CALIFORNIA": "CA",
    "COLORADO": "CO",
    "CONNECTICUT": "CT",
    "DELAWARE": "DE",
    "DISTRICT OF COLUMBIA": "DC",
    "WASHINGTON DC": "DC",
    "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",
}

STATE_ABBREV_SYNONYMS = {
    "N.Y.": "NY",
    "CALIF": "CA",
    "PENN": "PA",
}


def _normalize_state(value: Optional[str]) -> str:
    if not value:
        return ""
    upper = str(value).strip().upper()
    if not upper:
        return ""
    if upper in STATE_ABBREV_SYNONYMS:
        upper = STATE_ABBREV_SYNONYMS[upper]
    return STATE_NAME_TO_CODE.get(upper, upper)


def _assign_us_region(code: str) -> str:
    return REGION_BY_STATE.get(code, "")


def assign_us_region(value: str | None) -> Dict[str, Any]:
    """Return the Census region (Northeast, Midwest, South, West) for a US state code."""

    code = _normalize_state(value)
    if not code:
        return {"value": "", "meta": {"reason": "empty", "input": value}}

    if len(code) != 2:
        return {"value": "", "meta": {"reason": "invalid_state", "input": value}}

    region = _assign_us_region(code)
    if not region:
        return {"value": "", "meta": {"reason": "non_us_state", "code": code}}

    return {"value": region, "meta": {"state": code, "region": region}}
