"""Census data, geographic boundaries."""

import logging
import os

import geopandas as gpd
import pandas as pd
from pyproj import CRS

logger = logging.getLogger(__name__)


def _project_if_needed(gdf, target_epsg):
    """Project a GeoDataFrame only when its CRS differs from the target EPSG."""
    if gdf.crs is not None:
        current_epsg = gdf.crs.to_epsg()
        if current_epsg == int(target_epsg):
            return gdf
    return gdf.to_crs(epsg=target_epsg)


def collect_census_data(state_fips_code, county_fips_codes, year, census_data_file, geo_level='county', census_api_key=None):
    """
    Collect census data at specified geographic level (county, tract, or CBG).

    Parameters
    ----------
    state_fips_code : str
        FIPS code for the state
    county_fips_codes : list or str
        List of county FIPS codes or comma-separated string
    year : int
        Census year
    census_data_file : str
        Path to save the CSV output
    geo_level : str
        Geographic level for data collection: 'county', 'tract', or 'cbg'
        Default is 'county'
    census_api_key : str, optional
        Census API key. If not provided, reads from the ``CENSUS_API_KEY``
        environment variable.

    Returns
    -------
    pandas.DataFrame
        DataFrame containing population data for the specified geographic level
    """
    # Validate geo_level parameter
    valid_levels = ['county', 'tract', 'cbg']
    if geo_level.lower() not in valid_levels:
        raise ValueError(f"Invalid geo_level '{geo_level}'. Must be one of: {', '.join(valid_levels)}")

    geo_level = geo_level.lower()

    # Check if the output file already exists
    if os.path.exists(census_data_file):
        logger.info("Loading existing %s data from %s", geo_level, census_data_file)
        return pd.read_csv(census_data_file, dtype={'GEOID': str})

    # Resolve Census API key: parameter > env var
    if not census_api_key:
        census_api_key = os.environ.get("CENSUS_API_KEY", "").strip()
    if not census_api_key:
        raise ValueError(
            "Census API key not found. Provide it via the census_api_key parameter "
            "or the CENSUS_API_KEY environment variable. "
            "Get a free key at https://api.census.gov/data/key_signup.html"
        )

    logger.info("Collecting %s data for year %s...", geo_level.upper(), year)

    # Initialize the Census API
    from census import Census
    c = Census(census_api_key, year=year)

    # Convert list of county FIPS to comma-separated string if it's a list
    if isinstance(county_fips_codes, list):
        county_fips_string = ','.join(county_fips_codes)
    else:
        county_fips_string = county_fips_codes

    logger.info("Downloading population data for %ss...", geo_level)

    try:
        # Different API calls based on geographic level
        if geo_level == 'county':
            census_data = c.acs5.state_county(
                fields=('NAME', 'B01003_001E'),  # B01003_001E is total population
                state_fips=state_fips_code,
                county_fips=county_fips_string
            )
        elif geo_level == 'tract':
            census_data = c.acs5.state_county_tract(
                fields=('NAME', 'B01003_001E'),
                state_fips=state_fips_code,
                county_fips=county_fips_string,
                tract='*'  # Request all tracts
            )
        elif geo_level == 'cbg':
            census_data = c.acs5.state_county_blockgroup(
                fields=('NAME', 'B01003_001E'),
                state_fips=state_fips_code,
                county_fips=county_fips_string,
                blockgroup='*'  # Request all block groups
            )
        else:
            raise ValueError(f"Invalid geo_level '{geo_level}'. Must be one of: {', '.join(valid_levels)}")

        # Create a DataFrame from the census data
        df = pd.DataFrame(census_data)

        # Rename columns for clarity
        df = df.rename(columns={'B01003_001E': 'population', 'NAME': 'name'})

        # Create GEOID based on geographic level
        if geo_level == 'county':
            df['GEOID'] = df['state'] + df['county']
        elif geo_level == 'tract':
            df['GEOID'] = df['state'] + df['county'] + df['tract']
        elif geo_level == 'cbg':
            df['GEOID'] = df['state'] + df['county'] + df['tract'] + df['block group']

        # Convert population to numeric
        df['population'] = pd.to_numeric(df['population'], errors='coerce')

        # Save the raw census data
        if census_data_file:
            os.makedirs(os.path.dirname(census_data_file), exist_ok=True)
            df.to_csv(census_data_file, index=False)
            logger.info("%s population data saved to %s", geo_level.capitalize(), census_data_file)

        return df

    except Exception as e:
        logger.error("Error downloading Census data: %s", e)
        raise


