wrolpi/controller/main.py

"""
WROLPi Controller - FastAPI Application

The Controller provides:
- System status monitoring (CPU, memory, drives, network)
- Service management (start/stop/restart systemd services)
- Disk mounting and management
- Admin operations (hotspot, throttle, shutdown, reboot)
- Simple emergency UI when main app is down
"""

import asyncio
import logging
from contextlib import asynccontextmanager
from pathlib import Path

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(name)s %(levelname)s %(message)s')
logger = logging.getLogger(__name__)

from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates

from controller import __version__
from fastapi.middleware.cors import CORSMiddleware

from controller.api.admin import router as admin_router
from controller.api.disks import router as disks_router
from controller.api.onboarding import router as onboarding_router
from controller.api.ready import router as ready_router
from controller.api.schemas import HealthResponse
from controller.api.scripts import router as scripts_router
from controller.api.services import router as services_router
from controller.api.status import router as status_router
from controller.lib.config import (
    is_docker_mode,
    is_primary_drive_mounted,
    reload_config_from_drive,
)
from controller.lib.docker_services import (
    can_manage_containers,
    get_all_containers_status,
)
from controller.lib.status import (
    get_cpu_status,
    get_iostat_status,
    get_load_status,
    get_memory_status,
    get_primary_drive_status,
    get_uptime_status,
)
from controller.lib.status_worker import status_worker_loop
from controller.lib.systemd import get_all_services_status

# Templates directory
TEMPLATES_DIR = Path(__file__).parent / "templates"
templates = Jinja2Templates(directory=str(TEMPLATES_DIR))


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Application lifespan handler.
    Runs on startup and shutdown.
    """
    # Startup
    logger.info("WROLPi Controller v%s starting...", __version__)
    logger.info("Docker mode: %s", is_docker_mode())

    # Try to load config from drive if mounted
    if is_primary_drive_mounted():
        if reload_config_from_drive():
            logger.info("Loaded configuration from drive")
        else:
            logger.info("Using default configuration (no controller.yaml found)")

        # Apply timezone from wrolpi.yaml to the system.
        from controller.lib.admin import apply_timezone_from_config
        apply_timezone_from_config()
    else:
        logger.info("Primary drive not mounted - using default configuration")

    # Initialize cached status
    app.state.cached_status = {}

    # Start status worker background task
    status_task = asyncio.create_task(status_worker_loop(app))

    yield

    # Shutdown
    logger.info("WROLPi Controller shutting down...")
    status_task.cancel()
    try:
        await status_task
    except asyncio.CancelledError:
        pass


# Create FastAPI app
app = FastAPI(
    title="WROLPi Controller",
    description="System management service for WROLPi",
    version=__version__,
    lifespan=lifespan,
)

# CORS middleware - matches Caddy's permissive policy for offline/local use.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include routers
app.include_router(admin_router)
app.include_router(disks_router)
app.include_router(onboarding_router)
app.include_router(ready_router)
app.include_router(scripts_router)
app.include_router(services_router)
app.include_router(status_router)


def get_host(request: Request) -> str:
    """Get the host for links (use request host, fallback to localhost)."""
    host = request.headers.get("host", "localhost")
    # Strip port if present
    if ":" in host:
        host = host.split(":")[0]
    return host


@app.get("/", response_class=HTMLResponse)
async def dashboard(request: Request):
    """Render the main dashboard UI with real status data."""
    # Try to use cached status first
    cached = getattr(request.app.state, "cached_status", {})

    if cached:
        cpu = cached.get("cpu_stats") or {}
        memory = cached.get("memory_stats") or {}
        load = cached.get("load_stats") or {}
        iostat = cached.get("iostat_stats") or {}
        uptime = cached.get("uptime_stats") or {}
        drives = cached.get("drives_stats") or []
        primary_drive = None
        for drive in drives:
            if drive.get("mount") == "/media/wrolpi":
                primary_drive = drive
                break
    else:
        # Fallback to direct calls if cache not yet populated
        cpu = get_cpu_status()
        memory = get_memory_status()
        load = get_load_status()
        iostat = get_iostat_status()
        uptime = get_uptime_status()
        primary_drive = get_primary_drive_status()

    # Format storage data
    storage = {}
    if primary_drive:
        free_bytes = primary_drive["size"] - primary_drive["used"]
        storage = {
            "percent": primary_drive["percent"],
            "free_gb": round(free_bytes / (1024 ** 3), 1),
        }

    # Format uptime
    uptime_seconds = uptime.get("uptime_seconds", 0)
    if uptime_seconds:
        days = uptime_seconds // 86400
        hours = (uptime_seconds % 86400) // 3600
        uptime_formatted = f"{days}d {hours}h" if days > 0 else f"{hours}h {(uptime_seconds % 3600) // 60}m"
    else:
        uptime_formatted = "--"

    # Get service status
    if is_docker_mode():
        if can_manage_containers():
            services = get_all_containers_status()
        else:
            services = []
    else:
        services = get_all_services_status()

    # Sort services by name for consistent display
    services = sorted(services, key=lambda s: s["name"])

    context = {
        "version": __version__,
        "docker_mode": is_docker_mode(),
        "drive_mounted": is_primary_drive_mounted(),
        "host": get_host(request),

        # Real status data
        "cpu": {
            "percent": cpu["percent"],
            "temperature_c": cpu["temperature"],  # Field renamed
        },
        "memory": {
            "percent": round((memory["used"] / memory["total"]) * 100, 1) if memory["total"] > 0 else 0,
            "used_gb": round(memory["used"] / (1024 ** 3), 1),
        },
        "load": {
            "load_1min": load["minute_1"],  # Field renamed
        },
        "storage": storage,
        "iostat": {
            "percent_iowait": iostat.get("percent_iowait"),
        },
        "uptime": {
            "formatted": uptime_formatted,
            "seconds": uptime_seconds,
        },

        # Real service status
        "services": services,
    }

    return templates.TemplateResponse(request, "index.html", context)


@app.get("/api/health", response_model=HealthResponse)
async def health_check() -> HealthResponse:
    """
    Health check endpoint.
    Returns basic status for monitoring and load balancer checks.
    """
    return HealthResponse(
        status="healthy",
        version=__version__,
        docker_mode=is_docker_mode(),
        drive_mounted=is_primary_drive_mounted(),
    )