"""State and county lookup functions."""

import json
from pathlib import Path
import us

# load county lookup table
_COUNTY_DATA = None

def _load_county_data():
    '''Load county lookup table (lazy loading).'''
    global _COUNTY_DATA
    if _COUNTY_DATA is None:
        data_file = Path(__file__).parent.parent / 'data' / 'counties.json'
        with open(data_file) as f:
            _COUNTY_DATA = json.load(f)
    return _COUNTY_DATA

def state(name_or_fips):
    '''
    Look up state by name, abbreviation, or FIPS code.

    Parameters
    ----------
    name_or_fips : str
        State name ('California'), abbreviation ('CA'), or FIPS code ('06')

    Returns
    -------
    dict
        {'name': str, 'abbr': str, 'fips': str}

    Examples
    --------
    >>> state('California')
    {'name': 'California', 'abbr': 'CA', 'fips': '06'}
    >>> state('CA')
    {'name': 'California', 'abbr': 'CA', 'fips': '06'}
    >>> state('06')
    {'name': 'California', 'abbr': 'CA', 'fips': '06'}
    '''
    # try FIPS first
    if name_or_fips.isdigit() and len(name_or_fips) == 2:
        st = us.states.lookup(name_or_fips)
    else:
        st = us.states.lookup(name_or_fips)

    if st is None:
        raise ValueError(f"State '{name_or_fips}' not found.")

    return {
        'name': st.name,
        'abbr': st.abbr,
        'fips': st.fips
    }

def county(name_or_fips, state=None):
    '''
    Look up county by name or FIPS code.

    Parameters
    ----------
    name_or_fips : str
        County name or FIPS code
    state : str, optional
        State name, abbreviation, or FIPS code to narrow search

    Returns
    -------
    dict or str
        If state specified: {'name': str, 'fips': str, 'state_fips': str}
        If state not specified and multiple matches: prints matches

    Examples
    --------
    >>> county('Alameda', state='CA')
    {'name': 'Alameda', 'fips': '001', 'state_fips': '06'}
    '''
    counties = _load_county_data()

    # if state specified, resolve it
    if state:
        state_info = globals()['state'](state)
        state_fips = state_info['fips']

        if state_fips not in counties:
            raise ValueError(f"State FIPS '{state_fips}' not found.")

        # look up county in this state
        name_lower = name_or_fips.lower().replace(' county', '').replace(' parish', '')

        if name_or_fips.isdigit():
            county_fips = name_or_fips
        elif name_lower in counties[state_fips]:
            county_fips = counties[state_fips][name_lower]
        else:
            raise ValueError(f"County '{name_or_fips}' not found in state {state_info['name']} ({state_fips}).")

        unique_counties = {}
        for name, fips in counties[state_fips].items():
            if fips.isdigit() and fips not in unique_counties:
                unique_counties[fips] = name.title()

        county_name = unique_counties.get(county_fips, name_or_fips.title())
        return {
            'name': county_name,
            'fips': county_fips,
            'state_fips': state_fips
        }

    # if no state specified, search all states
    name_lower = name_or_fips.lower().replace(' county', '').replace(' parish', '')
    matches = []

    for state_fips, state_counties in counties.items():
        if name_lower in state_counties:
            st = us.states.lookup(state_fips)
            matches.append({
                'name': name_or_fips.title(),
                'state': st.name if st else state_fips,
                'state_fips': state_fips,
                'fips': state_counties[name_lower]
            })

    if len(matches) == 0:
        raise ValueError(f"County '{name_or_fips}' not found.")
    elif len(matches) == 1:
        return matches[0]
    else:
        # multiple matches
        msg = f"Multiple counties named '{name_or_fips}' found. Specify state or use FIPS code.\n\n"
        for m in matches:
            msg += f"  {name_or_fips.title()} County, {m['state']} (FIPS: {m['state_fips']}{m['fips']})\n"
        print(msg)
        return None

def list_counties(state):
    '''
    List all counties in a state.

    Parameters
    ----------
    state : str
        State name, abbreviation, or FIPS code

    Returns
    -------
    list of dict
        [{'name': str, 'fips': str}, ...]

    Examples
    --------
    >>> counties = list_counties('CA')
    >>> len(counties)
    58
    '''
    state_info = globals()['state'](state)
    state_fips = state_info['fips']

    counties_data = _load_county_data()

    if state_fips not in counties_data:
        raise ValueError(f"State FIPS '{state_fips}' not found.")

    # extract unique counties
    unique_counties = {}
    for name, fips in counties_data[state_fips].items():
        if not fips.isdigit():
            continue  # Skip non-FIPS entries
        if fips not in unique_counties:
            unique_counties[fips] = name.title()

    return [{'name': name, 'fips': fips} for fips, name in sorted(unique_counties.items())]
