src/database/crud.py

"""
Database CRUD operations for the Climate Risk Index Engine.

Uses psycopg2 (sync) for PostgreSQL when DATABASE_URL is set, with an in-memory
fallback for demo/testing.  Matches Weather AI 2's PgConnection pattern exactly.
"""

from __future__ import annotations

import json
import logging
import os
import uuid
from collections import defaultdict
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from src.database.schema import get_full_ddl, get_table_names

log = logging.getLogger(__name__)


def _to_date(val):
    """Convert a string or date to a datetime.date for psycopg2 DATE columns."""
    if val is None:
        return None
    if isinstance(val, datetime):
        return val.date()
    if hasattr(val, "date") and callable(getattr(val, "date", None)):
        return val.date()
    if hasattr(val, "year") and hasattr(val, "month") and hasattr(val, "day"):
        return val  # already a date
    if isinstance(val, str):
        return datetime.fromisoformat(val.replace("Z", "+00:00")).date() if "T" in val else datetime.strptime(val, "%Y-%m-%d").date()
    return val


def _to_datetime(val) -> Optional[datetime]:
    """Convert a string to a datetime for psycopg2 TIMESTAMPTZ columns."""
    if val is None:
        return None
    if isinstance(val, datetime):
        return val
    if isinstance(val, str):
        try:
            return datetime.fromisoformat(val.replace("Z", "+00:00"))
        except ValueError:
            return datetime.strptime(val, "%Y-%m-%d")
    return val


# ── Connection wrapper (Weather AI 2 pattern) ──────────────────────────

_pool = None


def _get_pool(dsn: str):
    """Get or create a SimpleConnectionPool for the given DSN."""
    global _pool
    from psycopg2.pool import SimpleConnectionPool
    if _pool is None or _pool.closed:
        _pool = SimpleConnectionPool(minconn=2, maxconn=10, dsn=dsn)
    return _pool


class PgConnection:
    """Thin psycopg2 wrapper matching Weather AI 2's PgConnection exactly.

    Uses SimpleConnectionPool. All methods are sync.
    conn.execute(sql, params) returns a cursor.
    """

    def __init__(self, dsn: str):
        pool = _get_pool(dsn)
        self._conn = pool.getconn()
        self._conn.autocommit = True
        self._pool = pool
        self._last_cur = None
        self._refresh_conn()

    def _refresh_conn(self):
        """Handle Neon cold starts — test with SELECT 1, reconnect if stale.

        Retries up to 3 times. SimpleConnectionPool initializes minconn=2
        fresh connections at startup, which typically go stale together
        after Neon's ~5-min idle timeout; a single retry can hand back a
        second dead conn.
        """
        try:
            cur = self._conn.cursor()
            cur.execute("SELECT 1")
            cur.close()
            return
        except Exception:
            pass
        log.info("Connection stale, reconnecting...")
        last_exc = None
        for _ in range(3):
            try:
                self._pool.putconn(self._conn, close=True)
            except Exception:
                pass
            try:
                self._conn = self._pool.getconn()
                self._conn.autocommit = True
                cur = self._conn.cursor()
                cur.execute("SELECT 1")
                cur.close()
                return
            except Exception as exc:
                last_exc = exc
        raise RuntimeError(
            f"Could not acquire a healthy DB connection after 3 attempts: {last_exc}"
        )

    def execute(self, sql: str, params=None):
        """Execute SQL with %s placeholders. Returns cursor.

        To avoid leaking server-side cursor resources over long-lived
        connections (e.g. during a pipeline run with hundreds of inserts),
        we close the previously-returned cursor before creating a new one.
        Callers that need to keep two cursors open simultaneously must use
        ``self._conn.cursor()`` directly.
        """
        if self._last_cur is not None:
            try:
                self._last_cur.close()
            except Exception:
                pass
            self._last_cur = None
        cur = self._conn.cursor()
        cur.execute(sql, params)
        self._last_cur = cur
        return cur

    def close(self):
        """Return connection to pool."""
        if self._last_cur is not None:
            try:
                self._last_cur.close()
            except Exception:
                pass
            self._last_cur = None
        if self._pool and self._conn:
            self._pool.putconn(self._conn)
            self._conn = None

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


