"""Main data acquisition functions."""

import inspect
import re
import pandas as pd
import numpy as np
import logging
from tqdm import tqdm
from pycen.utils.api import census_api_call
from pycen.utils.cache import get_cache_path, save_to_cache, load_from_cache
from pycen.geography.lookup import state as lookup_state
from pycen.geography.catalog import load_geography_catalog
from pycen.datasets import normalize_dataset

# Census API geography specifications (for tabular data)
CENSUS_GEO_ALIASES = {
    'cbsa': 'metropolitan statistical area/micropolitan statistical area',
    'core based statistical area': 'metropolitan statistical area/micropolitan statistical area',
    'csa': 'combined statistical area',
    'zcta': 'zip code tabulation area',
    'puma': 'public use microdata area',
    'cbg': 'block group',
}

# GEOID component widths for standardization
GEOID_WIDTHS = {
    'state': 2,
    'county': 3,
    'tract': 6,
    'block group': 1,
    'block': 4,
}

HIERARCHICAL_GEOS = {'state', 'county', 'tract', 'block group', 'block'}

STATE_REQUIRED_GEOS = {
    'county', 'tract', 'block group', 'block', 'place', 'county subdivision',
    'puma', 'public use microdata area', 'congressional district',
}
COUNTY_REQUIRED_GEOS = {'tract', 'block group', 'block'}

PYGRIS_FUNCTIONS = {
    'nation',
    'divisions',
    'regions',
    'states',
    'counties',
    'tracts',
    'block_groups',
    'blocks',
    'places',
    'pumas',
    'school_districts',
    'zctas',
    'congressional_districts',
    'state_legislative_districts',
    'voting_districts',
    'area_water',
    'linear_water',
    'coastline',
    'core_based_statistical_areas',
    'combined_statistical_areas',
    'metro_divisions',
    'new_england',
    'county_subdivisions',
    'urban_areas',
    'primary_roads',
    'primary_secondary_roads',
    'roads',
    'rails',
    'native_areas',
    'alaska_native_regional_corporations',
    'tribal_block_groups',
    'tribal_census_tracts',
    'tribal_subdivisions_national',
    'landmarks',
    'military',
}

CONUS_STATES = {
    '01', '04', '05', '06', '08', '09', '10', '11', '12', '13',
    '16', '17', '18', '19', '20', '21', '22', '23', '24', '25',
    '26', '27', '28', '29', '30', '31', '32', '33', '34', '35',
    '36', '37', '38', '39', '40', '41', '42', '44', '45', '46',
    '47', '48', '49', '50', '51', '53', '54', '55', '56',
}

NONCONUS_STATES = {'02', '15', '60', '66', '69', '72', '78'}

BOUNDARY_GEO_ALIASES = {
    'us': 'nation',
    'region': 'regions',
    'division': 'divisions',
    'state': 'states',
    'county': 'counties',
    'tract': 'tracts',
    'block group': 'block_groups',
    'block': 'blocks',
    'place': 'places',
    'puma': 'pumas',
    'public use microdata area': 'pumas',
    'zcta': 'zctas',
    'zip code tabulation area': 'zctas',
    'cbsa': 'core_based_statistical_areas',
    'metropolitan statistical area/micropolitan statistical area': 'core_based_statistical_areas',
    'csa': 'combined_statistical_areas',
    'combined statistical area': 'combined_statistical_areas',
    'urban area': 'urban_areas',
    'congressional district': 'congressional_districts',
    'county subdivision': 'county_subdivisions',
    'school district (elementary)': 'school_districts',
    'school district (secondary)': 'school_districts',
    'school district (unified)': 'school_districts',
    'state legislative district': 'state_legislative_districts',
    'state legislative districts': 'state_legislative_districts',
    'voting district': 'voting_districts',
    'voting districts': 'voting_districts',
}

