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

from __future__ import annotations

import time
from pathlib import Path
from typing import Any

import aiosqlite

_SCHEMA = """
CREATE TABLE IF NOT EXISTS system_health_snapshot (
    component    TEXT PRIMARY KEY,
    last_status  TEXT NOT NULL,
    last_seen_at REAL NOT NULL,
    payload      TEXT NOT NULL DEFAULT '{}'
);

CREATE TABLE IF NOT EXISTS entity_baseline (
    entity_id   TEXT PRIMARY KEY,
    -- state when entity first entered unavailable/unknown
    state       TEXT NOT NULL,
    -- timestamp when the entity FIRST entered its current bad state (INSERT OR IGNORE)
    first_seen  REAL NOT NULL,
    -- kept for legacy compat; not used by UnavailableEntitiesCheck
    attributes  TEXT NOT NULL DEFAULT '{}',
    updated_at  REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS check_history (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    check_name  TEXT NOT NULL,
    ran_at      REAL NOT NULL,
    healthy     INTEGER NOT NULL,
    message     TEXT NOT NULL DEFAULT '',
    payload     TEXT NOT NULL DEFAULT '{}'
);

CREATE TABLE IF NOT EXISTS alerts_sent (
    alert_key   TEXT PRIMARY KEY,
    sent_at     REAL NOT NULL
);
"""

_MIGRATE_ENTITY_BASELINE = """
ALTER TABLE entity_baseline ADD COLUMN first_seen REAL NOT NULL DEFAULT 0;
"""


class Storage:
    def __init__(self, db_path: Path) -> None:
        self._db_path = db_path
        self._db: aiosqlite.Connection | None = None

    async def open(self) -> None:
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._db = await aiosqlite.connect(self._db_path)
        self._db.row_factory = aiosqlite.Row
        await self._db.executescript(_SCHEMA)
        # Add first_seen column to existing databases that pre-date Phase 3
        try:
            await self._db.execute(_MIGRATE_ENTITY_BASELINE)
        except Exception:
            pass  # column already exists
        await self._db.commit()

    async def close(self) -> None:
        if self._db:
            await self._db.close()
            self._db = None

    def _conn(self) -> aiosqlite.Connection:
        if self._db is None:
            raise RuntimeError("Storage not open — call await storage.open() first")
        return self._db

    # ------------------------------------------------------------------
    # entity_baseline — tracks entities currently in bad state
    # ------------------------------------------------------------------

    async def set_entity_unavailable_since(
        self, entity_id: str, state: str, first_seen: float
    ) -> None:
        """Record when an entity first entered unavailable/unknown state.

        INSERT OR IGNORE: if the entity is already tracked, preserves the
        original first_seen timestamp so duration is computed correctly.
        """
        await self._conn().execute(
            """
            INSERT OR IGNORE INTO entity_baseline
                (entity_id, state, first_seen, attributes, updated_at)
            VALUES (?, ?, ?, '{}', ?)
            """,
            (entity_id, state, first_seen, first_seen),
        )
        await self._conn().commit()

    async def get_entity_first_unavailable_at(self, entity_id: str) -> float | None:
        """Return when the entity first entered its bad state, or None if not tracked."""
        async with self._conn().execute(
            "SELECT first_seen FROM entity_baseline WHERE entity_id = ?",
            (entity_id,),
        ) as cur:
            row = await cur.fetchone()
            return float(row["first_seen"]) if row else None

    async def clear_entity_unavailable(self, entity_id: str) -> None:
        """Remove entity from unavailable tracking (entity has recovered)."""
        await self._conn().execute(
            "DELETE FROM entity_baseline WHERE entity_id = ?",
            (entity_id,),
        )
        await self._conn().commit()

    async def get_all_tracked_entity_ids(self) -> list[str]:
        """Return all entity IDs currently tracked as unavailable/unknown."""
        async with self._conn().execute(
            "SELECT entity_id FROM entity_baseline"
        ) as cur:
            rows = await cur.fetchall()
            return [r["entity_id"] for r in rows]

    # Legacy upsert — kept for backwards compat with existing callers
    async def upsert_entity_baseline(
        self, entity_id: str, state: str, attributes: str, updated_at: float
    ) -> None:
        await self._conn().execute(
            """
            INSERT INTO entity_baseline (entity_id, state, first_seen, attributes, updated_at)
            VALUES (?, ?, ?, ?, ?)
            ON CONFLICT(entity_id) DO UPDATE SET
                state = excluded.state,
                attributes = excluded.attributes,
                updated_at = excluded.updated_at
            """,
            (entity_id, state, updated_at, attributes, updated_at),
        )
        await self._conn().commit()

    async def get_entity_baseline(self, entity_id: str) -> dict[str, Any] | None:
        async with self._conn().execute(
            "SELECT * FROM entity_baseline WHERE entity_id = ?", (entity_id,)
        ) as cur:
            row = await cur.fetchone()
            return dict(row) if row else None

    # ------------------------------------------------------------------
    # check_history
    # ------------------------------------------------------------------

    async def record_check(
        self,
        check_name: str,
        ran_at: float,
        healthy: bool,
        message: str,
        payload: str,
    ) -> None:
        await self._conn().execute(
            """
            INSERT INTO check_history (check_name, ran_at, healthy, message, payload)
            VALUES (?, ?, ?, ?, ?)
            """,
            (check_name, ran_at, int(healthy), message, payload),
        )
        await self._conn().commit()

    # ------------------------------------------------------------------
    # alerts_sent (dedup gate)
    # ------------------------------------------------------------------

    async def was_alert_sent(self, alert_key: str, within_seconds: float) -> bool:
        cutoff = time.time() - within_seconds
        async with self._conn().execute(
            "SELECT sent_at FROM alerts_sent WHERE alert_key = ? AND sent_at > ?",
            (alert_key, cutoff),
        ) as cur:
            return (await cur.fetchone()) is not None

    async def mark_alert_sent(self, alert_key: str) -> None:
        await self._conn().execute(
            """
            INSERT INTO alerts_sent (alert_key, sent_at) VALUES (?, ?)
            ON CONFLICT(alert_key) DO UPDATE SET sent_at = excluded.sent_at
            """,
            (alert_key, time.time()),
        )
        await self._conn().commit()

    async def clear_alert(self, alert_key: str) -> None:
        """Delete an alert record so the next occurrence triggers immediately."""
        await self._conn().execute(
            "DELETE FROM alerts_sent WHERE alert_key = ?", (alert_key,)
        )
        await self._conn().commit()

    # ------------------------------------------------------------------
    # system_health_snapshot — tracks last-known per-component status
    # ------------------------------------------------------------------

    async def get_system_health_snapshot(
        self, component: str
    ) -> dict[str, Any] | None:
        """Return the stored snapshot for a component, or None if unseen."""
        async with self._conn().execute(
            "SELECT * FROM system_health_snapshot WHERE component = ?",
            (component,),
        ) as cur:
            row = await cur.fetchone()
            return dict(row) if row else None

    async def upsert_system_health_snapshot(
        self, component: str, last_status: str, payload: str
    ) -> None:
        """Insert or replace the snapshot for a component."""
        await self._conn().execute(
            """
            INSERT INTO system_health_snapshot
                (component, last_status, last_seen_at, payload)
            VALUES (?, ?, ?, ?)
            ON CONFLICT(component) DO UPDATE SET
                last_status  = excluded.last_status,
                last_seen_at = excluded.last_seen_at,
                payload      = excluded.payload
            """,
            (component, last_status, time.time(), payload),
        )
        await self._conn().commit()