def collect_tract_boundaries(state_fips_code, county_fips_codes, year):
    """
    Download census tract boundaries from TIGER/Line shapefiles.

    Parameters
    ----------
    state_fips_code : str
        FIPS code for the state
    county_fips_codes : list
        List of county FIPS codes
    year : int
        Reference year for boundaries

    Returns
    -------
    geopandas.GeoDataFrame
        GeoDataFrame containing tract boundaries
    """
    try:
        # Download geographic boundaries
        geo_url = f"https://www2.census.gov/geo/tiger/TIGER{year}/TRACT/tl_{year}_{state_fips_code}_tract.zip"
        geo_data = gpd.read_file(geo_url)

        # Filter for counties of interest
        geo_data = geo_data[geo_data['COUNTYFP'].isin(county_fips_codes)]
    except Exception as e:
        logger.error("Failed to retrieve geographic boundaries: %s", e)
        raise
    return geo_data

def collect_geographic_boundaries(
    state_fips_code,
    county_fips_codes,
    year,
    area_name,
    geo_level,
    work_dir,
    target_epsg=4326,
    cartographic=False,
):
    """Collect geographic boundaries at the specified level, caching to GeoJSON.

    Parameters
    ----------
    state_fips_code : str
        FIPS code for the state.
    county_fips_codes : list
        List of county FIPS codes to filter by.
    year : int
        Reference year for boundary data.
    area_name : str
        Human-readable name for the study area (used in the output filename).
    geo_level : str
        Geographic level: ``'county'``, ``'cbg'``, ``'tract'``, or ``'taz'``.
    work_dir : str
        Directory where the cached GeoJSON file will be stored.
    target_epsg : int, optional
        EPSG code for the cached/output boundary CRS. Defaults to ``4326``.
    cartographic : bool, optional
        If ``True`` and ``geo_level='county'``, use generalized cartographic
        boundary files (`cb=True`). If ``False``, use full TIGER/Line county
        boundaries.

    Returns
    -------
    geopandas.GeoDataFrame
        Boundaries in the requested CRS.
    """
    CRS.from_user_input(target_epsg)
    crs_suffix = "wgs84" if int(target_epsg) == 4326 else f"epsg{int(target_epsg)}"
    boundary_suffix = "cb" if cartographic else "tiger"
    study_area_boundary_geo_path = f"{work_dir}/{area_name}_{geo_level}_{year}_{boundary_suffix}_{crs_suffix}.geojson"
    if os.path.exists(study_area_boundary_geo_path):
        cached_geo = gpd.read_file(study_area_boundary_geo_path)
        return _project_if_needed(cached_geo, target_epsg)
    else:
        os.makedirs(os.path.dirname(study_area_boundary_geo_path), exist_ok=True)
        from pygris import counties, block_groups

        if geo_level == 'county':
            geo_data = counties(state=state_fips_code, year=year, cb=cartographic, cache=True)
        elif geo_level == 'cbg':
            # Define fips code for selected counties
            geo_data = block_groups(state=state_fips_code, year=year, cb=True, cache=True)
        elif geo_level == 'taz':
            geo_data = collect_taz_boundaries(state_fips_code, year, os.path.dirname(study_area_boundary_geo_path))
        elif geo_level == "tract":
            geo_data = collect_tract_boundaries(state_fips_code, county_fips_codes, year)
        else:
            raise ValueError("Unsupported geographic level. Choose 'counties' or 'cbgs'.")

        countyfp_columns = [col for col in geo_data.columns if col.startswith('COUNTYFP')]
        mask = geo_data[countyfp_columns].apply(lambda x: x.isin(county_fips_codes)).any(axis=1)
        selected_geo = geo_data[mask]

        selected_geo_projected = _project_if_needed(selected_geo, target_epsg)
        selected_geo_projected.to_file(study_area_boundary_geo_path, driver="GeoJSON")
        return selected_geo_projected

