import geopandas as gpd
import rasterio
from rasterio.features import shapes
from shapely.geometry import shape
import zipfile
import os
import tempfile
import pandas as pd
import numpy as np

def analyze_exposed_populated_places(raster_path: str, pp_path: str):
    """
    Identify populated places exposed to hazard (intersecting non-zero raster values).
    
    Args:
        raster_path (str): Path to the hazard raster file
        pp_path (str): Path to the HOTOSM populated places shapefile zip
        
    Returns:
        tuple: (top_10_message, type_counts_message) or (None, error_message)
    """
    if not pp_path or not os.path.exists(pp_path):
        return None, "⚠️ No populated places data available for analysis."
    
    try:
        # Extract shapefile to temp directory
        temp_dir = tempfile.mkdtemp()
        with zipfile.ZipFile(pp_path, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # Find the .shp file
        shp_file = next((os.path.join(temp_dir, f) for f in os.listdir(temp_dir) if f.endswith('.shp')), None)
        if not shp_file:
            return None, "⚠️ No shapefile found in populated places data."
        
        # Load populated places
        pp_gdf = gpd.read_file(shp_file)
        
        # Check required columns
        if not {'name_en', 'population', 'place'}.issubset(pp_gdf.columns):
            return None, "⚠️ Populated places data missing required fields (name_en, population, place)."
        
        # Cleanup names
        pp_gdf['name_en'] = pp_gdf['name_en'].replace("None", np.nan)
        pp_gdf['name_en'] = pp_gdf['name_en'].fillna(pp_gdf.get('name', 'Unknown'))
        
        # Load raster and get non-zero areas as polygons
        with rasterio.open(raster_path) as src:
            raster_data = src.read(1)
            raster_crs = src.crs
            nodata = src.nodata if src.nodata is not None else -9999
            hazard_mask = (raster_data != nodata) & (raster_data != 0) & (raster_data > 0)
            hazard_geoms = [shape(geom) for geom, value in shapes(hazard_mask.astype('uint8'), transform=src.transform) if value == 1]
        
        if not hazard_geoms:
            return None, "⚠️ No hazard exposure detected in raster."
        
        hazard_gdf = gpd.GeoDataFrame({'geometry': hazard_geoms}, crs=raster_crs)
        if pp_gdf.crs != hazard_gdf.crs:
            pp_gdf = pp_gdf.to_crs(hazard_gdf.crs)
                
        # Find intersecting populated places
        exposed_pp = gpd.sjoin(pp_gdf, hazard_gdf, predicate='intersects', how='inner')

        if exposed_pp.empty:
            return None, "✅ No populated places exposed to this hazard."

        # Keep full exposed_pp for type counts (includes all places)
        full_exposed_pp = exposed_pp.copy()

        # Top 10 by population (only keep places with population)
        top_pp = full_exposed_pp.copy()
        top_pp['population'] = pd.to_numeric(top_pp['population'], errors='coerce')
        top_pp = top_pp.dropna(subset=['population'])
        top_10 = top_pp.nlargest(10, 'population')[['name_en', 'population', 'place']]

        top_10_lines = ["**📍 Top Exposed Populated Places by Population:**"]
        for idx, row in enumerate(top_10.itertuples(), 1):
            pop = int(row.population)
            place_type = row.place if row.place else "Unknown"
            top_10_lines.append(f"{idx}. **{row.name_en}** ({place_type}) - Population: {pop:,}")
        top_10_message = "\n".join(top_10_lines)

        # Type counts (include all exposed places, even missing population)
        type_counts = full_exposed_pp['place'].value_counts().to_dict()
        type_lines = ["\n\n**🏘️ Exposed Populated Places by Type:**", "", f"**Total:** {len(full_exposed_pp):,}"]
        for place_type, count in sorted(type_counts.items(), key=lambda x: x[1], reverse=True):
            cleaned_type = place_type.replace("_", " ").capitalize()
            type_lines.append(f"- {cleaned_type}: {count:,}")
        type_counts_message = "\n".join(type_lines)
        
        return top_10_message, type_counts_message
        
    except Exception as e:
        return None, f"⚠️ Error analyzing populated places: {str(e)}"


# Metadata for dynamic discovery
ANALYSIS_METADATA = {
    "name": "Exposed Populated Places",
    "function": analyze_exposed_populated_places,
    "required_files": ["raster_path", "pp_path"],
    "enabled": True
}