import geopandas as gpd
import os

# -----------------------------
# Settings
# -----------------------------
output_dir = "/mnt/storage/docker-data/pupmap/www/data/sources/"
os.makedirs(output_dir, exist_ok=True)

# Europe + Nordics country codes (for cities)
europe_nordics = [
    'AL','AD','AT','BE','BA','BG','HR','CY','CZ','DK','EE','FI','FR','DE',
    'GR','HU','IS','IE','IT','LV','LI','LT','LU','MT','MD','MC','ME','NL',
    'MK','NO','PL','PT','RO','RU','SM','RS','SK','SI','ES','SE','CH','UA','GB'
]

# -----------------------------
# Process Cities
# -----------------------------
print("Processing cities...")
cities = gpd.read_file("ne_10m_populated_places/ne_10m_populated_places.shp")
cities_filtered = cities[cities['ADM0_A3'].isin(europe_nordics)].copy()

# Extract coordinates from geometry
cities_filtered['lon'] = cities_filtered.geometry.x
cities_filtered['lat'] = cities_filtered.geometry.y

# Keep needed columns + geometry
cities_filtered = cities_filtered[['NAME','ADM0NAME','lat','lon','POP_MAX','geometry']]
cities_filtered = cities_filtered.rename(columns={
    'NAME':'name',
    'ADM0NAME':'country',
    'POP_MAX':'population'
})

# Drop any duplicate columns
cities_filtered = cities_filtered.loc[:, ~cities_filtered.columns.duplicated()]

cities_geojson_path = os.path.join(output_dir, "world_cities_europe.geojson")
cities_filtered.to_file(cities_geojson_path, driver='GeoJSON')
print(f"Saved {cities_geojson_path}")

# -----------------------------
# Process Roads
# -----------------------------
print("Processing roads...")
roads = gpd.read_file("ne_10m_roads/ne_10m_roads.shp")

# Filter by Europe/Nordics bounding box
minx, miny, maxx, maxy = -25.0, 34.0, 40.0, 72.0
roads_filtered = roads.cx[minx:maxx, miny:maxy].copy()

# Remove duplicate columns
roads_filtered = roads_filtered.loc[:, ~roads_filtered.columns.duplicated()]

# Pick columns: use FULLNAME/FCLASS if they exist
if 'FULLNAME' in roads_filtered.columns and 'FCLASS' in roads_filtered.columns:
    roads_filtered = roads_filtered[['FULLNAME','FCLASS','geometry']]
    roads_filtered = roads_filtered.rename(columns={'FULLNAME':'name','FCLASS':'type'})
else:
    # fallback: take first two non-geometry columns
    geom_col = 'geometry'
    other_cols = [c for c in roads_filtered.columns if c != geom_col]
    roads_filtered = roads_filtered[other_cols[:2] + [geom_col]]
    roads_filtered = roads_filtered.rename(columns={other_cols[0]:'name', other_cols[1]:'type'})

roads_geojson_path = os.path.join(output_dir, "roads_europe.geojson")
roads_filtered.to_file(roads_geojson_path, driver='GeoJSON')
print(f"Saved {roads_geojson_path}")

print("✅ All GeoJSON files are ready for MapLibre!")