# Dataset-specific geography restrictions
def _normalize_geo_name(value):
    text = value.strip().lower()
    text = text.replace('›', '>')
    parts = [part.strip() for part in text.split('>')]
    if len(parts) > 1:
        return ' > '.join(parts)
    return re.sub(r'\s+', ' ', text)


def _build_catalog_index(geos):
    key_map = {}
    base_map = {}
    for geo in geos:
        key = _normalize_geo_name(geo)
        parts = key.split(' > ')
        base = parts[-1]
        required_in = parts[:-1]
        entry = {
            'key': key,
            'base': base,
            'required_in': required_in,
            'for_token': f"{base}:*",
        }
        key_map[key] = entry
        base_map.setdefault(base, []).append(entry)
    return key_map, base_map


def _resolve_catalog_entry(geography, geos):
    name = _normalize_geo_name(geography)
    name = CENSUS_GEO_ALIASES.get(name, name)

    if name == 'school district':
        raise ValueError(
            "geography 'school district' is ambiguous. Use one of: "
            "'school district (elementary)', 'school district (secondary)', "
            "'school district (unified)'."
        )

    key_map, base_map = _build_catalog_index(geos)
    if name in key_map:
        return key_map[name]
    if name in base_map:
        entries = base_map[name]
        if len(entries) == 1:
            return entries[0]
        best = min(entries, key=lambda e: len(e['required_in']))
        choices = [e['key'] for e in entries]
        if best['base'] not in HIERARCHICAL_GEOS:
            def _parent_chain(entry):
                return ' > '.join(reversed(entry['required_in']))

            logging.warning(
                "PYCEN_GEO_AMBIGUOUS: base '%s' appears under multiple parent chains; "
                "defaulting to parents '%s'. Choose parents: %s",
                best['base'],
                _parent_chain(best),
                '; '.join(_parent_chain(entry) for entry in entries)
            )
        return best

    supported = ', '.join(sorted(base_map.keys()))
    raise ValueError(
        f"Geography '{geography}' not supported. Supported geographies: {supported}"
    )


def _resolve_geography(state, county):
    '''Resolve state and county to FIPS codes.'''
    state_fips = None
    county_fips = None

    if state:
        state_info = lookup_state(state)
        state_fips = state_info['fips']
        print(f"  State: {state_info['name']} ({state_info['abbr']}, FIPS: {state_fips})")

        if county:
            from pycen.geography.lookup import county as lookup_county
            county_info = lookup_county(county, state=state)
            county_fips = county_info['fips']
            print(f"  County: {county_info['name']} County (FIPS: {county_fips})")

    return state_fips, county_fips


def get_census(variables, geography, state=None, county=None, year=None,
               years=None, dataset='acs5', merge=None, clean=True,
               conus=True, cache=True, cache_dir='./pycen_cache'):
    '''
    Download Census data (non-spatial).

    Parameters
    ----------
    variables : list or dict
        Variable codes (list) or {code: name} mapping (dict)
    geography : str
        Geography level ('tract', 'block group', 'block', 'county', 'state', etc.)
    state : str, optional
        State name, abbreviation, or FIPS code
    county : str, optional
        County name or FIPS code
    year : int
        Census year (required if years is None)
    years : list of int, optional
        Multiple years to download
    dataset : str, default 'acs5'
        Data product ('acs5', 'acs1', 'dec_pl', 'dec_sf1')
    merge : str, optional
        'long' to merge multi-year data, None for separate dict
    clean : bool, default True
        Clean data (handle error codes, convert types)
    conus : bool, default True
        If True, only include continental US (48 states + DC). If False, include all states and territories
    cache : bool, default True
        Enable caching
    cache_dir : str, default './pycen_cache'
        Cache directory

    Returns
    -------
    DataFrame or dict
        Census data (non-spatial)

    Examples
    --------
    >>> df = get_census(['B19013_001E'], 'tract', state='CA', county='Alameda', year=2021, dataset='acs5')
    '''
    if years is None and year is None:
        raise ValueError("year is required")

    # multi-year support
    if years is not None:
        return _get_multi_year(variables, geography, state, county, False,
                               years, dataset, merge, clean, False, conus, cache, cache_dir)

    # single year
    return _get_single_year(variables, geography, state, county, False,
                            year, dataset, clean, False, conus, cache, cache_dir)