def init_db(database_url: str | None = None):
    """Initialize database connection. Returns PgConnection or None.

    If no DATABASE_URL, returns None (caller should use InMemoryStore).
    Creates schema tables on first connect.
    """
    url = database_url or os.environ.get("DATABASE_URL", "")
    if not url:
        log.info("No DATABASE_URL set, DB disabled")
        return None
    try:
        conn = PgConnection(url)
        # Create schema
        conn.execute(get_full_ddl())
        log.info("Database connected, schema initialized (%d tables)", len(get_table_names()))
        return conn
    except Exception as exc:
        log.warning("Database connection failed: %s", exc)
        return None


# ── In-memory fallback ───────────────────────────────────────────────────

class InMemoryStore:
    """
    Simple in-memory storage for demo mode.
    Stores rows as dicts keyed by table name.
    """

    def __init__(self):
        self.tables: Dict[str, list[dict]] = defaultdict(list)
        self._id_counters: Dict[str, int] = defaultdict(int)

    def insert(self, table: str, row: dict) -> int:
        """Insert a row, returning a synthetic ID."""
        self._id_counters[table] += 1
        row_copy = dict(row)
        row_copy["id"] = self._id_counters[table]
        if "created_at" not in row_copy:
            row_copy["created_at"] = datetime.now(timezone.utc).isoformat()
        self.tables[table].append(row_copy)
        return row_copy["id"]

    def query(
        self, table: str, filters: Optional[Dict[str, Any]] = None, limit: int = 100
    ) -> list[dict]:
        """Query rows with optional simple equality filters."""
        rows = self.tables.get(table, [])
        if filters:
            rows = [
                r for r in rows
                if all(r.get(k) == v for k, v in filters.items())
            ]
        return rows[:limit]

    def count(self, table: str) -> int:
        return len(self.tables.get(table, []))

    # ── Convenience helpers (mirror CRUD API for easy testing) ──

    def insert_zone(self, zone_id: str, data: dict) -> int:
        """Insert or overwrite a zone."""
        row = dict(data)
        row["zone_id"] = zone_id
        existing = [r for r in self.tables["zones"] if r["zone_id"] == zone_id]
        if existing:
            existing[0].update(row)
            return existing[0].get("id", 0)
        return self.insert("zones", row)

    def get_zone(self, zone_id: str) -> Optional[dict]:
        rows = self.query("zones", {"zone_id": zone_id}, limit=1)
        return rows[0] if rows else None

    def insert_heat_index(self, zone_id: str, date: str, data: dict) -> int:
        row = dict(data)
        row["zone_id"] = zone_id
        row["date"] = date
        return self.insert("heat_indices", row)

    def insert_prediction(self, zone_id: str, date: str, data: dict) -> int:
        row = dict(data)
        row["zone_id"] = zone_id
        row["date"] = date
        return self.insert("predictions", row)

    def get_recent_heat_indices(self, zone_id: str, limit: int = 90) -> list[dict]:
        rows = self.query("heat_indices", {"zone_id": zone_id}, limit=limit)
        return sorted(rows, key=lambda r: r.get("date", ""), reverse=True)

    def get_recent_predictions(self, zone_id: str, limit: int = 30) -> list[dict]:
        rows = self.query("predictions", {"zone_id": zone_id}, limit=limit)
        return sorted(rows, key=lambda r: r.get("date", ""), reverse=True)


# ── CRUD functions (all sync) ──────────────────────────────────────────

# --- Zones ---

