from fastapi import APIRouter
import psutil
import time
import os

router = APIRouter()

def get_size(bytes, suffix="B"):
    """
    Scale bytes to its proper format
    e.g:
        1253656 => '1.20MB'
        1253656678 => '1.17GB'
    """
    factor = 1024
    for unit in ["", "K", "M", "G", "T", "P"]:
        if bytes < factor:
            return f"{bytes:.2f}{unit}{suffix}"
        bytes /= factor

@router.get("/stats")
def get_system_stats():
    # CPU
    cpu_usage = psutil.cpu_percent(interval=1)
    
    # Memory
    svmem = psutil.virtual_memory()
    memory_stats = {
        "total": get_size(svmem.total),
        "available": get_size(svmem.available),
        "percent": svmem.percent,
        "used": get_size(svmem.used)
    }
    
    # Disk (Root)
    try:
        partition = psutil.disk_usage("/")
        disk_stats = {
            "total": get_size(partition.total),
            "free": get_size(partition.free),
            "percent": partition.percent,
            "used": get_size(partition.used)
        }
    except Exception:
        disk_stats = {"error": "Unavailable"}

    # Uptime
    boot_time = psutil.boot_time()
    uptime_seconds = time.time() - boot_time
    uptime_string = time.strftime("%H:%M:%S", time.gmtime(uptime_seconds))
    if uptime_seconds > 86400:
        days = int(uptime_seconds // 86400)
        uptime_string = f"{days} days, " + uptime_string

    return {
        "cpu": cpu_usage,
        "memory": memory_stats,
        "disk": disk_stats,
        "uptime": uptime_string
    }