rest_api.py

"""
REST API layer for the Farm Layout Model.

Why this exists:
    The underlying engine (`design_api.process_farm_design`) only accepts a
    GeoJSON FeatureCollection string. Callers prefer a flatter, intent-driven
    schema:

        { farm, plots[], water_sources[] }

    This module owns the translation between that caller schema and the
    internal GeoJSON representation, and exposes a single one-shot endpoint
    ``POST /api/v1/design`` (plus a tiny health probe).

The Gradio UI is unaffected — `app.py` mounts this router alongside the
existing Gradio Blocks app.
"""
from __future__ import annotations

import json
from enum import Enum
from typing import Any, Dict, List, Optional

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field
from shapely.geometry import MultiPolygon, Polygon
from shapely.ops import unary_union

from design_api import process_farm_design
from unit_converter import m2_to_area, area_unit_label, supported_area_units, UnitError


# ──────────────────────────────────────────────────────────────────────────────
# Enums
# ──────────────────────────────────────────────────────────────────────────────


class DesignType(str, Enum):
    """Irrigation design type: centralized (valves at pump) or distributed (valves at zones)."""
    CENTRALIZED = "centralized"
    DISTRIBUTED = "distributed"


# ──────────────────────────────────────────────────────────────────────────────
# Request models — mirror the caller's schema 1:1 so docs/422s read naturally.
# ──────────────────────────────────────────────────────────────────────────────


class LatLon(BaseModel):
    latitude: float
    longitude: float


class FarmLocation(BaseModel):
    address: Optional[str] = None
    latitude: float
    longitude: float


class FarmSize(BaseModel):
    value: float
    unit: str


class Farm(BaseModel):
    name: str
    location: FarmLocation
    size: Optional[FarmSize] = None
    crop_name: str


class Plot(BaseModel):
    plot_id: str
    name: Optional[str] = None
    area_sq_m: Optional[float] = None
    centroid: Optional[LatLon] = None
    boundaries: List[LatLon] = Field(..., min_length=3)


class WaterSource(BaseModel):
    water_source_id: str
    type: Optional[str] = None
    name: Optional[str] = None
    location: LatLon
    plot_id: Optional[str] = None


class DesignRequest(BaseModel):
    farm: Farm
    plots: List[Plot] = Field(..., min_length=1)
    water_sources: List[WaterSource] = Field(..., min_length=1)
    # Engine-tuning knobs are optional with sensible defaults so callers can
    # ignore them entirely.
    pump_hp: float = 5.0
    design_type: Optional[DesignType] = None  # If None, will default based on farm size
    headland_buffer_m: float = 1.0
    override_lateral_spacing_m: Optional[float] = None
    area_unit: str = Field(
        default="ha",
        description="Unit for area values in the response. "
        "Accepts: ha, acres, m2, bigha, guntha, cent, and 20+ others. "
        "See GET /rest/v1/units for the full list.",
    )


# ──────────────────────────────────────────────────────────────────────────────
# Caller-schema  →  GeoJSON FeatureCollection
# ──────────────────────────────────────────────────────────────────────────────


def _ring_lonlat(boundary: List[LatLon]) -> List[List[float]]:
    """Convert {lat, lon} points to a closed [lon, lat] ring (RFC 7946 order)."""
    ring = [[p.longitude, p.latitude] for p in boundary]
    if ring[0] != ring[-1]:
        ring.append(ring[0])
    return ring


def _farm_boundary_polygon(plots: List[Plot]) -> Polygon:
    """
    Build a single Polygon that bounds every plot.

    `geojson_io.extract_farm_boundary` only accepts a Polygon. If the plots
    don't overlap, the union is a MultiPolygon — fall back to its convex hull
    so we still hand the engine something sane.
    """
    polys = [Polygon(_ring_lonlat(p.boundaries)) for p in plots]
    merged = unary_union(polys)
    if isinstance(merged, MultiPolygon):
        merged = merged.convex_hull
    if not isinstance(merged, Polygon):
        # Degenerate input (e.g. all collinear points) — fail fast with a
        # readable error instead of letting shapely surface a cryptic one.
        raise ValueError("Could not derive a single farm boundary polygon from plots")
    return merged


def _polygon_to_coords(poly: Polygon) -> List[List[List[float]]]:
    """Shapely Polygon → GeoJSON Polygon coordinates (single outer ring)."""
    return [[list(xy) for xy in poly.exterior.coords]]


