RadioPropagationApi/api/app/services/buildings.py

from __future__ import annotations

import json
from typing import Any

from sqlalchemy import text
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import Session

from app.core.geo import PathPoint
from app.models.buildings import BuildingsQueryRequest, BuildingsQueryResponse


def _feature_from_value(value: Any) -> dict[str, Any]:
    if isinstance(value, str):
        return json.loads(value)
    return dict(value)


def _geometry_from_geojson(geojson: dict[str, Any]) -> dict[str, Any]:
    if geojson.get("type") == "Feature":
        geometry = geojson.get("geometry")
    else:
        geometry = geojson

    if not isinstance(geometry, dict) or geometry.get("type") not in {
        "LineString",
        "MultiLineString",
        "Polygon",
        "MultiPolygon",
    }:
        raise ValueError("path must be a GeoJSON geometry or Feature")
    return geometry


def query_buildings(
    request: BuildingsQueryRequest,
    db: Session,
    limit: int = 5000,
) -> BuildingsQueryResponse:
    try:
        if request.bbox is not None:
            rows = _query_by_bbox(db, request.bbox, limit)
        else:
            rows = _query_by_path(db, request.path or {}, request.buffer_m, limit)
    except SQLAlchemyError as exc:
        raise NotImplementedError(
            "PostGIS buildings table is not ready. Run scripts/load_buildings.sh first."
        ) from exc

    return BuildingsQueryResponse(features=[_feature_from_value(row["feature"]) for row in rows])


def building_heights_along(points: list[PathPoint], db: Session) -> list[float]:
    if not points:
        return []

    points_json = json.dumps(
        [{"i": index, "lon": point.lon, "lat": point.lat} for index, point in enumerate(points)]
    )
    statement = text(
        """
        WITH points AS (
          SELECT *
          FROM jsonb_to_recordset(CAST(:points_json AS jsonb))
            AS p(i integer, lon double precision, lat double precision)
        )
        SELECT p.i, COALESCE(MAX(b.height_m), 0) AS height_m
        FROM points p
        LEFT JOIN buildings b
          ON ST_Intersects(b.geom, ST_SetSRID(ST_MakePoint(p.lon, p.lat), 4326))
        GROUP BY p.i
        ORDER BY p.i
        """
    )
    rows = db.execute(statement, {"points_json": points_json}).mappings().all()
    heights = [0.0] * len(points)
    for row in rows:
        heights[int(row["i"])] = float(row["height_m"] or 0.0)
    return heights


def _query_by_bbox(
    db: Session,
    bbox: tuple[float, float, float, float],
    limit: int,
) -> list[dict[str, Any]]:
    minlon, minlat, maxlon, maxlat = bbox
    statement = text(
        """
        SELECT jsonb_build_object(
          'type', 'Feature',
          'id', osm_id,
          'geometry', ST_AsGeoJSON(geom)::jsonb,
          'properties', jsonb_build_object(
            'height_m', height_m,
            'levels', levels,
            'building_type', building_type,
            'source', source
          )
        ) AS feature
        FROM buildings
        WHERE geom && ST_MakeEnvelope(:minlon, :minlat, :maxlon, :maxlat, 4326)
          AND ST_Intersects(
            geom,
            ST_MakeEnvelope(:minlon, :minlat, :maxlon, :maxlat, 4326)
          )
        LIMIT :limit
        """
    )
    return list(
        db.execute(
            statement,
            {
                "minlon": minlon,
                "minlat": minlat,
                "maxlon": maxlon,
                "maxlat": maxlat,
                "limit": limit,
            },
        )
        .mappings()
        .all()
    )


def _query_by_path(
    db: Session,
    path_geojson: dict[str, Any],
    buffer_m: float,
    limit: int,
) -> list[dict[str, Any]]:
    geometry_json = json.dumps(_geometry_from_geojson(path_geojson))
    statement = text(
        """
        WITH query_geom AS (
          SELECT ST_Transform(
            ST_Buffer(
              ST_Transform(ST_SetSRID(ST_GeomFromGeoJSON(:geometry_json), 4326), 3857),
              :buffer_m
            ),
            4326
          ) AS geom
        )
        SELECT jsonb_build_object(
          'type', 'Feature',
          'id', b.osm_id,
          'geometry', ST_AsGeoJSON(b.geom)::jsonb,
          'properties', jsonb_build_object(
            'height_m', b.height_m,
            'levels', b.levels,
            'building_type', b.building_type,
            'source', b.source
          )
        ) AS feature
        FROM buildings b, query_geom q
        WHERE ST_Intersects(b.geom, q.geom)
        LIMIT :limit
        """
    )
    return list(
        db.execute(
            statement,
            {"geometry_json": geometry_json, "buffer_m": buffer_m, "limit": limit},
        )
        .mappings()
        .all()
    )