def upsert_zone(conn, zone_data: dict) -> None:
    """Insert or update a zone."""
    conn.execute(
        """
        INSERT INTO zones (zone_id, name, city, country, latitude, longitude,
                           elevation_m, area_km2, population_est, settlement_type,
                           worker_population_est, outdoor_exposure_pct,
                           heat_vulnerability, hot_months, notes)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON CONFLICT (zone_id) DO UPDATE SET
            name = EXCLUDED.name,
            population_est = EXCLUDED.population_est,
            worker_population_est = EXCLUDED.worker_population_est,
            outdoor_exposure_pct = EXCLUDED.outdoor_exposure_pct,
            heat_vulnerability = EXCLUDED.heat_vulnerability,
            hot_months = EXCLUDED.hot_months,
            notes = EXCLUDED.notes
        """,
        (zone_data["zone_id"], zone_data["name"], zone_data["city"],
         zone_data["country"], zone_data["latitude"], zone_data["longitude"],
         zone_data.get("elevation_m"), zone_data.get("area_km2"),
         zone_data.get("population_est"), zone_data["settlement_type"],
         zone_data.get("worker_population_est"),
         zone_data.get("outdoor_exposure_pct"),
         zone_data["heat_vulnerability"],
         zone_data.get("hot_months", []),
         zone_data.get("notes", "")),
    )


def get_zone(conn, zone_id: str) -> Optional[dict]:
    """Fetch a single zone."""
    cur = conn.execute("SELECT * FROM zones WHERE zone_id = %s", (zone_id,))
    cols = [d[0] for d in cur.description] if cur.description else []
    row = cur.fetchone()
    return dict(zip(cols, row)) if row else None


def get_all_zones(conn) -> list[dict]:
    """Fetch all zones."""
    cur = conn.execute("SELECT * FROM zones ORDER BY city, name")
    cols = [d[0] for d in cur.description] if cur.description else []
    rows = cur.fetchall()
    return [dict(zip(cols, r)) for r in rows]


# --- Daily readings ---

def insert_daily_reading(conn, reading: dict) -> Optional[int]:
    """Insert a daily reading. Returns the row ID."""
    cur = conn.execute(
        """
        INSERT INTO daily_readings (zone_id, date, temp_mean_c, temp_max_c,
            temp_min_c, humidity_pct, wind_speed_ms, solar_rad_wm2,
            precip_mm, source, data_quality)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON CONFLICT (zone_id, date) DO UPDATE SET
            temp_mean_c = COALESCE(EXCLUDED.temp_mean_c, daily_readings.temp_mean_c),
            temp_max_c = COALESCE(EXCLUDED.temp_max_c, daily_readings.temp_max_c),
            temp_min_c = COALESCE(EXCLUDED.temp_min_c, daily_readings.temp_min_c),
            humidity_pct = COALESCE(EXCLUDED.humidity_pct, daily_readings.humidity_pct),
            wind_speed_ms = COALESCE(EXCLUDED.wind_speed_ms, daily_readings.wind_speed_ms),
            solar_rad_wm2 = COALESCE(EXCLUDED.solar_rad_wm2, daily_readings.solar_rad_wm2),
            precip_mm = COALESCE(EXCLUDED.precip_mm, daily_readings.precip_mm),
            source = EXCLUDED.source,
            data_quality = GREATEST(EXCLUDED.data_quality, daily_readings.data_quality)
        RETURNING id
        """,
        (reading["zone_id"], _to_date(reading["date"]),
         reading.get("temp_mean_c"), reading.get("temp_max_c"),
         reading.get("temp_min_c"), reading.get("humidity_pct"),
         reading.get("wind_speed_ms"), reading.get("solar_rad_wm2"),
         reading.get("precip_mm"),
         reading.get("source", "unknown"),
         reading.get("data_quality", 0.0)),
    )
    row = cur.fetchone()
    return row[0] if row else None


