db.py

"""
SQLite layer for hive registry + inspection history.

Storage is ephemeral on HF Spaces' free tier (container restart wipes
the filesystem), that's fine for a hackathon demo. For real-world
deployment we'd back this with HF Hub persistence or a real DB.
"""

from __future__ import annotations

import json
import sqlite3
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator

DB_PATH = Path(__file__).parent / "apiarist.sqlite"


SCHEMA = """
CREATE TABLE IF NOT EXISTS hives (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL UNIQUE,
    location TEXT,
    queen_marker TEXT,
    notes TEXT,
    created_at REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS inspections (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    hive_id INTEGER NOT NULL,
    queen_detected INTEGER,
    varroa_mites_visible INTEGER,
    swarm_cells_detected INTEGER,
    brood_pattern TEXT,
    frame_health TEXT,
    notes TEXT,
    raw_response TEXT,
    structured_json TEXT,
    created_at REAL NOT NULL,
    FOREIGN KEY (hive_id) REFERENCES hives(id) ON DELETE CASCADE
);

CREATE INDEX IF NOT EXISTS idx_inspections_hive ON inspections(hive_id, created_at);
"""


@contextmanager
def conn() -> Iterator[sqlite3.Connection]:
    c = sqlite3.connect(DB_PATH)
    c.row_factory = sqlite3.Row
    c.execute("PRAGMA foreign_keys = ON")
    try:
        yield c
        c.commit()
    finally:
        c.close()


def init_db() -> None:
    with conn() as c:
        c.executescript(SCHEMA)


def add_hive(name: str, location: str = "", queen_marker: str = "", notes: str = "") -> int:
    with conn() as c:
        cur = c.execute(
            "INSERT INTO hives (name, location, queen_marker, notes, created_at) "
            "VALUES (?, ?, ?, ?, ?)",
            (name, location, queen_marker, notes, time.time()),
        )
        return cur.lastrowid or 0


def get_or_create_hive(name: str) -> int:
    """Look up a hive by name; create with defaults if missing."""
    name = name.strip()
    if not name:
        name = "Unnamed Hive"
    with conn() as c:
        row = c.execute("SELECT id FROM hives WHERE name = ?", (name,)).fetchone()
        if row:
            return row["id"]
        cur = c.execute(
            "INSERT INTO hives (name, created_at) VALUES (?, ?)",
            (name, time.time()),
        )
        return cur.lastrowid or 0


def list_hives() -> list[dict]:
    with conn() as c:
        rows = c.execute(
            """
            SELECT h.id, h.name, h.location, h.queen_marker, h.notes, h.created_at,
                   COUNT(i.id) AS inspection_count,
                   MAX(i.created_at) AS last_inspected
            FROM hives h
            LEFT JOIN inspections i ON i.hive_id = h.id
            GROUP BY h.id
            ORDER BY h.name
            """
        ).fetchall()
    return [dict(r) for r in rows]


def delete_hive(hive_id: int) -> None:
    with conn() as c:
        c.execute("DELETE FROM hives WHERE id = ?", (hive_id,))


def add_inspection(hive_id: int, results: dict, raw_response: str = "") -> int:
    with conn() as c:
        cur = c.execute(
            """
            INSERT INTO inspections (
                hive_id, queen_detected, varroa_mites_visible,
                swarm_cells_detected, brood_pattern, frame_health,
                notes, raw_response, structured_json, created_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                hive_id,
                int(bool(results.get("queen_detected"))),
                int(results.get("varroa_mites_visible", 0) or 0),
                int(bool(results.get("swarm_cells_detected"))),
                results.get("brood_pattern", ""),
                results.get("frame_health", ""),
                results.get("notes", ""),
                raw_response,
                json.dumps(results),
                time.time(),
            ),
        )
        return cur.lastrowid or 0


def get_inspections_for_hive(hive_id: int, limit: int = 50) -> list[dict]:
    with conn() as c:
        rows = c.execute(
            """
            SELECT id, queen_detected, varroa_mites_visible, swarm_cells_detected,
                   brood_pattern, frame_health, notes, created_at, structured_json
            FROM inspections
            WHERE hive_id = ?
            ORDER BY created_at DESC
            LIMIT ?
            """,
            (hive_id, limit),
        ).fetchall()
    return [dict(r) for r in rows]


def hive_stats() -> dict:
    """High-level apiary summary for the dashboard."""
    with conn() as c:
        total_hives = c.execute("SELECT COUNT(*) AS n FROM hives").fetchone()["n"]
        total_inspections = c.execute(
            "SELECT COUNT(*) AS n FROM inspections"
        ).fetchone()["n"]
        recent = c.execute(
            """
            SELECT h.name, i.frame_health, i.varroa_mites_visible, i.created_at
            FROM inspections i
            JOIN hives h ON h.id = i.hive_id
            ORDER BY i.created_at DESC
            LIMIT 5
            """
        ).fetchall()
    return {
        "total_hives": total_hives,
        "total_inspections": total_inspections,
        "recent": [dict(r) for r in recent],
    }