homelab-codex-ws/services/ha-diag-agent/src/ha_diag/main.py

from __future__ import annotations

import asyncio
import json
import logging
import time
from datetime import datetime

import structlog
import uvicorn
from apscheduler.schedulers.asyncio import AsyncIOScheduler

from .api import app, register_checks
from .checks.heartbeat import HeartbeatCheck
from .config import Settings
from .event_emitter import EventEmitter
from .ha_client import HAClient
from .storage import Storage

_log = structlog.get_logger()


def _configure_structlog(log_level: str) -> None:
    structlog.configure(
        processors=[
            structlog.processors.add_log_level,
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.processors.StackInfoRenderer(),
            structlog.processors.format_exc_info,
            structlog.processors.JSONRenderer(),
        ],
        logger_factory=structlog.PrintLoggerFactory(),
    )
    logging.basicConfig(level=getattr(logging, log_level.upper(), logging.INFO))


async def _run_check_and_emit(check, emitter: EventEmitter, storage: Storage) -> None:
    try:
        result = await check.run()
        await storage.record_check(
            check_name=check.name,
            ran_at=time.time(),
            healthy=result.healthy,
            message=result.message,
            payload=json.dumps(result.payload),
        )
        if result.event_type:
            emitter.emit(
                event_type=result.event_type,
                severity=result.severity.value,
                service="homeassistant",
                message=result.message,
                payload=result.payload,
            )
            _log.warning(
                "check_unhealthy",
                check=check.name,
                event=result.event_type,
                msg=result.message,
            )
        else:
            _log.info("check_ok", check=check.name)
    except Exception as exc:
        _log.error("check_error", check=check.name, error=str(exc), exc_info=True)


async def run(settings: Settings) -> None:
    _configure_structlog(settings.log_level)
    _log.info(
        "ha_diag_agent_starting",
        node=settings.node_name,
        location=settings.location_tag,
        ha_url=settings.ha_url,
        interval=settings.check_interval,
    )

    storage = Storage(settings.data_dir / "ha_diag.db")
    await storage.open()

    emitter = EventEmitter(settings.events_dir, settings.node_name, settings.location_tag)
    ha_client = HAClient(settings.ha_url, settings.ha_token)

    checks = [HeartbeatCheck(ha_client)]
    register_checks(checks, settings.node_name, settings.location_tag)

    scheduler = AsyncIOScheduler()
    for check in checks:
        scheduler.add_job(
            _run_check_and_emit,
            "interval",
            seconds=settings.check_interval,
            args=[check, emitter, storage],
            id=f"check_{check.name}",
            next_run_time=datetime.now(),
        )
    scheduler.start()
    _log.info("scheduler_started", checks=[c.name for c in checks])

    config = uvicorn.Config(
        app,
        host="0.0.0.0",
        port=settings.port,
        log_level=settings.log_level.lower(),
    )
    server = uvicorn.Server(config)
    try:
        await server.serve()
    finally:
        scheduler.shutdown(wait=False)
        await storage.close()


def main() -> None:
    settings = Settings.load()
    asyncio.run(run(settings))


if __name__ == "__main__":
    main()