server.py

"""
server.py

MCP surface for Hydraulic Solver Teaching Mode.

Local stdio:
    python server.py

Remote HTTP on HF Spaces:
    TRANSPORT=http PORT=7860 CLIENT_API_KEY=your-secret python server.py

Expected endpoints with recent MCP Python SDK:
    Streamable HTTP: /mcp
    SSE: usually /sse when TRANSPORT=sse

Note:
- This file exposes deterministic teaching/hydraulic tools from hydraulic_core.py.
- It is prepared for agentic clients such as Claude, Cursor, LM Studio, n8n,
  HuggingChat, Codex, LangChain/LangGraph, and AutoGen.
"""

from __future__ import annotations

import os
from typing import List

import uvicorn
from starlette.responses import HTMLResponse, JSONResponse
from starlette.requests import Request
from starlette.middleware.base import BaseHTTPMiddleware

from mcp.server.fastmcp import FastMCP

import hydraulic_core as hc


HOST = os.getenv("HOST", "0.0.0.0")
PORT = int(os.getenv("PORT", "7860"))
TRANSPORT = os.getenv("TRANSPORT", "stdio").lower()
CLIENT_API_KEY = os.getenv("CLIENT_API_KEY", "")

mcp = FastMCP("hydraulic-solver-teaching-mode", host=HOST, port=PORT)


@mcp.tool()
def load_network(inp_text: str, title: str = "uploaded_network") -> dict:
    """Load EPANET INP text into a temporary session and return a session_id."""
    return hc.load_network(inp_text, title)


@mcp.tool()
def upload_network(inp_text: str, title: str = "uploaded_network") -> dict:
    """Alias of load_network for platforms that prefer upload terminology."""
    return hc.load_network(inp_text, title)


@mcp.tool()
def close_session(session_id: str) -> dict:
    """Close a loaded network session."""
    return hc.close_session(session_id)


@mcp.tool()
def network_summary(session_id: str) -> dict:
    """Return parsed session summary: units and section counts."""
    return hc.network_summary(session_id)


@mcp.tool()
def solve_single_pipe(unit_system: str, length: float, diameter: float, c_hw: float,
                      target_headloss: float, initial_flow: float,
                      max_iter: int = 25, tolerance: float = 0.0001) -> dict:
    """Solve a single Hazen-Williams pipe flow/headloss problem."""
    return hc.solve_single_pipe(unit_system, length, diameter, c_hw, target_headloss, initial_flow, max_iter, tolerance)


@mcp.tool()
def solve_hardy_cross_loop(unit_system: str, flows: List[float], lengths: List[float],
                           diameters: List[float], c_values: List[float],
                           max_iter: int = 25, tolerance: float = 0.00001) -> dict:
    """Run a single-loop Hardy Cross teaching calculation."""
    return hc.solve_hardy_cross_loop(unit_system, flows, lengths, diameters, c_values, max_iter, tolerance)


@mcp.tool()
def solve_two_loop_hardy_cross(unit_system: str, flows: List[float], common_length: float,
                               common_diameter: float, c_hw: float = 120.0,
                               max_iter: int = 25, tolerance: float = 0.00001) -> dict:
    """Run a two-loop Hardy Cross teaching calculation with shared pipe P2."""
    return hc.solve_two_loop_hardy_cross(unit_system, flows, common_length, common_diameter, c_hw, max_iter, tolerance)


@mcp.tool()
def solve_three_reservoir(unit_system: str, 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) -> dict:
    """Solve the three-reservoir central junction head problem."""
    return hc.solve_three_reservoir(unit_system, reservoir_heads, demand, initial_head, lengths, diameters, c_values, max_iter, tolerance)


@mcp.tool()
def solve_pdd_demand(unit_system: str, required_demand: float, available_pressure: float,
                     minimum_pressure: float, required_pressure: float, exponent: float = 0.5) -> dict:
    """Calculate pressure-dependent demand satisfaction."""
    return hc.solve_pdd_demand(unit_system, required_demand, available_pressure, minimum_pressure, required_pressure, exponent)


@mcp.tool()
def simulate_tank_eps(unit_system: str, diameter: float, initial_level: float, min_level: float,
                      max_level: float, timestep_hr: float,
                      inflows: List[float], outflows: List[float]) -> dict:
    """Run an isolated tank EPS water-balance teaching simulation."""
    return hc.simulate_tank_eps(unit_system, diameter, initial_level, min_level, max_level, timestep_hr, inflows, outflows)


@mcp.tool()
def evaluate_valve_behavior(unit_system: str, valve_type: str, upstream_head: float,
                            setting: float, flow: float, diameter: float, minor_loss_k: float = 0.0) -> dict:
    """Evaluate isolated PRV/PSV/FCV/TCV/check-valve teaching behavior."""
    return hc.evaluate_valve_behavior(unit_system, valve_type, upstream_head, setting, flow, diameter, minor_loss_k)


@mcp.tool()
def solve_pump_operating_point(unit_system: str, shutoff_head: float, design_flow: float,
                               static_head: float, pump_curve_k: float, system_curve_k: float) -> dict:
    """Estimate a pump/system curve operating point."""
    return hc.solve_pump_operating_point(unit_system, shutoff_head, design_flow, static_head, pump_curve_k, system_curve_k)


@mcp.tool()
def pressure_zone_analysis(unit_system: str, 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) -> dict:
    """Analyze simple pressure-zone pressure categories under a PRV/source head."""
    return hc.pressure_zone_analysis(unit_system, source_head, prv_setting, node_elevations, demand_multiplier, min_pressure, max_pressure)