def collect_taz_boundaries(state_fips_code, year, output_dir):
    """Download and read TAZ boundaries from a TIGER/Line shapefile ZIP.

    Parameters
    ----------
    state_fips_code : str
        FIPS code for the state.
    year : int
        Reference year for boundaries.
    output_dir : str
        Directory to cache the downloaded ZIP file.

    Returns
    -------
    geopandas.GeoDataFrame
        TAZ boundary geometries.
    """
    import shutil
    import tempfile
    from zipfile import ZipFile

    state_geo_zip = output_dir + f"/tl_{year}_{state_fips_code}_taz10.zip"
    if not os.path.exists(state_geo_zip):
        state_geo_zip = download_taz_shapefile(state_fips_code, year, output_dir)

    # Extract the shapefile from the ZIP archive
    with ZipFile(state_geo_zip, 'r') as zip_ref:
        temp_dir = tempfile.mkdtemp()
        try:
            zip_ref.extractall(temp_dir)

            # Find the .shp file in the extracted files
            shapefile_name = [f for f in os.listdir(temp_dir) if f.endswith('.shp')][0]
            shapefile_path = os.path.join(temp_dir, shapefile_name)

            # Read the shapefile into a GeoDataFrame
            gdf = gpd.read_file(shapefile_path)
        finally:
            shutil.rmtree(temp_dir)

    return gdf

def download_taz_shapefile(state_fips_code, year, output_dir):
    """Download TAZ shapefiles for a given state-level FIPS code.

    Parameters
    ----------
    state_fips_code : str
        State-level FIPS code.
    year : int
        Reference year for the shapefile.
    output_dir : str
        Directory to save the downloaded ZIP file.

    Returns
    -------
    str
        Path to the downloaded ZIP file.
    """
    import requests
    # Ensure the FIPS code is a string, padded to 2 characters
    fips_code_str = str(state_fips_code).zfill(2)

    # Construct the download URL
    base_url = f"https://www2.census.gov/geo/tiger/TIGER2010/TAZ/2010/"
    filename = f"tl_{year}_{fips_code_str}_taz10.zip"
    download_url = base_url + filename

    # Make the output directory if it doesn't exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Full path for saving the file
    output_path = os.path.join(output_dir, filename)

    # Start the download
    logger.info("Downloading TAZ shapefile for FIPS code %s from %s", state_fips_code, download_url)
    try:
        response = requests.get(download_url)
        response.raise_for_status()  # This will check for errors

        # Write the content of the response to a ZIP file
        with open(output_path, 'wb') as file:
            file.write(response.content)

        logger.info("File saved to %s", output_path)

    except requests.RequestException as e:
        logger.error("Error downloading the file: %s", e)

    return output_path