def insert_daily_readings_batch(conn, readings: list[dict]) -> int:
    """Bulk-insert daily readings with a single round-trip per chunk.

    Uses psycopg2's executemany (single statement, many rows). Much faster
    than insert_daily_reading() in a loop during pipeline runs. Does not
    return row IDs — use insert_daily_reading() if you need the ID.
    """
    if not readings:
        return 0
    rows = [
        (r["zone_id"], _to_date(r["date"]),
         r.get("temp_mean_c"), r.get("temp_max_c"),
         r.get("temp_min_c"), r.get("humidity_pct"),
         r.get("wind_speed_ms"), r.get("solar_rad_wm2"),
         r.get("precip_mm"),
         r.get("source", "unknown"),
         r.get("data_quality", 0.0))
        for r in readings
    ]
    cur = conn._conn.cursor()
    try:
        cur.executemany(
            """
            INSERT INTO daily_readings (zone_id, date, temp_mean_c, temp_max_c,
                temp_min_c, humidity_pct, wind_speed_ms, solar_rad_wm2,
                precip_mm, source, data_quality)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT (zone_id, date) DO UPDATE SET
                temp_mean_c = COALESCE(EXCLUDED.temp_mean_c, daily_readings.temp_mean_c),
                temp_max_c = COALESCE(EXCLUDED.temp_max_c, daily_readings.temp_max_c),
                temp_min_c = COALESCE(EXCLUDED.temp_min_c, daily_readings.temp_min_c),
                humidity_pct = COALESCE(EXCLUDED.humidity_pct, daily_readings.humidity_pct),
                wind_speed_ms = COALESCE(EXCLUDED.wind_speed_ms, daily_readings.wind_speed_ms),
                solar_rad_wm2 = COALESCE(EXCLUDED.solar_rad_wm2, daily_readings.solar_rad_wm2),
                precip_mm = COALESCE(EXCLUDED.precip_mm, daily_readings.precip_mm),
                source = EXCLUDED.source,
                data_quality = GREATEST(EXCLUDED.data_quality, daily_readings.data_quality)
            """,
            rows,
        )
        return len(rows)
    finally:
        cur.close()


def get_daily_readings(conn, zone_id: str, limit: int = 90) -> list[dict]:
    """Fetch recent daily readings for a zone."""
    cur = conn.execute(
        "SELECT * FROM daily_readings WHERE zone_id = %s ORDER BY date DESC LIMIT %s",
        (zone_id, limit),
    )
    cols = [d[0] for d in cur.description] if cur.description else []
    rows = cur.fetchall()
    return [dict(zip(cols, r)) for r in rows]


# --- Healed readings ---

def insert_healed_reading(conn, reading: dict) -> Optional[int]:
    """Insert a healed reading."""
    cur = conn.execute(
        """
        INSERT INTO healed_readings (zone_id, date, raw_reading_id,
            temp_mean_c, temp_max_c, temp_min_c,
            humidity_pct, wind_speed_ms, quality_score, heal_action, fields_corrected)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON CONFLICT (zone_id, date) DO UPDATE SET
            temp_mean_c = EXCLUDED.temp_mean_c,
            temp_max_c = EXCLUDED.temp_max_c,
            temp_min_c = EXCLUDED.temp_min_c,
            humidity_pct = EXCLUDED.humidity_pct,
            wind_speed_ms = EXCLUDED.wind_speed_ms,
            quality_score = EXCLUDED.quality_score,
            heal_action = EXCLUDED.heal_action,
            fields_corrected = EXCLUDED.fields_corrected,
            healed_at = NOW()
        RETURNING id
        """,
        (reading["zone_id"], _to_date(reading["date"]), reading.get("raw_reading_id"),
         reading.get("temp_mean_c"),
         reading.get("temp_max_c"), reading.get("temp_min_c"),
         reading.get("humidity_pct"), reading.get("wind_speed_ms"),
         reading.get("quality_score", 0.0),
         reading.get("heal_action", "passthrough"),
         reading.get("fields_corrected", [])),
    )
    row = cur.fetchone()
    return row[0] if row else None


# --- Healing log ---

