import osmnx as ox
import shapely
import random
import uuid
import networkx as nx
from matplotlib import cm

def generate_osm_city_data(lat, lon, dist=400, scale=1.0):
    print(f"🏙️ Fetching OSM buildings and network at ({lat}, {lon})")

    scale_factor = scale
    status_options = ["OK", "Warning", "Critical", "Offline"]

    # ————— STREET NETWORK —————
    G = ox.graph_from_point((lat, lon), dist=dist, network_type='drive').to_undirected()
    degree = dict(G.degree())
    max_degree = max(degree.values()) if degree else 1
    color_map = cm.get_cmap("plasma")

    # ————— BUILDINGS —————
    tags = {"building": True}
    gdf = ox.features_from_point((lat, lon), tags=tags, dist=dist)
    gdf = gdf[gdf.geometry.type.isin(["Polygon", "MultiPolygon"])].to_crs(epsg=3857)
    gdf["centroid"] = gdf.geometry.centroid

    raw_buildings = []
    for i, row in gdf.iterrows():
        centroid = row["centroid"]
        polygon = row["geometry"]
        building_id = f"BLD-{uuid.uuid4().hex[:6].upper()}"
        status = random.choice(status_options)

        try:
            height = float(row.get("height", None))
        except:
            height = float(row.get("building:levels", 3)) * 3.2 if row.get("building:levels") else 10.0

        try:
            node = ox.distance.nearest_nodes(G, X=centroid.x, Y=centroid.y)
            node_degree = degree.get(node, 0)
        except:
            node_degree = 0

        norm_value = node_degree / max_degree
        rgba = color_map(norm_value)
        hex_color = '#%02x%02x%02x' % tuple(int(c * 255) for c in rgba[:3])

        raw_buildings.append({
            "id": building_id,
            "raw_x": centroid.x,
            "raw_z": centroid.y,
            "width": polygon.bounds[2] - polygon.bounds[0],
            "depth": polygon.bounds[3] - polygon.bounds[1],
            "height": height,
            "color": hex_color,
            "status": status,
        })

    # ————— CENTER AND SCALE —————
    if raw_buildings:
        avg_x = sum(b['raw_x'] for b in raw_buildings) / len(raw_buildings)
        avg_z = sum(b['raw_z'] for b in raw_buildings) / len(raw_buildings)

        buildings = [{
            "id": b['id'],
            "position_x": (b['raw_x'] - avg_x) * scale_factor,
            "position_z": (b['raw_z'] - avg_z) * scale_factor,
            "width": b['width'] * scale_factor,
            "depth": b['depth'] * scale_factor,
            "height": b['height'] * scale_factor,
            "color": b['color'],
            "status": b['status'],
        } for b in raw_buildings]
    else:
        buildings = []

    return {"buildings": buildings}