def get_censhp(variables, geography, state=None, county=None, year=None,
               years=None, dataset='acs5', merge=None, clean=True,
               add_area=False, conus=True, cache=True, cache_dir='./pycen_cache',
               clip_to=None, place=None, cbsa=None, csa=None, zcta=None):
    '''
    Download Census data and boundaries as a GeoDataFrame.

    Parameters
    ----------
    variables : list or dict
        Variable codes (list) or {code: name} mapping (dict)
    geography : str
        Geography level ('tract', 'block group', 'block', 'county', 'state', etc.)
    state : str or list, optional
        State name(s), abbreviation(s), or FIPS code(s)
    county : str, optional
        County name or FIPS code
    year : int
        Census year (required if years is None)
    years : list of int, optional
        Multiple years to download
    dataset : str, default 'acs5'
        Data product ('acs5', 'acs1', 'dec_pl', 'dec_sf1')
    merge : str, optional
        'long' to merge multi-year data, None for separate dict
    clean : bool, default True
        Clean data (handle error codes, convert types)
    add_area : bool, default False
        Add area_sqkm column (area in square kilometers)
    conus : bool, default True
        If True, only include continental US (48 states + DC). If False, include all states and territories
    cache : bool, default True
        Enable caching
    cache_dir : str, default './pycen_cache'
        Cache directory
    clip_to : str, optional
        Geography type to clip results to ('place', 'cbsa', 'csa', 'zcta')
    place : str, optional
        Place name or FIPS code to clip to (requires clip_to='place' or place alone)
    cbsa : str, optional
        CBSA name or code to clip to (requires clip_to='cbsa' or cbsa alone)
    csa : str, optional
        CSA name or code to clip to (requires clip_to='csa' or csa alone)
    zcta : str, optional
        ZCTA code to clip to (requires clip_to='zcta' or zcta alone)

    Returns
    -------
    GeoDataFrame or dict
        Census data with geometries (non-spatial)

    Examples
    --------
    >>> gdf = get_censhp({'B19013_001E': 'median_income'}, 'tract', state='CA', county='Alameda', year=2021, dataset='acs5')
    >>> gdf = get_censhp(['B19013_001E'], 'tract', state='CA', clip_to='place', place='Oakland city', year=2021)
    '''
    if years is None and year is None:
        raise ValueError("year is required")

    if years is not None:
        return _get_multi_year(
            variables, geography, state, county, True,
            years, dataset, merge, clean, add_area, conus, cache, cache_dir,
            clip_to, place, cbsa, csa, zcta
        )

    return _get_single_year(
        variables, geography, state, county, True,
        year, dataset, clean, add_area, conus, cache, cache_dir,
        clip_to, place, cbsa, csa, zcta
    )

def join_data(census_df, boundaries_gdf, add_area=False):
    '''
    Join Census data to boundaries on GEOID.
    '''
    census_df = census_df.copy()
    boundaries_gdf = boundaries_gdf.copy()

    census_df['GEOID'] = census_df['GEOID'].astype(str)
    boundaries_gdf['GEOID'] = boundaries_gdf['GEOID'].astype(str)

    gdf = boundaries_gdf.merge(census_df, on='GEOID', how='inner')
    if add_area:
        gdf['area_sqkm'] = gdf.to_crs(epsg=5070).area / 1e6
    return gdf