def insert_healing_log(conn, entry: dict) -> Optional[int]:
    """Insert a healing log entry."""
    cur = conn.execute(
        """
        INSERT INTO healing_log (zone_id, date, healed_reading_id,
            agent_type, reasoning, corrections, tools_used,
            confidence, tokens_used, latency_ms)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        RETURNING id
        """,
        (entry["zone_id"], _to_date(entry["date"]), entry.get("healed_reading_id"),
         entry.get("agent_type", "rule_based"), entry.get("reasoning"),
         json.dumps(entry.get("corrections", {})),
         entry.get("tools_used", []),
         entry.get("confidence"), entry.get("tokens_used", 0),
         entry.get("latency_ms", 0)),
    )
    row = cur.fetchone()
    return row[0] if row else None


# --- Heat indices ---

def insert_heat_index(conn, record: dict) -> Optional[int]:
    """Insert a daily heat index record."""
    cur = conn.execute(
        """
        INSERT INTO heat_indices (zone_id, date, grid_temp_c, uhi_delta_c,
            corrected_temp_c, wbgt_c, heat_index_c, heat_risk_score,
            risk_level, consecutive_hot_days)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON CONFLICT (zone_id, date) DO UPDATE SET
            grid_temp_c = EXCLUDED.grid_temp_c,
            uhi_delta_c = EXCLUDED.uhi_delta_c,
            corrected_temp_c = EXCLUDED.corrected_temp_c,
            wbgt_c = EXCLUDED.wbgt_c,
            heat_index_c = EXCLUDED.heat_index_c,
            heat_risk_score = EXCLUDED.heat_risk_score,
            risk_level = EXCLUDED.risk_level,
            consecutive_hot_days = EXCLUDED.consecutive_hot_days,
            computed_at = NOW()
        RETURNING id
        """,
        (record["zone_id"], _to_date(record["date"]),
         record.get("grid_temp_c"), record.get("uhi_delta_c"),
         record.get("corrected_temp_c"), record.get("wbgt_c"),
         record.get("heat_index_c"), record.get("heat_risk_score"),
         record.get("risk_level"), record.get("consecutive_hot_days", 0)),
    )
    row = cur.fetchone()
    return row[0] if row else None


def get_recent_heat_indices(conn, zone_id: str, limit: int = 90) -> list[dict]:
    """Fetch recent heat index records for a zone."""
    cur = conn.execute(
        "SELECT * FROM heat_indices WHERE zone_id = %s ORDER BY date DESC LIMIT %s",
        (zone_id, limit),
    )
    cols = [d[0] for d in cur.description] if cur.description else []
    rows = cur.fetchall()
    return [dict(zip(cols, r)) for r in rows]


# --- Predictions ---

def insert_prediction(conn, record: dict) -> Optional[int]:
    """Insert a daily prediction record."""
    cur = conn.execute(
        """
        INSERT INTO predictions (zone_id, date, trigger_probability_7d,
            prediction_confidence, model_tier, xgb_probability,
            lstm_probability, ensemble_method,
            annual_cost_per_worker, payout_factor, learned_frequency)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON CONFLICT (zone_id, date) DO UPDATE SET
            trigger_probability_7d = EXCLUDED.trigger_probability_7d,
            prediction_confidence = EXCLUDED.prediction_confidence,
            model_tier = EXCLUDED.model_tier,
            xgb_probability = EXCLUDED.xgb_probability,
            lstm_probability = EXCLUDED.lstm_probability,
            ensemble_method = EXCLUDED.ensemble_method,
            annual_cost_per_worker = EXCLUDED.annual_cost_per_worker,
            payout_factor = EXCLUDED.payout_factor,
            learned_frequency = EXCLUDED.learned_frequency,
            predicted_at = NOW()
        RETURNING id
        """,
        (record["zone_id"], _to_date(record["date"]),
         record.get("trigger_probability_7d"),
         record.get("prediction_confidence"),
         record.get("model_tier", "climatology"),
         record.get("xgb_probability"),
         record.get("lstm_probability"),
         record.get("ensemble_method", "average"),
         record.get("annual_cost_per_worker"),
         record.get("payout_factor"),
         record.get("learned_frequency")),
    )
    row = cur.fetchone()
    return row[0] if row else None


