import psutil
import asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect,HTTPException
import docker

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")


# GET DOCKER CONTAINERS
@app.get("/docker/containers")
def get_containers():

    try:
        # Connect to the Docker daemon running on this VM
        client = docker.from_env()

        # all=True means:
        # running + stopped + exited containers
        containers = client.containers.list(all=True)

        result = []

        for container in containers:

            result.append({
                "id": container.short_id,
                "name": container.name,
                "image": (
                    container.image.tags[0]
                    if container.image.tags
                    else container.image.short_id
                ),
                "status": container.status,
                "ports": container.attrs[
                    "NetworkSettings"
                ]["Ports"]
            })

        return {
            "count": len(result),
            "containers": result
        }

    except docker.errors.DockerException as e:

        raise HTTPException(
            status_code=500,
            detail=f"Docker error: {str(e)}"
        )

# websocket for docker containers
def get_container_status():

    client = docker.from_env()

    containers = client.containers.list(all=True)

    result = []

    for container in containers:

        result.append({
            "id": container.short_id,
            "name": container.name,
            "status": container.status
        })

    return result

@app.websocket("/ws/docker/containers")
async def docker_container_status(websocket: WebSocket):

    await websocket.accept()

    try:

        while True:

            # Docker SDK is synchronous/blocking,
            # so run it in a worker thread.
            containers = await asyncio.to_thread(
                get_container_status
            )

            await websocket.send_json({
                "containers": containers
            })

            # Update every 2 seconds
            await asyncio.sleep(2)

    except WebSocketDisconnect:

        print("Docker WebSocket disconnected")

    except Exception as e:

        try:
            await websocket.send_json({
                "error": str(e)
            })

        except Exception:
            pass



if __name__ == "__main__":

    import uvicorn

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