from __future__ import annotations import math import sqlite3 from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Any def utc_now() -> datetime: return datetime.now(UTC) def to_db_time(value: datetime) -> str: return value.astimezone(UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z") def from_db_time(value: str) -> datetime: return datetime.fromisoformat(value.replace("Z", "+00:00")) class Storage: def __init__(self, path: str): self.path = path self.db: sqlite3.Connection | None = None def open(self) -> None: Path(self.path).parent.mkdir(parents=True, exist_ok=True) self.db = sqlite3.connect(self.path, check_same_thread=False) self.db.row_factory = sqlite3.Row self.db.execute("PRAGMA journal_mode=WAL") self.db.execute("PRAGMA synchronous=NORMAL") self.db.executescript( """ CREATE TABLE IF NOT EXISTS samples ( id INTEGER PRIMARY KEY AUTOINCREMENT, sampled_at TEXT NOT NULL, watts REAL NOT NULL, voltage REAL, current_amps REAL, socket_on INTEGER NOT NULL ); CREATE INDEX IF NOT EXISTS idx_samples_at ON samples(sampled_at); CREATE TABLE IF NOT EXISTS cycles ( id INTEGER PRIMARY KEY AUTOINCREMENT, started_at TEXT NOT NULL, ended_at TEXT, duration_seconds REAL, average_watts REAL, maximum_watts REAL, energy_wh REAL, sample_count INTEGER, incomplete_start INTEGER NOT NULL DEFAULT 0, incomplete_end INTEGER NOT NULL DEFAULT 0 ); CREATE INDEX IF NOT EXISTS idx_cycles_started ON cycles(started_at); """ ) self.db.commit() def close(self) -> None: if self.db is not None: self.db.close() self.db = None def _conn(self) -> sqlite3.Connection: if self.db is None: raise RuntimeError("Database is not open") return self.db def record_sample( self, sampled_at: datetime, watts: float, voltage: float | None, current_amps: float | None, socket_on: bool, ) -> None: db = self._conn() db.execute( "INSERT INTO samples(sampled_at, watts, voltage, current_amps, socket_on) VALUES (?, ?, ?, ?, ?)", (to_db_time(sampled_at), watts, voltage, current_amps, int(socket_on)), ) db.commit() def start_cycle(self, started_at: datetime, incomplete_start: bool = False) -> int: db = self._conn() cursor = db.execute( "INSERT INTO cycles(started_at, incomplete_start) VALUES (?, ?)", (to_db_time(started_at), int(incomplete_start)), ) db.commit() return int(cursor.lastrowid) def finish_cycle(self, cycle_id: int, ended_at: datetime, incomplete_end: bool = False) -> dict[str, Any]: db = self._conn() cycle = db.execute("SELECT * FROM cycles WHERE id = ?", (cycle_id,)).fetchone() if cycle is None: raise RuntimeError(f"Cycle {cycle_id} does not exist") started_at = from_db_time(cycle["started_at"]) if ended_at < started_at: ended_at = started_at start_text = to_db_time(started_at) end_text = to_db_time(ended_at) rows = db.execute( "SELECT sampled_at, watts FROM samples WHERE sampled_at >= ? AND sampled_at < ? ORDER BY sampled_at", (start_text, end_text), ).fetchall() values = [(from_db_time(row["sampled_at"]), float(row["watts"])) for row in rows] duration = max(0.0, (ended_at - started_at).total_seconds()) average = sum(watts for _, watts in values) / len(values) if values else 0.0 maximum = max((watts for _, watts in values), default=0.0) energy_wh = self._integrate_energy(values, ended_at) db.execute( """ UPDATE cycles SET ended_at = ?, duration_seconds = ?, average_watts = ?, maximum_watts = ?, energy_wh = ?, sample_count = ?, incomplete_end = ? WHERE id = ? """, (end_text, duration, average, maximum, energy_wh, len(values), int(incomplete_end), cycle_id), ) db.commit() return self.get_cycle(cycle_id) @staticmethod def _integrate_energy(values: list[tuple[datetime, float]], ended_at: datetime) -> float: if not values: return 0.0 watt_seconds = 0.0 for index, (sampled_at, watts) in enumerate(values): segment_end = values[index + 1][0] if index + 1 < len(values) else ended_at seconds = max(0.0, (segment_end - sampled_at).total_seconds()) watt_seconds += watts * seconds return watt_seconds / 3600.0 def get_cycle(self, cycle_id: int) -> dict[str, Any]: row = self._conn().execute("SELECT * FROM cycles WHERE id = ?", (cycle_id,)).fetchone() if row is None: raise RuntimeError(f"Cycle {cycle_id} does not exist") return dict(row) def close_open_cycles(self) -> int: db = self._conn() open_rows = db.execute("SELECT id, started_at FROM cycles WHERE ended_at IS NULL").fetchall() closed = 0 for row in open_rows: last_sample = db.execute( "SELECT sampled_at FROM samples WHERE sampled_at >= ? ORDER BY sampled_at DESC LIMIT 1", (row["started_at"],), ).fetchone() ended_at = from_db_time(last_sample["sampled_at"]) if last_sample else from_db_time(row["started_at"]) self.finish_cycle(int(row["id"]), ended_at, incomplete_end=True) closed += 1 return closed def recent_cycles(self, limit: int = 100) -> list[dict[str, Any]]: rows = self._conn().execute( "SELECT * FROM cycles WHERE ended_at IS NOT NULL ORDER BY started_at DESC LIMIT ?", (max(1, min(limit, 1000)),), ).fetchall() return [dict(row) for row in rows] def cycles_between(self, start: datetime, end: datetime) -> list[dict[str, Any]]: rows = self._conn().execute( """ SELECT * FROM cycles WHERE started_at < ? AND COALESCE(ended_at, ?) >= ? ORDER BY started_at DESC """, (to_db_time(end), to_db_time(end), to_db_time(start)), ).fetchall() return [dict(row) for row in rows] def summary_between(self, start: datetime, end: datetime) -> dict[str, Any]: cycles = self.cycles_between(start, end) durations: list[float] = [] energy_wh = 0.0 for cycle in cycles: cycle_start = max(start, from_db_time(cycle["started_at"])) cycle_end_raw = from_db_time(cycle["ended_at"]) if cycle["ended_at"] else min(end, utc_now()) cycle_end = min(end, cycle_end_raw) durations.append(max(0.0, (cycle_end - cycle_start).total_seconds())) energy_wh += float(cycle["energy_wh"] or 0.0) return { "cycle_count": len(cycles), "runtime_seconds": sum(durations), "average_duration_seconds": sum(durations) / len(durations) if durations else 0.0, "longest_duration_seconds": max(durations, default=0.0), "energy_wh": energy_wh, } def sample_history( self, since: datetime, until: datetime | None = None, max_points: int = 800, preserve_ranges: list[tuple[datetime, datetime]] | None = None, ) -> list[dict[str, Any]]: until = until or utc_now() rows = self._conn().execute( "SELECT sampled_at, watts FROM samples WHERE sampled_at >= ? AND sampled_at <= ? ORDER BY sampled_at", (to_db_time(since), to_db_time(until)), ).fetchall() if len(rows) <= max_points: selected = list(rows) else: # First/minimum/maximum/last preserves short power spikes that a # simple stride can skip entirely. All readings inside known pump # cycles are also retained for accurate hover details. bucket_count = max(1, max_points // 4) bucket_size = max(1, math.ceil(len(rows) / bucket_count)) selected_by_time: dict[str, sqlite3.Row] = {} for offset in range(0, len(rows), bucket_size): bucket = rows[offset : offset + bucket_size] for row in ( bucket[0], min(bucket, key=lambda item: float(item["watts"])), max(bucket, key=lambda item: float(item["watts"])), bucket[-1], ): selected_by_time[row["sampled_at"]] = row ranges = preserve_ranges or [] if ranges: normalized = [(start.astimezone(UTC), end.astimezone(UTC)) for start, end in ranges] for row in rows: sampled_at = from_db_time(row["sampled_at"]) if any(start <= sampled_at <= end for start, end in normalized): selected_by_time[row["sampled_at"]] = row selected = sorted(selected_by_time.values(), key=lambda row: row["sampled_at"]) return [{"sampled_at": row["sampled_at"], "watts": row["watts"]} for row in selected] def cleanup_samples(self, retention_days: int) -> int: cutoff = to_db_time(utc_now() - timedelta(days=retention_days)) db = self._conn() cursor = db.execute("DELETE FROM samples WHERE sampled_at < ?", (cutoff,)) db.commit() return int(cursor.rowcount)