#!/usr/bin/env python3
"""
Extracts building features from an OpenStreetMap .pbf file.

This script complements extract_roads.py by extracting building polygons
with attributes like building type, height, and building:levels for heatmap analysis.
"""

import argparse
import sys
from pathlib import Path
from pyrosm import OSM
import geopandas as gpd

def extract_buildings(location: str):
    location = location.lower()
    base_dir = Path.cwd()
    location_dir = base_dir / "output" / location
    
    # Check if location directory and PBF file exist
    pbf_file = location_dir / f"{location}.osm.pbf"
    if not pbf_file.exists():
        # Look for any .pbf file in the directory
        pbf_files = list(location_dir.glob("*.pbf"))
        if pbf_files:
            pbf_file = pbf_files[0]
            print(f"📦 Using PBF file: {pbf_file.name}")
        else:
            print(f"❌ No PBF file found in {location_dir}")
            print(f"   Run extract_roads.py first to download the OSM data.")
            sys.exit(1)
    
    output_path = location_dir / f"{location}_buildings.gpkg"
    
    # ------------------------------
    # EXTRACT BUILDINGS
    # ------------------------------
    print(f"\n🏢 Extracting buildings from {pbf_file.name}...")
    try:
        osm = OSM(str(pbf_file))
        buildings = osm.get_data_by_custom_criteria(
            custom_filter={"building": True},
            filter_type="keep",
            keep_nodes=False,
            keep_relations=False,
            # Extract key building attributes for heatmap analysis
            extra_attributes=[
                "building",      # building type (residential, commercial, etc.)
                "height",        # building height in meters
                "building:levels", # number of floors/stories
                "addr:housenumber",
                "addr:street",
                "amenity",
                "shop",
                "office"
            ]
        )

        if buildings is None or buildings.empty:
            print("⚠️ No building data found in this OSM file.")
            print("   This location may not have detailed building data in OSM.")
            return

        print(f"✅ Extracted {len(buildings):,} buildings from {location.capitalize()} data.")
        print(f"📋 Columns available: {list(buildings.columns)}")

        # Show some stats about the extracted data
        if 'building' in buildings.columns:
            building_types = buildings['building'].value_counts().head(10)
            print(f"🏘️ Top building types:\n{building_types}")
        
        if 'height' in buildings.columns:
            height_stats = buildings['height'].describe()
            print(f"📏 Height stats:\n{height_stats}")

    except Exception as e:
        print(f"❌ Building extraction failed: {e}")
        sys.exit(1)

    # ------------------------------
    # REPROJECT TO WGS84
    # ------------------------------
    try:
        buildings = buildings.to_crs(epsg=4326)
        print("🌐 Reprojected to WGS84 (EPSG:4326).")
    except Exception as e:
        print(f"⚠️ Reprojection failed: {e}")

    # ------------------------------
    # SAVE BUILDINGS DATA
    # ------------------------------
    try:
        buildings.to_file(output_path, layer="buildings", driver="GPKG")
        print(f"💾 Saved extracted buildings to: {output_path}")
    except Exception as e:
        print(f"❌ Saving failed: {e}")
        sys.exit(1)

    # ------------------------------
    # SUMMARY
    # ------------------------------
    print("\n✅ Building extraction complete!")
    print(f"📁 Folder: {location_dir}")
    print(f"📥 Input: {pbf_file.name}")
    print(f"💾 Output: {output_path.name}")
    print(f"🏢 Buildings extracted: {len(buildings):,}")
    print("-" * 60)
    print("Next steps:")
    print("1. Run load_buildings.py to reproject buildings to meters")
    print("2. Run visualize_all_attribute_grid.py to generate building heatmaps")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Extract OSM buildings for a location using a local PBF file."
    )
    parser.add_argument("location", help="Location name (e.g., alabama, egypt, thailand).")
    args = parser.parse_args()

    extract_buildings(args.location)