def _get_single_year(variables, geography, state, county, include_geometry, year,
                     dataset, clean, add_area, conus, cache, cache_dir,
                     clip_to=None, place=None, cbsa=None, csa=None, zcta=None):
    '''
    Get data for a single year.
    '''
    if isinstance(state, (list, tuple, set)):
        # Multi-state support: fetch per state and concatenate.
        dfs = []
        for state_item in state:
            df_state = _get_single_year(
                variables, geography, state_item, county, include_geometry, year,
                dataset, clean, add_area, conus, cache, cache_dir,
                clip_to, place, cbsa, csa, zcta
            )
            dfs.append(df_state)
        return pd.concat(dfs, ignore_index=True)

    # resolve geography
    print("Resolving geography...")
    state_fips, county_fips = _resolve_geography(state, county)
    if include_geometry:
        geos, _source = load_geography_catalog(dataset, year, cache_dir=cache_dir)
        _resolve_catalog_entry(geography, geos)
        _resolve_boundary_func(geography)

    # convert variables to list and dict
    if isinstance(variables, dict):
        var_list = list(variables.keys())
        var_names = variables
    else:
        var_list = list(variables)
        var_names = {v: v for v in var_list}
    request_vars = list(var_list)
    if not include_geometry and 'NAME' not in request_vars:
        request_vars.append('NAME')

    # check cache
    cache_params = {
        'variables': sorted(var_list),
        'geography': geography,
        'state': state_fips or '',
        'county': county_fips or '',
        'year': year,
        'dataset': dataset,
        'conus': conus
    }

    if cache:
        cache_path = get_cache_path(cache_dir, 'api', cache_params)
        cached_data = load_from_cache(cache_path, is_spatial=False)
        if cached_data is not None:
            if not include_geometry and 'NAME' not in cached_data.columns:
                cached_data = None
            elif include_geometry and 'NAME' in cached_data.columns:
                cached_data = cached_data.drop(columns=['NAME'])
        if cached_data is not None:
            print(f"Loading from cache: {cache_path}")
            df = cached_data

            # add geometry if requested
            if include_geometry:
                df = _add_geometry(df, geography, state_fips, county_fips, year,
                                  add_area, conus, cache, cache_dir, cache_params)
            return df

    # fetch from Census API
    print("\nFetching data...")
    print(f"  Product: {dataset} ({year})")
    print(f"  Geography: {geography.title()}")
    print(f"  Variables: {len(var_list)}")

    df = _fetch_census_data(
        request_vars,
        geography,
        state_fips,
        county_fips,
        year,
        dataset,
        conus,
        cache_dir=cache_dir
    )
    if include_geometry and 'NAME' in df.columns:
        df = df.drop(columns=['NAME'])

    # clean data
    if clean:
        df = _clean_data(df, var_list)

    # rename variables
    if var_names:
        rename_map = {k: v for k, v in var_names.items() if k != v}
        if rename_map:
            df = df.rename(columns=rename_map)

    # cache
    if cache:
        save_to_cache(df, cache_path, is_spatial=False)

    # add geometry
    if include_geometry:
        df = _add_geometry(
            df, geography, state_fips, county_fips, year,
            add_area, conus, cache, cache_dir, cache_params,
            clip_to, place, cbsa, csa, zcta
        )

    print("\nDone!")
    return df

