server/app.py

"""FastAPI server for the Space Manufacturing RL OpenEnv environment."""
from __future__ import annotations

import asyncio
import os
import sys
import uvicorn
from pathlib import Path
from typing import Any, Dict, Optional

# Make 'SpaceFactory' importable regardless of working directory.
# server/app.py is at <root>/SpaceFactory/server/app.py → parent×3 = <root>
_pkg_parent = str(Path(__file__).resolve().parent.parent.parent)
if _pkg_parent not in sys.path:
    sys.path.insert(0, _pkg_parent)

try:
    from fastapi import FastAPI, HTTPException, Query
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.responses import FileResponse, HTMLResponse
    from fastapi.staticfiles import StaticFiles
    from pydantic import BaseModel
except ImportError as exc:
    raise SystemExit(f"FastAPI is required: pip install fastapi uvicorn\n{exc}") from exc

from .ui_demo import ManufacturingUIDemo
from SpaceFactory.env import ManufacturingTaskEnv
from SpaceFactory.models import ManufacturingAction, ManufacturingObservation

# ── app ────────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Space Manufacturing RL",
    description="Custom space manufacturing frontend with mounted OpenEnv API routes.",
    version="0.1.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:4173", "http://127.0.0.1:4173", "http://localhost:5173", "http://127.0.0.1:5173"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
async def _suppress_win32_connection_reset() -> None:
    """Suppress spurious WinError 10054 from asyncio ProactorEventLoop on Windows."""
    if sys.platform != "win32":
        return
    loop = asyncio.get_event_loop()

    def _handler(loop: asyncio.AbstractEventLoop, context: dict) -> None:
        exc = context.get("exception")
        if isinstance(exc, ConnectionResetError):
            return  # benign Windows ProactorEventLoop artifact — silently ignore
        loop.default_exception_handler(context)

    loop.set_exception_handler(_handler)

# Single shared environment instance (stateful per process)
_env = ManufacturingTaskEnv(task_name="easy")

# UI demo instance
_ui_demo = ManufacturingUIDemo()

# Mount built frontend static assets if dist/ exists
_here   = Path(__file__).resolve().parent
_dist   = _here.parent / "frontend" / "dist"
_assets = _dist / "assets"
if _assets.exists():
    app.mount("/assets",     StaticFiles(directory=str(_assets)), name="assets")
    app.mount("/web/assets", StaticFiles(directory=str(_assets)), name="web-assets")


# ── helpers ────────────────────────────────────────────────────────────────────

class StepRequest(BaseModel):
    """Wrapper for a single environment step action."""

    action: Dict[str, Any]
    timeout_s: Optional[float] = 30.0


# ── OpenEnv core endpoints ─────────────────────────────────────────────────────

@app.post("/reset")
def reset(task: Optional[str] = Query(default=None)):
    """Reset the episode, optionally switching to a different task."""
    if task is not None:
        try:
            _env.set_task(task)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
    return _env.reset().model_dump()


@app.post("/step")
def step(body: StepRequest):
    """Advance the environment by one step."""
    try:
        action = ManufacturingAction.model_validate(body.action)
    except Exception as exc:
        raise HTTPException(status_code=422, detail=f"Invalid action: {exc}") from exc
    obs, reward, done, info = _env.step(action)
    return {"observation": obs.model_dump(), "reward": reward.model_dump(),
            "done": done, "info": info}


@app.get("/state")
def state():
    """Return the full current environment state snapshot."""
    return _env.state().model_dump()


@app.get("/tasks")
def list_tasks():
    """List all available task presets."""
    return ManufacturingTaskEnv.list_tasks()


@app.get("/schema")
def schema():
    """Return action, observation, and state JSON schemas."""
    return {
        "action_space":      ManufacturingAction.model_json_schema(),
        "observation_space": ManufacturingObservation.model_json_schema(),
        "state":             _env.state().model_json_schema(),
    }


@app.get("/health")
def health():
    """Health probe used by OpenEnv validators."""
    return {"status": "healthy"}


@app.get("/metadata")
def metadata():
    """Environment metadata required by the OpenEnv runtime standard."""
    return {
        "name": "space_manufacturing",
        "description": (
            "Orbital manufacturing platform coordination environment "
            "for production and delivery scheduling."
        ),
        "version": "0.1.0",
        "tasks": list(ManufacturingTaskEnv.list_tasks().keys()),
    }


# ── UI demo endpoints ──────────────────────────────────────────────────────────

@app.get("/api/ui/demo")
def ui_demo_snapshot():
    """Return the current UI demo snapshot (used by the React frontend)."""
    return _ui_demo.snapshot()


@app.post("/api/ui/demo/reset")
def ui_demo_reset(task_name: str = "medium"):
    """Reset the UI demo to a specific task."""
    if task_name not in ManufacturingTaskEnv.list_tasks():
        raise HTTPException(status_code=400, detail=f"Unknown task: {task_name}")
    return _ui_demo.reset(task_name)


@app.post("/api/ui/demo/step")
def ui_demo_step():
    """Advance the UI demo by one heuristic step."""
    return _ui_demo.step()


# ── Frontend SPA routes ────────────────────────────────────────────────────────

@app.get("/web", response_class=FileResponse, include_in_schema=False)
def serve_web_index():
    """Serve the React SPA index page."""
    index = _dist / "index.html"
    if index.exists():
        return FileResponse(str(index))
    return HTMLResponse(_fallback_html())


@app.get("/web/{full_path:path}", include_in_schema=False)
def serve_web_app(full_path: str):
    """Serve built SPA assets, falling back to index.html for client-side routing."""
    candidate = _dist / full_path
    if candidate.exists() and candidate.is_file():
        return FileResponse(str(candidate))
    index = _dist / "index.html"
    if index.exists():
        return FileResponse(str(index))
    return HTMLResponse(_fallback_html())


def _fallback_html() -> str:
    """Minimal status page shown when the frontend dist/ hasn't been built yet."""
    s = _env.state()
    platforms_html = "".join(
        f"<li>Platform {p.id}: energy={p.energy:.1f}, "
        f"mat={p.material_stock:.1f}, comp={p.component_stock:.1f}, "
        f"prod={p.product_stock}, last={p.last_action}</li>"
        for p in s.platforms
    )
    return (
        f"<!DOCTYPE html><html><head><title>Space Manufacturing RL</title>"
        f"<style>body{{font-family:monospace;background:#06111f;color:#c8deff;padding:24px}}"
        f"h1{{color:#7cf7c9}}li{{margin:4px 0}}"
        f".warn{{color:#ffb86b;font-size:.85rem;margin-top:16px}}</style></head><body>"
        f"<h1>Space Manufacturing RL — {s.task_name}</h1>"
        f"<p>Step: {s.step_count} / {s.max_steps} | "
        f"Total reward: {s.total_reward:.2f} | Done: {s.done}</p>"
        f"<h2>Platforms</h2><ul>{platforms_html}</ul>"
        f"<p class='warn'>Frontend not built. Run: "
        f"cd frontend &amp;&amp; npm install &amp;&amp; npm run build</p>"
        f"</body></html>"
    )


# ── Entry point ────────────────────────────────────────────────────────────────

def main() -> None:
    """Start the uvicorn server with Windows-compatible event loop policy."""
    # Windows ProactorEventLoop raises spurious ConnectionResetError (WinError 10054)
    # when clients close keep-alive connections. Selector loop avoids the issue.
    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

    uvicorn.run("SpaceFactory.server.app:app", host="0.0.0.0", port=8000, reload=False)


if __name__ == "__main__":
    main()