app/main.py

from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.openapi.utils import get_openapi

from app.core.storage import ensure_storage_dirs
from app.core.config import Settings

# Page renderers
from app.pages.landing import render_landing
from app.pages.phone import render_phone_page
from app.pages.gbridge import render_gbridge_page
from app.pages.windsurf import render_windsurf_page
from app.pages.receipts import render_receipts_page
from app.pages.chat import render_chat_page
from app.pages.m5 import render_m5_page
from app.pages.m5_landing import render_m5_landing
from app.pages.docs import render_docs
from app.pages.api_docs import render_api_docs
from app.pages.public_chat import render_public_chat

# API routes
from app.api.routes_sessions import router as sessions_router
from app.api.routes_jobs import router as jobs_router
from app.api.routes_receipts import router as receipts_router
from app.api.routes_gbridge import router as gbridge_router
from app.api.routes_windsurf import router as windsurf_router
from app.api.routes_health import router as health_router
from app.api.routes_chat import router as chat_router
from app.api.routes_mesh import router as mesh_router
from app.api.routes_accounts import router as accounts_router
from app.api.routes_openai import router as openai_router
from app.api.routes_auth import router as auth_router
from app.api.routes_chats_public import router as chats_public_router
from app.api.routes_pipeline import router as pipeline_router

# Phone relay
from app.phone.relay import (
    connect_worker_websocket, connect_client_websocket,
    disconnect_worker, disconnect_client,
    handle_worker_message, handle_client_message,
    handle_mac_node_message,
)
from app.mac_relay import connect_mac_node, disconnect_mac_node
from app.phone.sessions import verify_worker_join_token

# gbridge router
from app.gbridge.router import route_capability_call

app = FastAPI(title="M5 API Gateway", version="1.0.0", description="Local compute gateway with OpenAI-compatible API")


def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema
    openapi_schema = get_openapi(
        title="M5 API Gateway",
        version="1.0.0",
        description="Local compute gateway with OpenAI-compatible API",
        routes=app.routes,
    )
    openapi_schema["info"]["x-logo"] = {"url": "https://m5.hf.space/logo.png"}
    app.openapi_schema = openapi_schema
    return app.openapi_schema


app.openapi = custom_openapi

# Mount static assets
app.mount("/static", StaticFiles(directory="app/static", html=True), name="static")

# Include API routes
app.include_router(sessions_router)
app.include_router(jobs_router)
app.include_router(receipts_router)
app.include_router(gbridge_router)
app.include_router(windsurf_router)
app.include_router(health_router)
app.include_router(chat_router)
app.include_router(mesh_router)
app.include_router(accounts_router)
app.include_router(openai_router)
app.include_router(auth_router)
app.include_router(chats_public_router)
app.include_router(pipeline_router)


@app.on_event("startup")
async def startup_event():
    ensure_storage_dirs()
    # Attempt to start Ollama service (only works if Ollama is installed locally)
    try:
        from app.compute.ollama_manager import get_ollama_manager
        ollama = get_ollama_manager()
        await ollama.start()
        await ollama.ensure_default_models()
    except Exception as e:
        import logging
        logging.getLogger(__name__).warning(f"Ollama manager not available in this environment: {e}")


@app.on_event("shutdown")
async def shutdown_event():
    pass


# Page routes
@app.get("/", response_class=HTMLResponse)
async def root():
    return render_landing()


@app.get("/phone", response_class=HTMLResponse)
async def phone_page(request: Request):
    return render_phone_page(request)


@app.get("/gbridge", response_class=HTMLResponse)
async def gbridge_page(request: Request):
    return render_gbridge_page(request)


@app.get("/windsurf", response_class=HTMLResponse)
async def windsurf_page(request: Request):
    return render_windsurf_page(request)


@app.get("/chat", response_class=HTMLResponse)
async def chat_page(request: Request):
    return render_chat_page(request)


@app.get("/receipts", response_class=HTMLResponse)
async def receipts_page(request: Request):
    return render_receipts_page(request)


@app.get("/m5", response_class=HTMLResponse)
async def m5_page(request: Request):
    return render_m5_page(request)


@app.get("/m5-landing", response_class=HTMLResponse)
async def m5_landing_page(request: Request):
    return render_m5_landing()


@app.get("/docs", response_class=HTMLResponse)
async def docs_page(request: Request):
    return render_docs()


@app.get("/api", response_class=HTMLResponse)
async def api_docs_page(request: Request):
    return render_api_docs()


@app.get("/c/{public_id}", response_class=HTMLResponse)
async def public_chat_page(public_id: str, request: Request):
    return render_public_chat()


@app.get("/dashboard", response_class=HTMLResponse)
async def dashboard_page():
    with open("app/static/dashboard.html", "r") as f:
        return f.read()


@app.get("/worker")
async def worker_page(request: Request, session_id: str, token: str):
    if not verify_worker_join_token(session_id, token):
        return HTMLResponse(content="<h1>Invalid token</h1>", status_code=403)
    # Serve the iPhone worker PWA
    with open("app/static/worker/worker.html", "r") as f:
        return HTMLResponse(content=f.read(), status_code=200)


# gbridge capability endpoint — the core primitive
@app.post("/cap/{capability_name}")
async def capability_endpoint(capability_name: str, request: Request):
    body = await request.json()
    session_id = body.get("session_id")
    payload = body.get("payload", {})
    result = route_capability_call(capability_name, payload, session_id)
    return result


# WebSocket routes
@app.websocket("/ws/gbridge/worker/{session_id}/{worker_id}")
async def worker_websocket(ws: WebSocket, session_id: str, worker_id: str):
    await ws.accept()
    await connect_worker_websocket(session_id, worker_id, ws)
    try:
        while True:
            data = await ws.receive_json()
            await handle_worker_message(session_id, worker_id, data)
    except WebSocketDisconnect:
        await disconnect_worker(session_id, worker_id)


@app.websocket("/ws/gbridge/client/{session_id}")
async def client_websocket(ws: WebSocket, session_id: str):
    await ws.accept()
    await connect_client_websocket(session_id, ws)
    try:
        while True:
            data = await ws.receive_json()
            await handle_client_message(session_id, data)
    except WebSocketDisconnect:
        await disconnect_client(session_id)


@app.websocket("/ws/mac/{session_id}/{node_id}")
async def mac_websocket(ws: WebSocket, session_id: str, node_id: str):
    await ws.accept()
    await connect_mac_node(node_id, ws)
    try:
        while True:
            data = await ws.receive_json()
            await handle_mac_node_message(node_id, data)
    except WebSocketDisconnect:
        await disconnect_mac_node(node_id)