def _fetch_census_data(variables, geography, state_fips, county_fips, year, dataset, conus,
                       cache_dir='./pycen_cache'):
    '''
    Fetch data from Census API.
    '''
    # build API request
    dataset_path = normalize_dataset(dataset)
    url = f"https://api.census.gov/data/{year}/{dataset_path}"

    geos, _source = load_geography_catalog(dataset, year, cache_dir=cache_dir)
    entry = _resolve_catalog_entry(geography, geos)
    geo_lower = entry['base']

    # build for clause
    for_clause = entry['for_token']

    # build 'in' clause
    in_parts = []
    values = {
        'state': state_fips,
        'county': county_fips,
    }
    for parent in entry['required_in']:
        value = values.get(parent)
        if value:
            in_parts.append(f"{parent}:{value}")
        else:
            in_parts.append(f"{parent}:*")

    in_clause = ' '.join(in_parts) if in_parts else None

    # API parameters
    params = {
        'get': ','.join(variables),
        'for': for_clause,
    }

    if in_clause:
        params['in'] = in_clause

    # make API call
    with tqdm(total=1, desc="Downloading") as pbar:
        try:
            data = census_api_call(url, params)
        except ValueError as exc:
            msg = str(exc)
            if "wildcard not allowed for 'state'" in msg and not state_fips:
                raise ValueError(
                    f"{msg}\nHint: geography '{geography}' for {dataset} {year} requires state; "
                    "pass state='MI' (or FIPS/abbr)."
                ) from None
            if "wildcard not allowed for 'county'" in msg and not county_fips:
                raise ValueError(
                    f"{msg}\nHint: geography '{geography}' for {dataset} {year} requires county; "
                    "pass county='Wayne' (or FIPS)."
                ) from None
            raise
        pbar.update(1)

    # convert to DataFrame
    if not data or len(data) < 2:
        raise ValueError("No data returned from Census API.")

    df = pd.DataFrame(data[1:], columns=data[0])

    # create GEOID
    _build_geoid(df, geo_lower, geography)

    # filter by conus if applicable
    if conus and not state_fips and 'state' in df.columns:
        df['state'] = df['state'].astype(str).str.zfill(2)
        df = df[df['state'].isin(CONUS_STATES)]

    return df

def _build_geoid(df, geo_lower, geography):
    # Hierarchical geographies (state->county->tract->bg->block).
    if geo_lower in HIERARCHICAL_GEOS:
        hierarchical_columns = ['state', 'county', 'tract', 'block group', 'block']
        geoid_parts = [col for col in hierarchical_columns if col in df.columns]
        if geoid_parts:
            df['GEOID'] = df[geoid_parts].apply(
                lambda row: ''.join(
                    str(row[col]).zfill(GEOID_WIDTHS.get(col, len(str(row[col]))))
                    for col in geoid_parts
                ),
                axis=1
            )
        return df

    # Non-hierarchical geographies.
    if geo_lower in ['place', 'county subdivision']:
        if 'state' in df.columns and geo_lower in df.columns:
            df['GEOID'] = df['state'].str.zfill(2) + df[geo_lower].astype(str).str.zfill(5)
        return df

    if geo_lower in ['puma', 'public use microdata area']:
        puma_col = 'public use microdata area' if 'public use microdata area' in df.columns else 'puma'
        if 'state' in df.columns and puma_col in df.columns:
            df['GEOID'] = df['state'].str.zfill(2) + df[puma_col].astype(str).str.zfill(5)
        return df

    if geo_lower in ['zcta', 'zip code tabulation area']:
        zcta_col = 'zip code tabulation area' if 'zip code tabulation area' in df.columns else 'zcta'
        if zcta_col in df.columns:
            df['GEOID'] = df[zcta_col].astype(str).str.zfill(5)
        return df

    if geo_lower == 'congressional district':
        if 'state' in df.columns and 'congressional district' in df.columns:
            df['GEOID'] = df['state'].str.zfill(2) + df['congressional district'].astype(str).str.zfill(2)
        return df

    if geo_lower in ['cbsa', 'csa', 'metropolitan statistical area/micropolitan statistical area',
                     'combined statistical area', 'urban area']:
        alias_map = {
            'cbsa': 'metropolitan statistical area/micropolitan statistical area',
            'csa': 'combined statistical area',
        }
        candidate_cols = []
        if geo_lower in alias_map:
            candidate_cols.extend([alias_map[geo_lower], geo_lower])
        else:
            candidate_cols.append(geo_lower)
        for col in candidate_cols:
            if col in df.columns:
                df['GEOID'] = df[col].astype(str)
                break
        return df

    if geo_lower in ['us', 'region', 'division']:
        if geo_lower in df.columns:
            df['GEOID'] = df[geo_lower].astype(str)
        return df

    raise ValueError(
        f"GEOID creation not implemented for geography '{geography}'. "
        "Add a handler in _build_geoid."
    )

