rest_api.py

"""
rest_api.py

REST / OpenAPI surface for Hydraulic Solver Teaching Mode.

Run:
    CLIENT_API_KEY=your-secret uvicorn rest_api:app --host 0.0.0.0 --port 8000

OpenAPI:
    /openapi.json
Docs:
    /docs
"""

from __future__ import annotations

import os
from typing import List, Optional

from fastapi import Depends, FastAPI, Header, HTTPException
from pydantic import BaseModel, Field

import hydraulic_core as hc


API_KEY = os.getenv("CLIENT_API_KEY", "")

app = FastAPI(
    title="Hydraulic Solver Teaching Mode API",
    version="1.0.0",
    description="REST/OpenAPI surface exposing the same deterministic tools as the MCP server.",
)


def require_auth(authorization: Optional[str] = Header(default=None)) -> None:
    if not API_KEY:
        return
    expected = f"Bearer {API_KEY}"
    if authorization != expected:
        raise HTTPException(status_code=401, detail="Missing or invalid bearer token.")


class LoadNetworkRequest(BaseModel):
    inp_text: str
    title: str = "uploaded_network"


class SinglePipeRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    length: float
    diameter: float
    c_hw: float = 120.0
    target_headloss: float
    initial_flow: float
    max_iter: int = 25
    tolerance: float = 0.0001


class HardyLoopRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    flows: List[float]
    lengths: List[float]
    diameters: List[float]
    c_values: List[float]
    max_iter: int = 25
    tolerance: float = 0.00001


class TwoLoopRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    flows: List[float] = Field(..., description="Seven initial pipe flows P1..P7.")
    common_length: float
    common_diameter: float
    c_hw: float = 120.0
    max_iter: int = 25
    tolerance: float = 0.00001


class ThreeReservoirRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    reservoir_heads: List[float]
    demand: float
    initial_head: float
    lengths: List[float]
    diameters: List[float]
    c_values: List[float]
    max_iter: int = 25
    tolerance: float = 0.0001


class PDDRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    required_demand: float
    available_pressure: float
    minimum_pressure: float
    required_pressure: float
    exponent: float = 0.5


class TankEPSRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    diameter: float
    initial_level: float
    min_level: float
    max_level: float
    timestep_hr: float = 1.0
    inflows: List[float]
    outflows: List[float]


class ValveRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    valve_type: str = "PRV"
    upstream_head: float
    setting: float
    flow: float
    diameter: float
    minor_loss_k: float = 0.0


class PumpRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    shutoff_head: float
    design_flow: float
    static_head: float
    pump_curve_k: float
    system_curve_k: float


class PressureZoneRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    source_head: float
    prv_setting: float
    node_elevations: List[float]
    demand_multiplier: float = 1.0
    min_pressure: float = 14.0
    max_pressure: float = 56.0


class LeakageRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    average_pressure: float
    authorized_demand: float
    leakage_coefficient: float
    pressure_exponent: float = 1.0


class WaterAgeRequest(BaseModel):
    unit_system: str = "MKS / L/s"
    pipe_volume: float
    tank_volume: float
    demand: float
    dead_end_factor: float = 2.0


class ChlorineRequest(BaseModel):
    initial_chlorine_mg_l: float
    bulk_decay_per_day: float
    travel_time_hours: float
    wall_decay_factor: float = 0.0


@app.get("/health")
def health():
    return {"ok": True, "service": "hydraulic-solver-teaching-mode-api"}


@app.post("/sessions", dependencies=[Depends(require_auth)])
def create_session(req: LoadNetworkRequest):
    return hc.load_network(req.inp_text, req.title)


@app.get("/sessions/{session_id}", dependencies=[Depends(require_auth)])
def get_session_summary(session_id: str):
    return hc.network_summary(session_id)


@app.delete("/sessions/{session_id}", dependencies=[Depends(require_auth)])
def delete_session(session_id: str):
    return hc.close_session(session_id)


@app.post("/tools/solve_single_pipe", dependencies=[Depends(require_auth)])
def solve_single_pipe(req: SinglePipeRequest):
    return hc.solve_single_pipe(**req.model_dump())


@app.post("/tools/solve_hardy_cross_loop", dependencies=[Depends(require_auth)])
def solve_hardy_cross_loop(req: HardyLoopRequest):
    return hc.solve_hardy_cross_loop(**req.model_dump())


@app.post("/tools/solve_two_loop_hardy_cross", dependencies=[Depends(require_auth)])
def solve_two_loop_hardy_cross(req: TwoLoopRequest):
    return hc.solve_two_loop_hardy_cross(**req.model_dump())


@app.post("/tools/solve_three_reservoir", dependencies=[Depends(require_auth)])
def solve_three_reservoir(req: ThreeReservoirRequest):
    return hc.solve_three_reservoir(**req.model_dump())


@app.post("/tools/solve_pdd_demand", dependencies=[Depends(require_auth)])
def solve_pdd_demand(req: PDDRequest):
    return hc.solve_pdd_demand(**req.model_dump())


@app.post("/tools/simulate_tank_eps", dependencies=[Depends(require_auth)])
def simulate_tank_eps(req: TankEPSRequest):
    return hc.simulate_tank_eps(**req.model_dump())


@app.post("/tools/evaluate_valve_behavior", dependencies=[Depends(require_auth)])
def evaluate_valve_behavior(req: ValveRequest):
    return hc.evaluate_valve_behavior(**req.model_dump())


@app.post("/tools/solve_pump_operating_point", dependencies=[Depends(require_auth)])
def solve_pump_operating_point(req: PumpRequest):
    return hc.solve_pump_operating_point(**req.model_dump())


@app.post("/tools/pressure_zone_analysis", dependencies=[Depends(require_auth)])
def pressure_zone_analysis(req: PressureZoneRequest):
    return hc.pressure_zone_analysis(**req.model_dump())


@app.post("/tools/leakage_nrw_analysis", dependencies=[Depends(require_auth)])
def leakage_nrw_analysis(req: LeakageRequest):
    return hc.leakage_nrw_analysis(**req.model_dump())


@app.post("/tools/water_age_analysis", dependencies=[Depends(require_auth)])
def water_age_analysis(req: WaterAgeRequest):
    return hc.water_age_analysis(**req.model_dump())


@app.post("/tools/chlorine_decay_analysis", dependencies=[Depends(require_auth)])
def chlorine_decay_analysis(req: ChlorineRequest):
    return hc.chlorine_decay_analysis(**req.model_dump())


@app.get("/tools/generate_epanet_validation_inp", dependencies=[Depends(require_auth)])
def generate_epanet_validation_inp(case: str = "three_reservoir_mks"):
    return hc.generate_epanet_validation_inp(case)