@mcp.tool()
def leakage_nrw_analysis(unit_system: str, average_pressure: float, authorized_demand: float,
                         leakage_coefficient: float, pressure_exponent: float = 1.0) -> dict:
    """Estimate pressure-dependent leakage and NRW percentage."""
    return hc.leakage_nrw_analysis(unit_system, average_pressure, authorized_demand, leakage_coefficient, pressure_exponent)


@mcp.tool()
def water_age_analysis(unit_system: str, pipe_volume: float, tank_volume: float,
                       demand: float, dead_end_factor: float = 2.0) -> dict:
    """Estimate simple water-age/residence-time teaching values."""
    return hc.water_age_analysis(unit_system, pipe_volume, tank_volume, demand, dead_end_factor)


@mcp.tool()
def chlorine_decay_analysis(initial_chlorine_mg_l: float, bulk_decay_per_day: float,
                            travel_time_hours: float, wall_decay_factor: float = 0.0) -> dict:
    """Estimate first-order chlorine decay for teaching."""
    return hc.chlorine_decay_analysis(initial_chlorine_mg_l, bulk_decay_per_day, travel_time_hours, wall_decay_factor)


@mcp.tool()
def generate_epanet_validation_inp(case: str = "three_reservoir_mks") -> dict:
    """Generate a small EPANET INP validation model."""
    return hc.generate_epanet_validation_inp(case)


class BearerAuthMiddleware(BaseHTTPMiddleware):
    """Simple bearer-token guard for HTTP MCP endpoints on HF Spaces.

    Root and health endpoints stay public so the HF App tab can show status.
    MCP protocol endpoints are protected when CLIENT_API_KEY is set.
    """

    async def dispatch(self, request: Request, call_next):
        if request.url.path in {"/", "/health"} or request.method == "OPTIONS":
            return await call_next(request)

        if CLIENT_API_KEY:
            expected = f"Bearer {CLIENT_API_KEY}"
            if request.headers.get("authorization") != expected:
                return JSONResponse(
                    {"error": "Unauthorized. Use Authorization: Bearer <CLIENT_API_KEY>."},
                    status_code=401,
                )

        return await call_next(request)


async def root(request: Request):
    """Human-readable landing page for the Hugging Face App tab."""
    html = """
    <!doctype html>
    <html>
    <head>
      <title>Hydraulic Solver MCP Server</title>
      <style>
        body { font-family: system-ui, -apple-system, Segoe UI, sans-serif; max-width: 980px; margin: 48px auto; line-height: 1.55; }
        code { background: #f3f4f6; padding: 2px 5px; border-radius: 4px; }
        .card { border: 1px solid #e5e7eb; border-radius: 12px; padding: 18px 22px; margin: 16px 0; }
        .ok { color: #0a7f2e; font-weight: 700; }
      </style>
    </head>
    <body>
      <h1>💧 Hydraulic Solver MCP Server</h1>
      <p class="ok">Server is running.</p>
      <p>This Space is an agent-facing MCP server, not a Gradio visual app. The browser root page is only a status page.</p>

      <div class="card">
        <h2>Endpoints</h2>
        <p><b>Streamable HTTP MCP:</b> <code>/mcp</code></p>
        <p><b>Health check:</b> <code>/health</code></p>
      </div>

      <div class="card">
        <h2>How to connect</h2>
        <p>Use an MCP client such as Claude, Cursor, LM Studio, n8n, HuggingChat, Codex, LangChain, or AutoGen.</p>
        <p>Remote URL:</p>
        <p><code>https://YOUR-SPACE.hf.space/mcp</code></p>
        <p>If <code>CLIENT_API_KEY</code> is set as a Space secret, send:</p>
        <p><code>Authorization: Bearer your-secret</code></p>
      </div>

      <div class="card">
        <h2>Included tool families</h2>
        <ul>
          <li>Single Pipe Solver</li>
          <li>Hardy Cross Loop and Two-Loop Hardy Cross</li>
          <li>Three-Reservoir Solver</li>
          <li>PDD Demand Satisfaction</li>
          <li>Tank EPS Storage</li>
          <li>Valve Behavior</li>
          <li>Pump Operating Point</li>
          <li>Pressure Zone, Leakage, Water Age, and Chlorine teaching tools</li>
          <li>EPANET validation INP generator</li>
        </ul>
      </div>
    </body>
    </html>
    """
    return HTMLResponse(html)


async def health(request: Request):
    return JSONResponse(
        {
            "ok": True,
            "service": "hydraulic-solver-mcp-server",
            "transport": TRANSPORT,
            "mcp_endpoint": "/mcp",
            "auth_required": bool(CLIENT_API_KEY),
        }
    )


def create_streamable_http_app():
    """Create the Starlette ASGI app served on HF Spaces.

    The MCP SDK app already owns the /mcp route. We add only root and health
    routes so the Hugging Face App tab does not show plain "Not Found".
    """
    app = mcp.streamable_http_app()
    app.add_route("/", root, methods=["GET"])
    app.add_route("/health", health, methods=["GET"])
    app.add_middleware(BearerAuthMiddleware)
    return app



if __name__ == "__main__":
    if TRANSPORT in {"http", "streamable-http", "streamable_http"}:
        uvicorn.run(create_streamable_http_app(), host=HOST, port=PORT)
    elif TRANSPORT in {"sse"}:
        # SSE transport is kept for older clients. Streamable HTTP /mcp is preferred.
        mcp.run(transport="sse")
    else:
        mcp.run(transport="stdio")