def _clean_data(df, variables):
    '''
    Clean Census data (handle error codes, convert types).
    '''
    for var in variables:
        if var in df.columns:
            # convert to numeric
            df[var] = pd.to_numeric(df[var], errors='coerce')

            # replace Census error codes with NaN
            df.loc[df[var] < 0, var] = np.nan

    return df


def _normalize_boundary_geo(geography):
    geo_lower = geography.lower()
    if geo_lower in {'school district', 'school districts'}:
        raise ValueError(
            "geography 'school district' is ambiguous. Use one of: "
            "'school district (elementary)', 'school district (secondary)', "
            "'school district (unified)'."
        )
    if geo_lower in BOUNDARY_GEO_ALIASES:
        return BOUNDARY_GEO_ALIASES[geo_lower]
    normalized = re.sub(r'[^a-z0-9]+', '_', geo_lower).strip('_')
    return normalized


def _filter_kwargs(func, kwargs):
    params = inspect.signature(func).parameters
    return {k: v for k, v in kwargs.items() if k in params}


def _resolve_boundary_func(geography):
    try:
        import pygris
    except ImportError as exc:
        raise ImportError(
            "Spatial features require geopandas and pygris.\n"
            "Install with: pip install pycen"
        ) from exc
    func_name = _normalize_boundary_geo(geography)
    if func_name not in PYGRIS_FUNCTIONS:
        supported = ', '.join(sorted(PYGRIS_FUNCTIONS))
        raise ValueError(f"Geometry '{geography}' not supported for boundaries. Supported: {supported}")
    fetch_func = getattr(pygris, func_name, None)
    if fetch_func is None:
        raise ValueError(
            f"Geometry '{geography}' requires newer version of pygris with '{func_name}()' support. "
            f"Try: pip install --upgrade pygris"
        )
    return fetch_func, func_name

def _add_geometry(df, geography, state_fips, county_fips, year,
                  add_area, conus, cache, cache_dir, cache_params,
                  clip_to=None, place=None, cbsa=None, csa=None, zcta=None):
    '''
    Add geometry boundaries to data.
    '''
    try:
        import geopandas as gpd
        import pygris
    except ImportError:
        raise ImportError(
            "Spatial features require geopandas and pygris.\n"
            "Install with: pip install pycen"
        )

    # check geometry cache
    if cache:
        geo_cache_params = {**cache_params, 'type': 'geometry'}
        geo_cache_path = get_cache_path(cache_dir, 'geometries', geo_cache_params)
        cached_geo = load_from_cache(geo_cache_path, is_spatial=True)

        if cached_geo is not None:
            print(f"Loading boundaries from cache...")
            gdf_boundaries = _ensure_boundary_geoid(cached_geo)
        else:
            gdf_boundaries = _fetch_boundaries(geography, state_fips, county_fips, year, conus)
            save_to_cache(gdf_boundaries, geo_cache_path, is_spatial=True)
    else:
        gdf_boundaries = _fetch_boundaries(geography, state_fips, county_fips, year, conus)

    print("Merging data...")
    gdf = join_data(df, gdf_boundaries, add_area=add_area)
    clip_kind, clip_value = _resolve_clip_target(clip_to, place, cbsa, csa, zcta)
    if clip_kind:
        gdf = _clip_to_boundary(gdf, clip_kind, clip_value, state_fips, year, conus)
    return gdf

