import psutil
import asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
app = FastAPI()


# @app.get("/hardware/status")
def hardware_status():

    cpu = psutil.cpu_percent(interval=1)
    memory = psutil.virtual_memory()
    disk = psutil.disk_usage("/")

    return {
        "cpu_usage_percent": cpu,

        "memory_usage_percent": memory.percent,
        "memory_total_mb": round(
            memory.total / 1024 / 1024,
            2
        ),
        "memory_used_mb": round(
            memory.used / 1024 / 1024,
            2
        ),
        "memory_available_mb": round(
            memory.available / 1024 / 1024,
            2
        ),

        "disk_usage_percent": disk.percent,
        "disk_total_gb": round(
            disk.total / 1024 / 1024 / 1024,
            2
        ),
        "disk_used_gb": round(
            disk.used / 1024 / 1024 / 1024,
            2
        ),
        "disk_free_gb": round(
            disk.free / 1024 / 1024 / 1024,
            2
        )
    }

@app.websocket("/ws/hardware/status")
async def hardware_status_monitor(websocket: WebSocket):

    await websocket.accept()

    try:

        while True:

            try:

                result = await asyncio.wait_for(
                    asyncio.to_thread(
                        hardware_status
                    ),
                    timeout=3
                )

            except asyncio.TimeoutError:

                result = {
                    "cpu_usage_percent": None,
                    "memory_usage_percent": None,
                    "memory_total_mb": None,
                    "memory_used_mb": None,
                    "memory_available_mb": None,
                    "disk_usage_percent": None,
                    "disk_total_gb": None,
                    "disk_used_gb": None,
                    "disk_free_gb": None,
                    "error": "Hardware monitoring timed out"
                }

            await websocket.send_json(result)

            await asyncio.sleep(1)

    except WebSocketDisconnect:

        print("Hardware monitoring WebSocket disconnected")

if __name__ == "__main__":

    import uvicorn

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=9100
    )