def to_geojson_feature_collection(req: DesignRequest) -> str:
    """
    Translate a `DesignRequest` into the GeoJSON FeatureCollection shape that
    `process_farm_design` understands. Returns the serialized JSON string the
    engine expects.
    """
    crop = (req.farm.crop_name or "generic").strip().lower() or "generic"
    farm_boundary = _farm_boundary_polygon(req.plots)
    
    # Resolve design_type: explicit override or derive from farm size
    design_type = req.design_type
    if design_type is None:
        # Default: centralized for < 1ha, distributed for >= 1ha
        farm_area_ha = farm_boundary.area / 10000
        design_type = DesignType.DISTRIBUTED if farm_area_ha >= 1.0 else DesignType.CENTRALIZED

    features: List[Dict[str, Any]] = []

    # 1) Farm boundary (single Polygon — required by extract_farm_boundary).
    features.append({
        "type": "Feature",
        "properties": {
            "type": "farm_boundary",
            "name": req.farm.name,
        },
        "geometry": {
            "type": "Polygon",
            "coordinates": _polygon_to_coords(farm_boundary),
        },
    })

    # 2) Pump — engine consumes only the first source today; extras carried
    #    through as metadata-only Points so they round-trip in the response.
    primary_ws = req.water_sources[0]
    features.append({
        "type": "Feature",
        "properties": {
            "type": "pump",
            "name": primary_ws.name or primary_ws.water_source_id,
            "pump_hp": float(req.pump_hp),
            "horsepower": float(req.pump_hp),
            "water_source_id": primary_ws.water_source_id,
            "source_type": primary_ws.type,
        },
        "geometry": {
            "type": "Point",
            "coordinates": [primary_ws.location.longitude, primary_ws.location.latitude],
        },
    })
    for extra in req.water_sources[1:]:
        features.append({
            "type": "Feature",
            "properties": {
                "type": "water_source_extra",
                "name": extra.name or extra.water_source_id,
                "water_source_id": extra.water_source_id,
                "source_type": extra.type,
                "plot_id": extra.plot_id,
            },
            "geometry": {
                "type": "Point",
                "coordinates": [extra.location.longitude, extra.location.latitude],
            },
        })

    # 3) Crop zones — one per plot, all inheriting the farm's crop_name.
    for plot in req.plots:
        features.append({
            "type": "Feature",
            "properties": {
                "type": "crop_zone",
                "crop": crop,
                "name": plot.name or plot.plot_id,
                "plot_id": plot.plot_id,
            },
            "geometry": {
                "type": "Polygon",
                "coordinates": [_ring_lonlat(plot.boundaries)],
            },
        })

    fc = {
        "type": "FeatureCollection",
        "properties": {
            "farm_name": req.farm.name,
            "farm_address": req.farm.location.address,
            "pump_hp": float(req.pump_hp),
            "design_type": design_type.value,
            "headland_buffer_m": float(req.headland_buffer_m),
            "override_lateral_spacing_m": req.override_lateral_spacing_m,
        },
        "features": features,
    }
    return json.dumps(fc)


# ──────────────────────────────────────────────────────────────────────────────
# Router
# ──────────────────────────────────────────────────────────────────────────────


def _convert_area_fields(data: Any, unit: str) -> Any:
    """Recursively convert *_ha and area_m2 fields to the requested unit.

    Walks dicts and lists. For every key ending in ``_ha`` it adds a
    sibling ``_<unit>`` key with the converted value.  ``area_m2`` fields
    get a sibling ``area_<unit>`` key.
    """
    if isinstance(data, dict):
        out: Dict[str, Any] = {}
        for k, v in data.items():
            out[k] = _convert_area_fields(v, unit)
            # Convert hectare fields → requested unit
            if k.endswith("_ha") and isinstance(v, (int, float)):
                m2 = v * 10_000
                label = area_unit_label(unit)
                new_key = k.replace("_ha", f"_{unit}")
                out[new_key] = round(m2_to_area(m2, unit), 4)
                out[k.replace("_ha", "_unit")] = label
            # Convert m² fields → requested unit
            if k == "area_m2" and isinstance(v, (int, float)):
                out[f"area_{unit}"] = round(m2_to_area(v, unit), 4)
        return out
    if isinstance(data, list):
        return [_convert_area_fields(item, unit) for item in data]
    return data


def build_router() -> APIRouter:
    router = APIRouter(prefix="/rest/v1", tags=["design"])

    @router.get("/health")
    def health() -> Dict[str, str]:
        return {"status": "ok", "schema_version": "v1"}

    @router.get("/units")
    def units() -> Dict[str, Any]:
        """List all supported area units for the area_unit parameter."""
        return {
            "area_units": {
                k: label for k, label in supported_area_units().items()
            },
        }

    @router.post("/design")
    def design(req: DesignRequest) -> Dict[str, Any]:
        # Validate area_unit early so callers get a clear error.
        area_unit = req.area_unit.strip().lower() or "ha"
        try:
            area_unit_label(area_unit)  # raises UnitError if invalid
        except UnitError:
            raise HTTPException(
                status_code=422,
                detail=f"Unknown area_unit '{req.area_unit}'. "
                f"Use GET /rest/v1/units to see supported values.",
            )

        # Translate caller schema → GeoJSON the engine understands.
        try:
            geojson_str = to_geojson_feature_collection(req)
        except ValueError as e:
            raise HTTPException(status_code=422, detail=str(e))

        # Run the engine. It returns a GeoJSON FeatureCollection; on a
        # business-rule failure it tags the result with type=farm_design_error.
        try:
            result = process_farm_design(geojson_str)
        except Exception as e:  # noqa: BLE001 — surface engine errors verbatim
            raise HTTPException(status_code=500, detail=f"design pipeline failed: {e}")

        props = result.get("properties", {}) or {}
        if props.get("type") == "farm_design_error":
            err = props.get("error", {}) or {}
            raise HTTPException(
                status_code=422,
                detail={
                    "code": err.get("code", "design_error"),
                    "message": err.get("message", "design pipeline returned an error"),
                    "geojson": result,
                },
            )

        warnings: List[str] = []
        if len(req.water_sources) > 1:
            warnings.append(
                "Only the first water_source was used as the pump; "
                "multi-pump designs are not yet supported."
            )

        # Convert area fields if caller requested a non-default unit.
        design_summary = props.get("design_summary", {})
        bom = props.get("bom", {})
        zone_details = props.get("zone_details", [])
        geojson_out = result

        if area_unit != "ha":
            design_summary = _convert_area_fields(design_summary, area_unit)
            zone_details = _convert_area_fields(zone_details, area_unit)
            geojson_out = _convert_area_fields(result, area_unit)

        return {
            "farm_name": req.farm.name,
            "area_unit": area_unit,
            "design_summary": design_summary,
            "bom": bom,
            "zone_details": zone_details,
            "warnings": warnings,
            "geojson": geojson_out,
        }

    return router