def _resolve_clip_target(clip_to, place, cbsa, csa, zcta):
    options = {
        'place': place,
        'cbsa': cbsa,
        'csa': csa,
        'zcta': zcta,
    }
    provided = {k: v for k, v in options.items() if v}
    if clip_to and provided:
        if clip_to not in options:
            supported = ', '.join(sorted(options))
            raise ValueError(f"clip_to '{clip_to}' not supported. Supported: {supported}")
        if clip_to not in provided:
            raise ValueError(f"clip_to '{clip_to}' requires {clip_to}=...")
        if len(provided) > 1:
            raise ValueError("Provide only one clip target (place/cbsa/csa/zcta).")
        return clip_to, provided[clip_to]
    if clip_to:
        if clip_to not in options:
            supported = ', '.join(sorted(options))
            raise ValueError(f"clip_to '{clip_to}' not supported. Supported: {supported}")
        raise ValueError(f"clip_to '{clip_to}' requires {clip_to}=...")
    if len(provided) > 1:
        raise ValueError("Provide only one clip target (place/cbsa/csa/zcta).")
    if provided:
        kind = next(iter(provided))
        return kind, provided[kind]
    return None, None

def _clip_to_boundary(gdf, clip_kind, clip_value, state_fips, year, conus):
    clip_geo = clip_kind
    boundary_state = state_fips if clip_kind == 'place' else None
    gdf_clip = _fetch_boundaries(clip_geo, boundary_state, None, year, conus)
    clip_cols = [c for c in gdf_clip.columns]
    value = str(clip_value).strip()
    is_numeric = value.isdigit()
    def _norm_text(text):
        text = text.lower()
        text = re.sub(r'[^a-z0-9]+', ' ', text)
        return re.sub(r'\s+', ' ', text).strip()
    if clip_kind == 'place':
        if not state_fips:
            raise ValueError("place clipping requires state.")
        if 'STATEFP' in clip_cols:
            gdf_clip = gdf_clip[gdf_clip['STATEFP'].astype(str).str.zfill(2) == state_fips]
        if is_numeric:
            for col in ('PLACEFP', 'GEOID'):
                if col in clip_cols:
                    gdf_clip = gdf_clip[gdf_clip[col].astype(str) == value]
                    break
        else:
            name_cols = [c for c in ('NAME', 'NAMELSAD') if c in clip_cols]
            if name_cols:
                value_norm = _norm_text(value)
                mask = False
                for col in name_cols:
                    mask = mask | gdf_clip[col].str.contains(value, case=False, na=False)
                    if value_norm:
                        col_norm = gdf_clip[col].fillna("").map(_norm_text)
                        mask = mask | col_norm.str.contains(value_norm, na=False)
                gdf_clip = gdf_clip[mask]
    elif clip_kind in ('cbsa', 'csa'):
        code_cols = ['GEOID', 'CBSAFP', 'CSAFP']
        if is_numeric:
            for col in code_cols:
                if col in clip_cols:
                    gdf_clip = gdf_clip[gdf_clip[col].astype(str) == value]
                    break
        else:
            name_cols = [c for c in ('NAME', 'NAMELSAD') if c in clip_cols]
            if name_cols:
                value_norm = _norm_text(value)
                mask = False
                for col in name_cols:
                    mask = mask | gdf_clip[col].str.contains(value, case=False, na=False)
                    if value_norm:
                        col_norm = gdf_clip[col].fillna("").map(_norm_text)
                        mask = mask | col_norm.str.contains(value_norm, na=False)
                gdf_clip = gdf_clip[mask]
    elif clip_kind == 'zcta':
        code_cols = ['GEOID', 'ZCTA5CE', 'ZCTA5CE20', 'ZCTA5CE10']
        for col in code_cols:
            if col in clip_cols:
                gdf_clip = gdf_clip[gdf_clip[col].astype(str) == value]
                break
    if gdf_clip.empty:
        raise ValueError(f"No boundary match found for {clip_kind}='{clip_value}'.")
    if gdf.crs is not None and gdf_clip.crs is not None and gdf.crs != gdf_clip.crs:
        gdf_clip = gdf_clip.to_crs(gdf.crs)
    clip_geom = gdf_clip.geometry.unary_union
    return gdf[gdf.geometry.intersects(clip_geom)]