def filter_boundaries_by_density(geo_data, pop_data, utm_epsg, geo_level, min_density_per_km2, density_geo_file):
    """
    Collect census boundaries, calculate population density, and filter by density threshold.

    Parameters
    ----------
    geo_data: GeoDataFrame
        Boundaries
    pop_data: Dataframe
        census
    utm_epsg: int
        EPSG code for the projected coordinate system
    geo_level: str
        Geographic level ('tract', 'cbg', etc.)
    min_density_per_km2: float, optional
        Minimum population density threshold (people per km²)
    density_geo_file: str
        Path to the CSV file containing density based geometry

    Returns
    -------
    geopandas.GeoDataFrame
        Filtered geographic boundaries based on density threshold
    """
    # Check if filtered boundaries already exist
    if os.path.exists(density_geo_file):
        logger.info("Loading existing %s boundaries...", geo_level)
        filtered_geo = gpd.read_file(density_geo_file)
        logger.info("Loaded %d %ss from existing file", len(filtered_geo), geo_level)
        return filtered_geo

    # If not, we need to collect and process the data
    logger.info("Processing %s boundaries for density analysis...", geo_level)

    # Ensure GEOID column has consistent type for merging
    geo_data['GEOID'] = geo_data['GEOID'].astype(str)

    # Merge boundaries with population data
    geo_with_pop = geo_data.merge(pop_data, on='GEOID')

    # Calculate area and density
    geo_with_pop['area_sqkm'] = (
            geo_with_pop.to_crs(epsg=utm_epsg)
            .geometry.area / 1000000  # Convert m² to km²
    )
    geo_with_pop['density_per_km2'] = geo_with_pop['population'] / geo_with_pop['area_sqkm']

    # Calculate percentile ranks
    geo_with_pop['density_percentile'] = (
            geo_with_pop['density_per_km2'].rank(pct=True) * 100
    ).round(1)

    # Log density analysis summary
    stats = geo_with_pop['density_per_km2'].describe()
    logger.info("Population Density Analysis:")
    logger.info("==========================")
    logger.info("%s Density Summary (people/km²):", geo_level.capitalize())
    logger.info("--------------------------------")
    logger.info("Mean density:     %s", format(stats['mean'], ",.1f"))
    logger.info("Median density:   %s", format(stats['50%'], ",.1f"))
    logger.info("Standard deviation:  %s", format(stats['std'], ",.1f"))
    logger.info("Minimum density:  %s", format(stats['min'], ",.1f"))
    logger.info("Maximum density:  %s", format(stats['max'], ",.1f"))
    logger.info("Density Distribution Quartiles:")
    logger.info("----------------------------")
    for q in [0.25, 0.5, 0.75]:
        logger.info(
            "%sth percentile: %s",
            int(q * 100),
            format(geo_with_pop['density_per_km2'].quantile(q), ",.1f"),
        )

    # Filter by minimum density if specified
    if min_density_per_km2 > 0:
        logger.info(
            "Filtering %ss by minimum density: %s people/km²",
            geo_level,
            format(min_density_per_km2, ",.1f"),
        )

        # Apply density filter
        filtered_geo = geo_with_pop[geo_with_pop["density_per_km2"] >= min_density_per_km2]

        # Log selection results
        logger.info("Selection Results:")
        logger.info("----------------")
        logger.info("Selected %d out of %d %ss", len(filtered_geo), len(geo_with_pop), geo_level)
        logger.info(
            "Total population in selected %ss: %s",
            geo_level,
            format(filtered_geo['population'].sum(), ","),
        )

        # Calculate percentage of total population
        total_population = geo_with_pop['population'].sum()
        if total_population > 0:
            population_percentage = (filtered_geo['population'].sum() / total_population * 100)
            logger.info("Percentage of total population: %.1f%%", population_percentage)
        else:
            logger.warning("Total population is zero, cannot calculate percentage")
    else:
        # If no density filter, use all areas
        filtered_geo = geo_with_pop
        logger.info("Using all %d %ss (no density filter applied)", len(filtered_geo), geo_level)

    # Ensure output is in WGS84 for consistency
    filtered_geo = filtered_geo.to_crs(epsg=4326)

    # Save to file for future use
    os.makedirs(os.path.dirname(density_geo_file), exist_ok=True)
    filtered_geo.to_file(density_geo_file, driver="GeoJSON")
    logger.info("Saved filtered %s boundaries to %s", geo_level, density_geo_file)

    return filtered_geo