def get_recent_predictions(conn, zone_id: str, limit: int = 30) -> list[dict]:
    """Fetch recent prediction records for a zone."""
    cur = conn.execute(
        "SELECT * FROM predictions WHERE zone_id = %s ORDER BY date DESC LIMIT %s",
        (zone_id, limit),
    )
    cols = [d[0] for d in cur.description] if cur.description else []
    rows = cur.fetchall()
    return [dict(zip(cols, r)) for r in rows]


# --- Trigger events ---

def insert_trigger_event(conn, event: dict) -> Optional[int]:
    """Insert a trigger event."""
    cur = conn.execute(
        """
        INSERT INTO trigger_events (zone_id, trigger_level, triggered_at,
            max_temp_c, max_wbgt_c, consecutive_days, heat_risk_score,
            settlement_type, payout_per_worker_usd, enrolled_workers,
            total_payout_usd)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        RETURNING id
        """,
        (event["zone_id"], event["trigger_level"], _to_datetime(event["triggered_at"]),
         event.get("max_temp_c"), event.get("max_wbgt_c"),
         event.get("consecutive_days"), event.get("heat_risk_score"),
         event.get("settlement_type"), event.get("payout_per_worker_usd"),
         event.get("enrolled_workers"), event.get("total_payout_usd")),
    )
    row = cur.fetchone()
    return row[0] if row else None


def get_trigger_events(conn, zone_id: Optional[str] = None, limit: int = 50) -> list[dict]:
    """Fetch trigger events, optionally filtered by zone."""
    if zone_id:
        cur = conn.execute(
            "SELECT * FROM trigger_events WHERE zone_id = %s ORDER BY triggered_at DESC LIMIT %s",
            (zone_id, limit),
        )
    else:
        cur = conn.execute(
            "SELECT * FROM trigger_events ORDER BY triggered_at DESC LIMIT %s",
            (limit,),
        )
    cols = [d[0] for d in cur.description] if cur.description else []
    rows = cur.fetchall()
    return [dict(zip(cols, r)) for r in rows]


# --- Basis risk ---

def insert_basis_risk(conn, report: dict) -> Optional[int]:
    """Insert a basis risk assessment."""
    cur = conn.execute(
        """
        INSERT INTO basis_risk (zone_id, overall_score, false_positive_rate,
            false_negative_rate, correlation, mae, total_events,
            true_positives, true_negatives, false_positives, false_negatives,
            trigger_accuracy, tier_accuracy, recommendations,
            confidence_low, confidence_high)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        RETURNING id
        """,
        (report["zone_id"], report["overall_score"],
         report.get("false_positive_rate"), report.get("false_negative_rate"),
         report.get("correlation"), report.get("mae"),
         report.get("total_events"), report.get("true_positives"),
         report.get("true_negatives"), report.get("false_positives"),
         report.get("false_negatives"), report.get("trigger_accuracy"),
         json.dumps(report.get("tier_accuracy", {})),
         report.get("recommendations", []),
         report.get("confidence_low"), report.get("confidence_high")),
    )
    row = cur.fetchone()
    return row[0] if row else None


# --- Explanations ---

def insert_explanation(conn, explanation: dict) -> Optional[int]:
    """Insert a generated explanation."""
    cur = conn.execute(
        """
        INSERT INTO explanations (trigger_event_id, zone_id, trigger_level,
            english_text, swahili_text, payout_amount, payout_currency,
            settlement_type, protective_actions, provider)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        RETURNING id
        """,
        (explanation.get("trigger_event_id"), explanation["zone_id"],
         explanation["trigger_level"], explanation["english_text"],
         explanation["swahili_text"], explanation.get("payout_amount"),
         explanation.get("payout_currency", "KES"),
         explanation.get("settlement_type"),
         explanation.get("protective_actions", []),
         explanation.get("provider", "template")),
    )
    row = cur.fetchone()
    return row[0] if row else None


