app.py

import asyncio
import os
from datetime import datetime
from pathlib import Path
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse

# Initialize FastAPI application
app = FastAPI()

# Track brain module status
_brain_loaded = False
try:
    import brain_minimal
    _brain_loaded = True
except ImportError:
    _brain_loaded = False


@app.get("/health")
async def health_check():
    """Structured health check endpoint returning server status."""
    print(f"[HEALTH] GET /health request at {datetime.utcnow().isoformat()}")
    return {
        "timestamp": datetime.utcnow().isoformat(),
        "startup_mode": os.environ.get("STARTUP_MODE", "unknown"),
        "brain_loaded": _brain_loaded
    }


@app.get("/status")
async def status_check():
    """Comprehensive health check with db, env, and filesystem validation."""
    checks = {"db": "ok", "env": "ok", "fs": "ok"}
    overall_healthy = True

    # 1. Database check - lightweight dataset list call
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        # Lightweight call - just check if we can reach HF
        repo_id = os.environ.get("OPENCLAW_DATASET_REPO", "")
        if repo_id:
            api.repo_info(repo_id=repo_id, repo_type="dataset")
            checks["db"] = "ok"
        else:
            checks["db"] = "error: OPENCLAW_DATASET_REPO not set"
            overall_healthy = False
    except Exception as e:
        checks["db"] = f"error: {str(e)}"
        overall_healthy = False

    # 2. Environment check - verify critical variables
    critical_vars = ["HF_TOKEN", "OPENCLAW_DATASET_REPO"]
    missing_vars = [v for v in critical_vars if not os.environ.get(v)]
    if missing_vars:
        checks["env"] = f"error: missing {', '.join(missing_vars)}"
        overall_healthy = False
    else:
        checks["env"] = "ok"

    # 3. Filesystem check - touch and delete test file
    test_file = Path("/tmp/test_write")
    try:
        test_file.touch()
        test_file.unlink()
        checks["fs"] = "ok"
    except Exception as e:
        checks["fs"] = f"error: {str(e)}"
        overall_healthy = False

    status = "healthy" if overall_healthy else "degraded"

    # Return HTTP 503 if critical checks fail
    if not overall_healthy:
        raise HTTPException(status_code=503, detail={
            "status": status,
            "checks": checks
        })

    return {"status": status, "checks": checks}


@app.websocket("/ws/agents")
async def websocket_agents(websocket: WebSocket):
    """WebSocket endpoint for real-time agent status updates."""
    await websocket.accept()
    try:
        while True:
            # Send agent status every 2 seconds
            await websocket.send_json({
                "agents": [
                    {"name": "Adam", "status": "Thinking"},
                    {"name": "Eve", "status": "Writing"}
                ]
            })
            await asyncio.sleep(2)
    except WebSocketDisconnect:
        print("[WS] Client disconnected")


# Mount static files to serve frontend at root path
app.mount("/", StaticFiles(directory="frontend", html=True), name="frontend")


# Main entry point
if __name__ == "__main__":
    print("Cain: Server running on 7860")
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)