def _ensure_boundary_geoid(gdf):
    if 'GEOID' in gdf.columns:
        return gdf
    for col in ('GEOID20', 'GEOID10', 'GEOIDFQ'):
        if col in gdf.columns:
            gdf = gdf.copy()
            gdf['GEOID'] = gdf[col].astype(str)
            return gdf
    return gdf

def _fetch_boundaries(geography, state_fips, county_fips, year, conus):
    '''Fetch boundaries from pygris.'''
    fetch_func, func_name = _resolve_boundary_func(geography)

    # Build kwargs for pygris function
    kwargs = {
        'year': year,
        'cb': True,
        'state': state_fips,
        'county': county_fips,
    }

    if func_name == 'blocks' and (not state_fips or not county_fips):
        raise ValueError("Block geometry requires both state and county parameters")

    print("Fetching boundaries...")
    gdf = fetch_func(**_filter_kwargs(fetch_func, kwargs))
    gdf = _ensure_boundary_geoid(gdf)

    if 'GEOID' not in gdf.columns:
        cols = ", ".join(sorted(gdf.columns))
        raise ValueError(
            "Boundaries are missing GEOID for join. "
            f"Available columns: {cols}"
        )

    # filter by conus if applicable
    if conus and not state_fips:
        if 'STATEFP' in gdf.columns:
            gdf['STATEFP'] = gdf['STATEFP'].astype(str).str.zfill(2)
            gdf = gdf[gdf['STATEFP'].isin(CONUS_STATES)]
        else:
            # fallback for geometries without state codes (e.g., CBSA/CSA)
            gdf_ll = gdf
            if gdf.crs is not None and not gdf.crs.is_geographic:
                gdf_ll = gdf.to_crs(epsg=4326)
            points = gdf_ll.geometry.representative_point()
            gdf = gdf[
                (points.x >= -130) & (points.x <= -60) &
                (points.y >= 20) & (points.y <= 55)
            ]

    return gdf

def _get_multi_year(variables, geography, state, county, include_geometry,
                    years, dataset, merge, clean, add_area, conus, cache, cache_dir,
                    clip_to=None, place=None, cbsa=None, csa=None, zcta=None):
    '''
    Get data for multiple years.
    '''
    datasets = {}

    for year in tqdm(years, desc="Downloading years"):
        print(f"\n--- Year {year} ---")
        df = _get_single_year(
            variables, geography, state, county, include_geometry,
            year, dataset, clean, add_area, conus, cache, cache_dir,
            clip_to, place, cbsa, csa, zcta
        )
        datasets[year] = df

    if merge == 'long':
        # merge into long format
        dfs = []
        for year, df in datasets.items():
            df['year'] = year
            dfs.append(df)

        return pd.concat(dfs, ignore_index=True)
    else:
        return datasets

def get_boundaries(geography, state=None, county=None, year=None, conus=True):
    '''
    Get boundaries only (no Census data).

    Parameters
    ----------
    geography : str
        Geography level
    state : str or list, optional
        State name(s), abbreviation(s), or FIPS code(s)
    county : str, optional
        County name or FIPS
    year : int
        Year for boundaries (required)
    conus : bool, default True
        If True, only include continental US (48 states + DC). If False, include all states and territories

    Returns
    -------
    GeoDataFrame
        Boundaries with GEOID and geometry

    Examples
    --------
    >>> gdf = get_boundaries('tract', state='CA', county='Alameda', year=2021)
    '''
    if year is None:
        raise ValueError("year is required.")

    state_fips, county_fips = _resolve_geography(state, county)
    return _fetch_boundaries(geography, state_fips, county_fips, year, conus)