def get_explanations(conn, zone_id: Optional[str] = None, limit: int = 20) -> list[dict]:
    """Fetch explanations, optionally by zone."""
    if zone_id:
        cur = conn.execute(
            "SELECT * FROM explanations WHERE zone_id = %s ORDER BY generated_at DESC LIMIT %s",
            (zone_id, limit),
        )
    else:
        cur = conn.execute(
            "SELECT * FROM explanations ORDER BY generated_at DESC LIMIT %s",
            (limit,),
        )
    cols = [d[0] for d in cur.description] if cur.description else []
    rows = cur.fetchall()
    return [dict(zip(cols, r)) for r in rows]


# --- Notifications ---

def insert_notification(conn, notif: dict) -> Optional[int]:
    """Insert a notification delivery record."""
    cur = conn.execute(
        """
        INSERT INTO notifications (explanation_id, zone_id, recipient, channel,
            status, message_preview, message_sid, cost_estimate, error)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        RETURNING id
        """,
        (notif.get("explanation_id"), notif["zone_id"],
         notif["recipient"], notif["channel"],
         notif["status"], notif.get("message_preview"),
         notif.get("message_sid"), notif.get("cost_estimate", 0.0),
         notif.get("error")),
    )
    row = cur.fetchone()
    return row[0] if row else None


def get_notifications(conn, zone_id: Optional[str] = None, limit: int = 50) -> list[dict]:
    """Fetch notification records."""
    if zone_id:
        cur = conn.execute(
            "SELECT * FROM notifications WHERE zone_id = %s ORDER BY sent_at DESC LIMIT %s",
            (zone_id, limit),
        )
    else:
        cur = conn.execute(
            "SELECT * FROM notifications ORDER BY sent_at DESC LIMIT %s",
            (limit,),
        )
    cols = [d[0] for d in cur.description] if cur.description else []
    rows = cur.fetchall()
    return [dict(zip(cols, r)) for r in rows]


# --- Pipeline runs ---

def start_pipeline_run(conn, run_id: Optional[str] = None) -> str:
    """Record the start of a pipeline run. Returns run_id."""
    rid = run_id or f"run-{uuid.uuid4().hex[:12]}"
    now = datetime.now(timezone.utc)
    conn.execute(
        """
        INSERT INTO pipeline_runs (run_id, started_at, status)
        VALUES (%s, %s, 'running')
        """,
        (rid, now),
    )
    return rid


def finish_pipeline_run(
    conn,
    run_id: str,
    status: str = "completed",
    steps_completed: Optional[list[str]] = None,
    step_status: Optional[dict] = None,
    error: Optional[str] = None,
    zones_processed: int = 0,
    total_cost_usd: float = 0.0,
) -> None:
    """Record the completion of a pipeline run."""
    now = datetime.now(timezone.utc)
    conn.execute(
        """
        UPDATE pipeline_runs
        SET finished_at = %s, status = %s, steps_completed = %s,
            step_status = %s, error = %s, zones_processed = %s,
            total_cost_usd = %s
        WHERE run_id = %s
        """,
        (now, status,
         steps_completed or [],
         json.dumps(step_status or {}),
         error, zones_processed,
         total_cost_usd,
         run_id),
    )


def get_recent_runs(conn, limit: int = 10) -> list[dict]:
    """Fetch recent pipeline runs."""
    cur = conn.execute(
        "SELECT * FROM pipeline_runs ORDER BY started_at DESC LIMIT %s",
        (limit,),
    )
    cols = [d[0] for d in cur.description] if cur.description else []
    rows = cur.fetchall()
    return [dict(zip(cols, r)) for r in rows]