"""Geography listing helpers."""

from pycen.geography.lookup import state as lookup_state

def _filter_names(gdf, query, name_cols):
    if not query:
        return gdf
    mask = False
    for col in name_cols:
        if col in gdf.columns:
            mask = mask | gdf[col].str.contains(query, case=False, na=False)
    return gdf[mask]

def _apply_limit_offset(gdf, limit, offset):
    if offset:
        gdf = gdf.iloc[int(offset):]
    if limit is not None:
        gdf = gdf.iloc[:int(limit)]
    return gdf

def _resolve_state_fips(state):
    if state is None:
        return None
    info = lookup_state(state)
    return info['fips']

def list_places(state, query=None, year=2020, limit=50, offset=0):
    '''List places for a state, optionally filtered by name.'''
    import pygris
    state_fips = _resolve_state_fips(state)
    gdf = pygris.places(year=year, state=state_fips, cb=True)
    gdf = _filter_names(gdf, query, ['NAME', 'NAMELSAD'])
    cols = [c for c in ['NAME', 'NAMELSAD', 'GEOID', 'PLACEFP', 'STATEFP'] if c in gdf.columns]
    gdf = gdf[cols].sort_values(cols[0]) if cols else gdf
    return _apply_limit_offset(gdf, limit, offset)

def list_county_boundaries(state, query=None, year=2020, limit=50, offset=0):
    '''List county boundaries for a state, optionally filtered by name.'''
    import pygris
    state_fips = _resolve_state_fips(state)
    gdf = pygris.counties(year=year, state=state_fips, cb=True)
    gdf = _filter_names(gdf, query, ['NAME'])
    cols = [c for c in ['NAME', 'GEOID', 'COUNTYFP', 'STATEFP'] if c in gdf.columns]
    gdf = gdf[cols].sort_values(cols[0]) if cols else gdf
    return _apply_limit_offset(gdf, limit, offset)

def list_counties(state, query=None, year=2020, limit=50, offset=0):
    '''Backward-compatible alias for list_county_boundaries.'''
    return list_county_boundaries(state, query=query, year=year, limit=limit, offset=offset)

def list_cbsa(query=None, year=2020, limit=50, offset=0):
    '''List CBSAs, optionally filtered by name.'''
    import pygris
    gdf = pygris.core_based_statistical_areas(year=year, cb=True)
    gdf = _filter_names(gdf, query, ['NAME'])
    cols = [c for c in ['NAME', 'GEOID', 'CBSAFP'] if c in gdf.columns]
    gdf = gdf[cols].sort_values(cols[0]) if cols else gdf
    return _apply_limit_offset(gdf, limit, offset)

def list_csa(query=None, year=2020, limit=50, offset=0):
    '''List CSAs, optionally filtered by name.'''
    import pygris
    gdf = pygris.combined_statistical_areas(year=year, cb=True)
    gdf = _filter_names(gdf, query, ['NAME'])
    cols = [c for c in ['NAME', 'GEOID', 'CSAFP'] if c in gdf.columns]
    gdf = gdf[cols].sort_values(cols[0]) if cols else gdf
    return _apply_limit_offset(gdf, limit, offset)

def list_zcta(query=None, year=2020, limit=50, offset=0):
    '''List ZCTAs, optionally filtered by name.'''
    import pygris
    gdf = pygris.zctas(year=year, cb=True)
    gdf = _filter_names(gdf, query, ['NAME'])
    cols = [c for c in ['NAME', 'GEOID', 'ZCTA5CE'] if c in gdf.columns]
    gdf = gdf[cols].sort_values(cols[0]) if cols else gdf
    return _apply_limit_offset(gdf, limit, offset)
