db.py

import json
import os
import random
import re
import sqlite3
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

try:
    import yaml as _yaml
    _YAML_AVAILABLE = True
except ImportError:
    _YAML_AVAILABLE = False

DB_PATH = Path(os.getenv("DISPATCH_DB_PATH", str(Path(__file__).parent / "search_history.db")))

_local = threading.local()


def _get_conn() -> sqlite3.Connection:
    if not getattr(_local, "conn", None):
        _local.conn = sqlite3.connect(DB_PATH)
        _local.conn.row_factory = sqlite3.Row
        _local.conn.execute("PRAGMA journal_mode=WAL")
        _local.conn.execute("PRAGMA foreign_keys=ON")
        _init_schema(_local.conn)
    return _local.conn


def _init_schema(conn: sqlite3.Connection) -> None:
    # Migration: rename crons → dispatches (runs once on old DBs, no-op on new)
    try:
        conn.execute("ALTER TABLE crons RENAME TO dispatches")
        conn.commit()
    except sqlite3.OperationalError:
        pass

    conn.executescript("""
        CREATE TABLE IF NOT EXISTS sessions (
            id         INTEGER PRIMARY KEY AUTOINCREMENT,
            model      TEXT    NOT NULL,
            agent      TEXT,
            source     TEXT,
            started_at TEXT    NOT NULL
        );

        CREATE TABLE IF NOT EXISTS messages (
            id         INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id INTEGER NOT NULL REFERENCES sessions(id),
            role       TEXT    NOT NULL,
            content    TEXT,
            created_at TEXT    NOT NULL
        );

        CREATE TABLE IF NOT EXISTS tool_calls (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id  INTEGER NOT NULL REFERENCES sessions(id),
            tool_name   TEXT    NOT NULL,
            args        TEXT    NOT NULL,
            result      TEXT,
            duration_ms INTEGER,
            created_at  TEXT    NOT NULL
        );

        CREATE TABLE IF NOT EXISTS dispatches (
            id             INTEGER PRIMARY KEY AUTOINCREMENT,
            agent          TEXT    NOT NULL,
            prompt         TEXT    NOT NULL,
            schedule       TEXT    NOT NULL,
            label          TEXT    NOT NULL,
            enabled        INTEGER NOT NULL DEFAULT 1,
            run_on_create  INTEGER NOT NULL DEFAULT 0,
            created_at     TEXT    NOT NULL
        );

        CREATE TABLE IF NOT EXISTS run_events (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id  INTEGER NOT NULL,
            seq         INTEGER NOT NULL,
            t_ms        INTEGER,
            duration_ms INTEGER,
            kind        TEXT    NOT NULL,
            phase       TEXT,
            sub         TEXT,
            name        TEXT,
            ok          INTEGER,
            error       TEXT,
            data        TEXT,
            created_at  TEXT    NOT NULL
        );

        CREATE TABLE IF NOT EXISTS run_payloads (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id  INTEGER NOT NULL,
            seq         INTEGER NOT NULL,
            input       TEXT,
            output      TEXT,
            created_at  TEXT    NOT NULL
        );

        CREATE TABLE IF NOT EXISTS settings (
            key        TEXT PRIMARY KEY,
            value      TEXT,
            updated_at TEXT
        );

        CREATE TABLE IF NOT EXISTS agents (
            id                           INTEGER PRIMARY KEY AUTOINCREMENT,
            name                         TEXT    NOT NULL UNIQUE,
            label                        TEXT    NOT NULL DEFAULT '',
            description                  TEXT    DEFAULT '',
            tier                         INTEGER DEFAULT 1,
            enabled                      INTEGER DEFAULT 1,
            goal_template                TEXT    NOT NULL DEFAULT '',
            output_template              TEXT    DEFAULT '',
            system_prompt                TEXT    DEFAULT '',
            sources                      TEXT    DEFAULT '[]',
            sub_queries                  TEXT    DEFAULT NULL,
            tool_bindings                TEXT    DEFAULT '["web_search"]',
            direct_sources               TEXT    DEFAULT NULL,
            freshness                    TEXT    DEFAULT NULL,
            planner_query_count          INTEGER DEFAULT 4,
            max_validator_loops          INTEGER DEFAULT 1,
            validate_links               INTEGER DEFAULT 0,
            fetch_content                INTEGER DEFAULT 0,
            always_research              INTEGER DEFAULT 1,
            fast_model                   TEXT    DEFAULT '',
            recall_agent                 TEXT    DEFAULT '',
            synthesizer_snippet_chars    INTEGER DEFAULT -1,
            vault_agent                  INTEGER DEFAULT 0,
            query_builder_plugin         TEXT    DEFAULT NULL,
            query_decompose_prompt       TEXT    DEFAULT NULL,
            post_search_extractor_plugin TEXT    DEFAULT NULL,
            is_builtin                   INTEGER DEFAULT 0,
            created_at                   TEXT    DEFAULT (datetime('now')),
            updated_at                   TEXT    DEFAULT (datetime('now'))
        );
    """)
    conn.commit()
    # Migrate columns that may be absent in older databases before creating indexes
    for col, typedef in [
        ("color", "INTEGER"),
        ("tags", "TEXT"),  # JSON list of lowercase tags
        ("recall_runs", "INTEGER"),  # previous digests shown to the synthesizer; NULL = default (2), 0 = off
    ]:
        try:
            conn.execute(f"ALTER TABLE agents ADD COLUMN {col} {typedef}")
            conn.commit()
        except sqlite3.OperationalError:
            pass  # column already exists
    for col, typedef in [("agent", "TEXT"), ("source", "TEXT"), ("ended_at", "TEXT"), ("viewed_at", "TEXT"), ("logs", "TEXT"), ("cron_id", "INTEGER"), ("dispatch_id", "INTEGER"), ("metrics", "TEXT"),
                        ("title", "TEXT"), ("focus_session_id", "INTEGER")]:  # title/focus: digest chat threads
        try:
            conn.execute(f"ALTER TABLE sessions ADD COLUMN {col} {typedef}")
            conn.commit()
        except sqlite3.OperationalError:
            pass  # column already exists
    # Backfill dispatch_id from cron_id for sessions created before the rename
    try:
        conn.execute("UPDATE sessions SET dispatch_id = cron_id WHERE cron_id IS NOT NULL AND dispatch_id IS NULL")
        conn.commit()
    except sqlite3.OperationalError:
        pass
    for col, typedef in [
        ("trigger_type", "TEXT NOT NULL DEFAULT 'cron'"),
        ("repeat_until",  "TEXT"),
        ("max_runs",      "INTEGER"),
        ("run_count",     "INTEGER NOT NULL DEFAULT 0"),
        ("timezone",      "TEXT"),
        ("tags",          "TEXT"),  # JSON list; the agent's tags apply too
    ]:
        try:
            conn.execute(f"ALTER TABLE dispatches ADD COLUMN {col} {typedef}")
            conn.commit()
        except sqlite3.OperationalError:
            pass  # column already exists
    _localize_legacy_schedules(conn)
    _clear_auto_labels(conn)
    conn.executescript("""
        CREATE INDEX IF NOT EXISTS idx_sessions_agent        ON sessions(agent);
        CREATE INDEX IF NOT EXISTS idx_sessions_started_at   ON sessions(started_at DESC);
        CREATE INDEX IF NOT EXISTS idx_messages_session_role ON messages(session_id, role);
        CREATE INDEX IF NOT EXISTS idx_run_events_session    ON run_events(session_id, seq);
        CREATE INDEX IF NOT EXISTS idx_run_payloads_session  ON run_payloads(session_id, seq);
        CREATE INDEX IF NOT EXISTS idx_run_payloads_created  ON run_payloads(created_at);
    """)
    conn.commit()


def local_timezone() -> str:
    """IANA name for new schedules when the client sends none: DISPATCH_TZ, else the system zone, else UTC."""
    tz = os.getenv("DISPATCH_TZ")
    if tz:
        return tz
    path = os.path.realpath("/etc/localtime")
    return path.split("zoneinfo/", 1)[1] if "zoneinfo/" in path else "UTC"


def _localize_legacy_schedules(conn) -> None:
    """Schedules saved before the timezone column were UTC cron. Rewrite the simple ones
    (fixed minute/hour, any weekday set) into local wall-clock time; mark the rest UTC."""
    from zoneinfo import ZoneInfo
    tz = local_timezone()
    offset = datetime.now(ZoneInfo(tz)).utcoffset()
    offset_min = int(offset.total_seconds() // 60) if offset else 0
    rows = conn.execute("SELECT id, schedule, trigger_type FROM dispatches WHERE timezone IS NULL").fetchall()
    for row in rows:
        new_schedule, new_tz = row["schedule"], "UTC"
        parts = row["schedule"].split()
        if row["trigger_type"] != "once" and len(parts) == 5 and tz != "UTC":
            m, h, dom, mon, dow = parts
            if m.isdigit() and h.isdigit() and dom == "*" and mon == "*":
                total = int(h) * 60 + int(m) + offset_min
                day_shift, minute = divmod(total, 24 * 60)
                if dow == "*":
                    days = "*"
                else:
                    try:
                        utc_days = set()
                        for chunk in dow.split(","):
                            a, _, b = chunk.partition("-")
                            utc_days.update(range(int(a), int(b or a) + 1))
                        days = ",".join(str(d) for d in sorted({(d % 7 + day_shift) % 7 for d in utc_days}))
                    except ValueError:
                        days = None
                if days is not None:
                    new_schedule, new_tz = f"{minute % 60} {minute // 60} * * {days}", tz
        conn.execute("UPDATE dispatches SET schedule=?, timezone=? WHERE id=?", (new_schedule, new_tz, row["id"]))
    conn.commit()


_AUTO_LABEL = re.compile(r"^(Daily at|(Sun|Mon|Tue|Wed|Thu|Fri|Sat)s? at|(Sun|Mon|Tue|Wed|Thu|Fri|Sat)s?[/–])")


def _clear_auto_labels(conn) -> None:
    """The UI used to save the schedule text as the label when none was given; clear those so it shows the live schedule."""
    for row in conn.execute("SELECT id, label FROM dispatches").fetchall():
        if _AUTO_LABEL.match(row["label"] or ""):
            conn.execute("UPDATE dispatches SET label='' WHERE id=?", (row["id"],))
    conn.commit()


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def start_session(model: str, agent: str = "general", source: str = "cli", viewed: bool = False, dispatch_id: Optional[int] = None) -> int:
    conn = _get_conn()
    viewed_at = _now() if viewed else None
    cur = conn.execute(
        "INSERT INTO sessions (model, agent, source, started_at, viewed_at, dispatch_id) VALUES (?, ?, ?, ?, ?, ?)",
        (model, agent, source, _now(), viewed_at, dispatch_id),
    )
    conn.commit()
    return cur.lastrowid


def mark_session_viewed(session_id: int) -> None:
    conn = _get_conn()
    conn.execute(
        "UPDATE sessions SET viewed_at = ? WHERE id = ? AND viewed_at IS NULL",
        (_now(), session_id),
    )
    conn.commit()


def end_session(session_id: int) -> None:
    conn = _get_conn()
    conn.execute("UPDATE sessions SET ended_at = ? WHERE id = ?", (_now(), session_id))
    conn.commit()


def log_message(session_id: int, role: str, content: str) -> int:
    conn = _get_conn()
    cur = conn.execute(
        "INSERT INTO messages (session_id, role, content, created_at) VALUES (?, ?, ?, ?)",
        (session_id, role, content, _now()),
    )
    conn.commit()
    return cur.lastrowid


def log_tool_call(
    session_id: int,
    tool_name: str,
    args: dict,
    result: dict,
    duration_ms: int,
) -> None:
    conn = _get_conn()
    conn.execute(
        """INSERT INTO tool_calls
               (session_id, tool_name, args, result, duration_ms, created_at)
           VALUES (?, ?, ?, ?, ?, ?)""",
        (session_id, tool_name, json.dumps(args), json.dumps(result), duration_ms, _now()),
    )
    conn.commit()


def get_sessions(limit: int = 50, agent: Optional[str] = None) -> list[dict]:
    """Runs (digest chat threads are listed by list_chats instead)."""
    conn = _get_conn()
    where = "WHERE COALESCE(s.source, '') != 'chat'" + (" AND s.agent = ?" if agent else "")
    params = (agent, limit) if agent else (limit,)
    rows = conn.execute(
        f"""
        SELECT
            s.id, s.agent, s.source, s.model, s.started_at,
            (s.viewed_at IS NOT NULL) AS viewed,
            (SELECT content FROM messages WHERE session_id = s.id AND role = 'user' ORDER BY id LIMIT 1) AS prompt,
            (SELECT COUNT(*) FROM tool_calls WHERE session_id = s.id) AS tool_count,
            CASE
                WHEN s.ended_at IS NOT NULL
                THEN CAST(ROUND((JULIANDAY(s.ended_at) - JULIANDAY(s.started_at)) * 86400000) AS INTEGER)
                ELSE (SELECT COALESCE(SUM(duration_ms), 0) FROM tool_calls WHERE session_id = s.id)
            END AS total_duration_ms,
            s.dispatch_id,
            json_extract(s.metrics, '$.status')        AS run_status,
            json_extract(s.metrics, '$.sources')       AS run_sources,
            json_extract(s.metrics, '$.timing.run_ms') AS run_ms,
            json_extract(s.metrics, '$.source')        AS metrics_source,
            json_extract(s.metrics, '$.degraded')      AS run_degraded,
            json_extract(s.metrics, '$.checks')        AS run_checks
        FROM sessions s
        {where}
        ORDER BY s.started_at DESC, s.id DESC
        LIMIT ?
        """,
        params,
    ).fetchall()
    return [dict(r) for r in rows]


def set_session_logs(session_id: int, logs: str) -> None:
    conn = _get_conn()
    conn.execute("UPDATE sessions SET logs = ? WHERE id = ?", (logs, session_id))
    conn.commit()


# ── Run metrics ───────────────────────────────────────────────────────────────

def insert_run_events(session_id: int, events: list[dict]) -> None:
    if not events:
        return
    conn = _get_conn()
    now = _now()
    conn.executemany(
        """INSERT INTO run_events (session_id, seq, t_ms, duration_ms, kind, phase, sub, name, ok, error, data, created_at)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
        [(session_id, e["seq"], e["t_ms"], e["duration_ms"], e["kind"], e["phase"], e["sub"], e["name"],
          int(bool(e["ok"])), e["error"], json.dumps(e["data"], ensure_ascii=False, default=str), now) for e in events],
    )
    conn.commit()


def insert_run_payloads(session_id: int, payloads: list[dict]) -> None:
    if not payloads:
        return
    conn = _get_conn()
    now = _now()
    conn.executemany(
        "INSERT INTO run_payloads (session_id, seq, input, output, created_at) VALUES (?, ?, ?, ?, ?)",
        [(session_id, p["seq"], p["input"], p["output"], now) for p in payloads],
    )
    conn.commit()


def purge_run_payloads(days: int) -> None:
    from datetime import timedelta
    conn = _get_conn()
    conn.execute("DELETE FROM run_payloads WHERE created_at < ?", ((datetime.now(timezone.utc) - timedelta(days=days)).isoformat(),))
    conn.commit()


def set_session_metrics(session_id: int, metrics: dict) -> None:
    conn = _get_conn()
    conn.execute("UPDATE sessions SET metrics = ? WHERE id = ?", (json.dumps(metrics, ensure_ascii=False, default=str), session_id))
    conn.commit()


def get_run_events(session_id: int) -> list[dict]:
    conn = _get_conn()
    rows = conn.execute(
        """SELECT e.seq, e.t_ms, e.duration_ms, e.kind, e.phase, e.sub, e.name, e.ok, e.error, e.data,
                  EXISTS(SELECT 1 FROM run_payloads p WHERE p.session_id = e.session_id AND p.seq = e.seq) AS has_payload
           FROM run_events e WHERE e.session_id = ? ORDER BY e.seq""",
        (session_id,),
    ).fetchall()
    return [{**dict(r), "ok": bool(r["ok"]), "has_payload": bool(r["has_payload"]), "data": json.loads(r["data"] or "{}")} for r in rows]


def get_run_payload(session_id: int, seq: int) -> Optional[dict]:
    conn = _get_conn()
    row = conn.execute(
        "SELECT seq, input, output, created_at FROM run_payloads WHERE session_id = ? AND seq = ?",
        (session_id, seq),
    ).fetchone()
    return dict(row) if row else None


def sessions_without_metrics() -> list[dict]:
    conn = _get_conn()
    rows = conn.execute("SELECT id FROM sessions WHERE metrics IS NULL AND ended_at IS NOT NULL ORDER BY id").fetchall()
    return [dict(r) for r in rows]


def record_missed_run(dispatch_id: int, agent: str, model: str, prompt: str, scheduled_for: str) -> int:
    """A run that was due while dispatch was down, kept so the run list and dashboard show the gap."""
    metrics = {
        "v": 2, "source": "live", "status": "missed",
        "run": {"trigger": "scheduled", "dispatch_id": dispatch_id, "scheduled_for": scheduled_for},
        "checks": [{"rule": "missed", "detail": "dispatch was not running when this run was due"}],
        "degraded": True,
    }
    conn = _get_conn()
    cur = conn.execute(
        "INSERT INTO sessions (model, agent, source, started_at, ended_at, viewed_at, dispatch_id, metrics) VALUES (?, ?, 'dispatch', ?, ?, ?, ?, ?)",
        (model, agent, scheduled_for, scheduled_for, _now(), dispatch_id, json.dumps(metrics)),
    )
    conn.execute("INSERT INTO messages (session_id, role, content, created_at) VALUES (?, 'user', ?, ?)", (cur.lastrowid, prompt, _now()))
    conn.commit()
    return cur.lastrowid


def recent_run_ms(dispatch_id: int, limit: int = 20, exclude_session: Optional[int] = None) -> list[int]:
    """Durations of the dispatch's latest clean runs (the baseline for the "slow" self-check)."""
    conn = _get_conn()
    rows = conn.execute(
        """SELECT json_extract(metrics, '$.timing.run_ms') AS ms FROM sessions
           WHERE dispatch_id = ? AND id != ? AND json_extract(metrics, '$.status') = 'done'
             AND json_extract(metrics, '$.timing.run_ms') IS NOT NULL
             AND COALESCE(json_extract(metrics, '$.degraded'), 0) = 0
           ORDER BY id DESC LIMIT ?""",
        (dispatch_id, exclude_session or -1, limit),
    ).fetchall()
    return [r["ms"] for r in rows]


def sessions_needing_checks() -> list[dict]:
    conn = _get_conn()
    rows = conn.execute(
        "SELECT id, metrics FROM sessions WHERE metrics IS NOT NULL AND json_extract(metrics, '$.checks') IS NULL ORDER BY id"
    ).fetchall()
    return [{"id": r["id"], "metrics": json.loads(r["metrics"])} for r in rows]


def sessions_with_metrics(since: str) -> list[dict]:
    conn = _get_conn()
    rows = conn.execute(
        "SELECT id, agent, source, dispatch_id, started_at, metrics FROM sessions WHERE metrics IS NOT NULL AND started_at >= ? ORDER BY started_at",
        (since,),
    ).fetchall()
    return [{**dict(r), "metrics": json.loads(r["metrics"])} for r in rows]


def get_session_detail(session_id: int) -> Optional[dict]:
    conn = _get_conn()
    session = conn.execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
    if not session:
        return None
    messages = conn.execute(
        "SELECT role, content, created_at FROM messages WHERE session_id = ? ORDER BY id",
        (session_id,),
    ).fetchall()
    tool_calls = conn.execute(
        "SELECT tool_name, args, result, duration_ms, created_at FROM tool_calls WHERE session_id = ? ORDER BY id",
        (session_id,),
    ).fetchall()
    return {
        **dict(session),
        "metrics": json.loads(session["metrics"]) if session["metrics"] else None,
        "messages": [dict(m) for m in messages],
        "tool_calls": [dict(t) for t in tool_calls],
    }


def get_latest_session_for_dispatch(dispatch_id: int, include_missed: bool = False) -> Optional[dict]:
    conn = _get_conn()
    missed = "" if include_missed else "AND COALESCE(json_extract(metrics, '$.status'), '') != 'missed'"
    row = conn.execute(
        f"SELECT id, agent, started_at, ended_at FROM sessions WHERE dispatch_id = ? {missed} ORDER BY started_at DESC, id DESC LIMIT 1",
        (dispatch_id,),
    ).fetchone()
    if not row:
        return None
    s = dict(row)
    msg = conn.execute(
        "SELECT content FROM messages WHERE session_id = ? AND role = 'assistant' ORDER BY id DESC LIMIT 1",
        (s["id"],),
    ).fetchone()
    s["content"] = msg["content"] if msg else None
    return s


# Runs whose output counts as a digest: done (or pre-metrics runs with no status), never missed/error/empty/cancelled.
_DONE_RUN = "COALESCE(json_extract(s.metrics, '$.status'), 'done') = 'done'"


def recent_outputs(dispatch_id: int, n: int, before_session: Optional[int] = None) -> list[dict]:
    """The last n done runs of a dispatch, newest first: {id, started_at, content}. `before_session` skips that run and later ones."""
    conn = _get_conn()
    before = "AND s.id < ?" if before_session else ""
    params = (dispatch_id, before_session, n) if before_session else (dispatch_id, n)
    rows = conn.execute(
        f"""SELECT * FROM (
                SELECT s.id, s.started_at,
                       (SELECT content FROM messages WHERE session_id = s.id AND role = 'assistant' ORDER BY id DESC LIMIT 1) AS content
                FROM sessions s
                WHERE s.dispatch_id = ? AND {_DONE_RUN} {before})
            WHERE content IS NOT NULL AND content != ''
            ORDER BY started_at DESC, id DESC LIMIT ?""",
        params,
    ).fetchall()
    return [dict(r) for r in rows]


def session_dispatch_id(session_id: int) -> Optional[int]:
    row = _get_conn().execute("SELECT dispatch_id FROM sessions WHERE id = ?", (session_id,)).fetchone()
    return row["dispatch_id"] if row else None


def get_output(session_id: int) -> Optional[str]:
    """A run's final output (its last assistant message)."""
    row = _get_conn().execute(
        "SELECT content FROM messages WHERE session_id = ? AND role = 'assistant' ORDER BY id DESC LIMIT 1", (session_id,)
    ).fetchone()
    return row["content"] if row else None


def get_phase_payloads(session_id: int, phases: tuple[str, ...]) -> list[dict]:
    """LLM payloads (input JSON, output) of a run's calls in the given phases, in call order."""
    conn = _get_conn()
    marks = ",".join("?" for _ in phases)
    rows = conn.execute(
        f"""SELECT e.phase, e.sub, p.input, p.output FROM run_events e
            JOIN run_payloads p ON p.session_id = e.session_id AND p.seq = e.seq
            WHERE e.session_id = ? AND e.kind = 'llm' AND e.phase IN ({marks})
            ORDER BY e.seq""",
        (session_id, *phases),
    ).fetchall()
    return [dict(r) for r in rows]


def search_outputs(query: str, days: int = 14, limit: int = 12) -> list[dict]:
    """Done dispatch runs in the last `days` whose output contains every word of `query` (case-insensitive)."""
    from datetime import timedelta
    words = [w for w in re.split(r"\s+", query.strip().lower()) if w][:6]
    if not words:
        return []
    conn = _get_conn()
    since = (datetime.now(timezone.utc) - timedelta(days=days)).isoformat()
    likes = " AND ".join("LOWER(m.content) LIKE ?" for _ in words)
    rows = conn.execute(
        f"""SELECT s.id, s.dispatch_id, s.agent, s.started_at, m.content FROM sessions s
            JOIN messages m ON m.id = (SELECT id FROM messages WHERE session_id = s.id AND role = 'assistant' ORDER BY id DESC LIMIT 1)
            WHERE s.dispatch_id IS NOT NULL AND s.started_at >= ? AND {_DONE_RUN} AND {likes}
            ORDER BY s.started_at DESC LIMIT ?""",
        (since, *[f"%{w}%" for w in words], limit),
    ).fetchall()
    return [dict(r) for r in rows]


# ── Digest chat threads: sessions with source 'chat'; turns in messages, tool steps in tool_calls ──

def create_chat(title: str, focus_session_id: Optional[int] = None) -> int:
    conn = _get_conn()
    cur = conn.execute(
        "INSERT INTO sessions (model, agent, source, started_at, viewed_at, title, focus_session_id) VALUES ('', 'chat', 'chat', ?, ?, ?, ?)",
        (_now(), _now(), title[:120], focus_session_id),
    )
    conn.commit()
    return cur.lastrowid


def list_chats(dispatch_id: Optional[int] = None, limit: int = 100) -> list[dict]:
    conn = _get_conn()
    where = "AND f.dispatch_id = ?" if dispatch_id else ""
    rows = conn.execute(
        f"""SELECT s.id, s.title, s.started_at, s.focus_session_id, f.dispatch_id AS focus_dispatch_id,
                   COALESCE((SELECT MAX(created_at) FROM messages WHERE session_id = s.id), s.started_at) AS updated_at,
                   (SELECT COUNT(*) FROM messages WHERE session_id = s.id AND role = 'user') AS turns
            FROM sessions s LEFT JOIN sessions f ON f.id = s.focus_session_id
            WHERE s.source = 'chat' {where}
            ORDER BY updated_at DESC LIMIT ?""",
        (dispatch_id, limit) if dispatch_id else (limit,),
    ).fetchall()
    return [dict(r) for r in rows]


def chat_counts_by_dispatch() -> dict[int, int]:
    rows = _get_conn().execute(
        """SELECT f.dispatch_id, COUNT(*) AS n FROM sessions s JOIN sessions f ON f.id = s.focus_session_id
           WHERE s.source = 'chat' AND f.dispatch_id IS NOT NULL GROUP BY f.dispatch_id"""
    ).fetchall()
    return {r["dispatch_id"]: r["n"] for r in rows}


def get_chat(chat_id: int) -> Optional[dict]:
    conn = _get_conn()
    s = conn.execute(
        """SELECT s.id, s.title, s.started_at, s.focus_session_id, f.dispatch_id AS focus_dispatch_id
           FROM sessions s LEFT JOIN sessions f ON f.id = s.focus_session_id WHERE s.id = ? AND s.source = 'chat'""",
        (chat_id,),
    ).fetchone()
    if not s:
        return None
    msgs = conn.execute("SELECT id, role, content, created_at FROM messages WHERE session_id = ? ORDER BY id", (chat_id,)).fetchall()
    steps = conn.execute("SELECT id, tool_name, args, duration_ms, created_at FROM tool_calls WHERE session_id = ? ORDER BY id", (chat_id,)).fetchall()
    return {**dict(s), "messages": [dict(m) for m in msgs],
            "steps": [{**dict(t), "args": json.loads(t["args"] or "{}")} for t in steps]}


def delete_chat(chat_id: int) -> bool:
    conn = _get_conn()
    if not conn.execute("SELECT 1 FROM sessions WHERE id = ? AND source = 'chat'", (chat_id,)).fetchone():
        return False
    for table in ("messages", "tool_calls"):
        conn.execute(f"DELETE FROM {table} WHERE session_id = ?", (chat_id,))
    conn.execute("DELETE FROM sessions WHERE id = ?", (chat_id,))
    conn.commit()
    return True


def create_dispatch(
    agent: str,
    prompt: str,
    schedule: str,
    label: str,
    run_on_create: bool = False,
    trigger_type: str = "cron",
    repeat_until: Optional[str] = None,
    max_runs: Optional[int] = None,
    timezone_name: str = "UTC",
    tags: Optional[list] = None,
) -> int:
    conn = _get_conn()
    cur = conn.execute(
        """INSERT INTO dispatches
               (agent, prompt, schedule, label, enabled, run_on_create, created_at,
                trigger_type, repeat_until, max_runs, run_count, timezone, tags)
           VALUES (?,?,?,?,1,?,?,?,?,?,0,?,?)""",
        (agent, prompt, schedule, label, int(run_on_create), _now(),
         trigger_type, repeat_until, max_runs, timezone_name, json.dumps(clean_tags(tags))),
    )
    conn.commit()
    return cur.lastrowid


def get_dispatches(enabled_only: bool = False) -> list[dict]:
    conn = _get_conn()
    where = "WHERE d.enabled = 1" if enabled_only else ""
    rows = conn.execute(f"""
        SELECT d.*,
            (SELECT GROUP_CONCAT(has_run, ',')
             FROM (
               SELECT CASE WHEN EXISTS(
                   SELECT 1 FROM sessions WHERE dispatch_id = d.id AND DATE(started_at) = DATE('now', '-6 days')
               ) THEN '1' ELSE '0' END AS has_run UNION ALL
               SELECT CASE WHEN EXISTS(
                   SELECT 1 FROM sessions WHERE dispatch_id = d.id AND DATE(started_at) = DATE('now', '-5 days')
               ) THEN '1' ELSE '0' END UNION ALL
               SELECT CASE WHEN EXISTS(
                   SELECT 1 FROM sessions WHERE dispatch_id = d.id AND DATE(started_at) = DATE('now', '-4 days')
               ) THEN '1' ELSE '0' END UNION ALL
               SELECT CASE WHEN EXISTS(
                   SELECT 1 FROM sessions WHERE dispatch_id = d.id AND DATE(started_at) = DATE('now', '-3 days')
               ) THEN '1' ELSE '0' END UNION ALL
               SELECT CASE WHEN EXISTS(
                   SELECT 1 FROM sessions WHERE dispatch_id = d.id AND DATE(started_at) = DATE('now', '-2 days')
               ) THEN '1' ELSE '0' END UNION ALL
               SELECT CASE WHEN EXISTS(
                   SELECT 1 FROM sessions WHERE dispatch_id = d.id AND DATE(started_at) = DATE('now', '-1 days')
               ) THEN '1' ELSE '0' END UNION ALL
               SELECT CASE WHEN EXISTS(
                   SELECT 1 FROM sessions WHERE dispatch_id = d.id AND DATE(started_at) = DATE('now')
               ) THEN '1' ELSE '0' END
             )
            ) AS last_7_days,
            -- every run (scheduled, catch-up, manual); run_count only counts the scheduled ones, for max_runs
            (SELECT COUNT(*) FROM sessions WHERE dispatch_id = d.id
               AND COALESCE(json_extract(metrics, '$.status'), '') != 'missed') AS runs_total,
            ls.id AS last_session_id, ls.started_at AS last_started_at,
            json_extract(ls.metrics, '$.status')        AS last_status,
            json_extract(ls.metrics, '$.sources')       AS last_sources,
            json_extract(ls.metrics, '$.timing.run_ms') AS last_run_ms,
            json_extract(ls.metrics, '$.degraded')      AS last_degraded,
            json_extract(ls.metrics, '$.checks')        AS last_checks
        FROM dispatches d
        LEFT JOIN sessions ls ON ls.id = (SELECT id FROM sessions WHERE dispatch_id = d.id AND ended_at IS NOT NULL ORDER BY started_at DESC, id DESC LIMIT 1)
        {where}
        ORDER BY d.id DESC
    """).fetchall()
    return [dict(r) for r in rows]


def delete_dispatch(dispatch_id: int) -> None:
    conn = _get_conn()
    conn.execute("DELETE FROM dispatches WHERE id = ?", (dispatch_id,))
    conn.commit()


def update_dispatch_run(dispatch_id: int) -> dict:
    conn = _get_conn()
    conn.execute("UPDATE dispatches SET run_count = run_count + 1 WHERE id = ?", (dispatch_id,))
    conn.commit()
    row = dict(conn.execute("SELECT * FROM dispatches WHERE id = ?", (dispatch_id,)).fetchone())
    should_disable = row.get("trigger_type") == "once"
    if row.get("max_runs") is not None and row["run_count"] >= row["max_runs"]:
        should_disable = True
    if row.get("repeat_until") is not None and _now() >= row["repeat_until"]:
        should_disable = True
    if should_disable:
        conn.execute("UPDATE dispatches SET enabled = 0 WHERE id = ?", (dispatch_id,))
        conn.commit()
        row["enabled"] = 0
    return row


def toggle_dispatch(dispatch_id: int, enabled: bool) -> None:
    conn = _get_conn()
    conn.execute("UPDATE dispatches SET enabled = ? WHERE id = ?", (int(enabled), dispatch_id))
    conn.commit()


def update_dispatch(
    dispatch_id: int,
    agent: str,
    prompt: str,
    schedule: str,
    label: str,
    trigger_type: str = "cron",
    repeat_until: Optional[str] = None,
    max_runs: Optional[int] = None,
    timezone_name: str = "UTC",
    tags: Optional[list] = None,
) -> None:
    conn = _get_conn()
    conn.execute(
        """UPDATE dispatches
           SET agent=?, prompt=?, schedule=?, label=?,
               trigger_type=?, repeat_until=?, max_runs=?, timezone=?, tags=?
           WHERE id=?""",
        (agent, prompt, schedule, label, trigger_type, repeat_until, max_runs, timezone_name,
         json.dumps(clean_tags(tags)), dispatch_id),
    )
    conn.commit()


def search_sessions(
    agent: Optional[str] = None,
    keyword: Optional[str] = None,
    limit: int = 5,
) -> list[dict]:
    conn = _get_conn()
    conditions: list[str] = []
    params: list = []
    if agent:
        conditions.append("s.agent = ?")
        params.append(agent)
    if keyword:
        conditions.append(
            "(EXISTS (SELECT 1 FROM messages WHERE session_id = s.id AND role = 'user'      AND content LIKE ?)"
            " OR EXISTS (SELECT 1 FROM messages WHERE session_id = s.id AND role = 'assistant' AND content LIKE ?))"
        )
        params.extend([f"%{keyword}%", f"%{keyword}%"])
    where = ("WHERE " + " AND ".join(conditions)) if conditions else ""
    params.append(min(limit, 20))
    rows = conn.execute(
        f"""
        SELECT
            s.id, s.agent, s.source, s.started_at,
            (SELECT content FROM messages WHERE session_id = s.id AND role = 'user'      ORDER BY id LIMIT 1) AS prompt,
            (SELECT content FROM messages WHERE session_id = s.id AND role = 'assistant' ORDER BY id LIMIT 1) AS response
        FROM sessions s
        {where}
        ORDER BY s.id DESC
        LIMIT ?
        """,
        params,
    ).fetchall()
    results = []
    for r in rows:
        d = dict(r)
        if d.get("response") and len(d["response"]) > 600:
            d["response"] = d["response"][:600] + "…"
        results.append(d)
    return results


_DEFAULT_DISPATCHES = [
    {
        "agent":    "news",
        "prompt":   "Find today's most important news stories from the last 24 hours. Cover: breaking/major news, viral stories, events affecting many people, surprising or unusual stories, and local SF/Bay Area/California news.",
        "schedule": "0 7 * * *",
        "label":    "Daily news — 7am",
    },
    {
        "agent":    "reddit",
        "prompt":   "Run the Reddit monitor: digest configured subreddits and search the keyword list.",
        "schedule": "30 7 * * *",
        "label":    "Reddit digest — 7:30am",
    },
    {
        "agent":    "ai",
        "prompt":   "Find the latest developments in AI/ML: recent research papers, technology breakthroughs, business and product trends, and operational/MLOps/LLMOps challenges.",
        "schedule": "0 9 * * 2",
        "label":    "AI briefing — Tue 9am",
    },
    {
        "agent":    "local-events",
        "prompt":   "Find upcoming SF events this week, this month, and this year; active local SF communities; and notable new or reopening SF restaurants.",
        "schedule": "0 9 * * 4",
        "label":    "SF events — Thu 9am",
    },
    {
        "agent":    "media",
        "prompt":   "Find what movies, TV shows, streaming titles, and video games are newly released or playing this week in SF and on major streaming platforms.",
        "schedule": "0 10 * * 6",
        "label":    "Media — Sat 10am",
    },
    {
        "agent":    "camping",
        "prompt":   "Find available campsite reservations in California over the next 60 days to 6 months.",
        "schedule": "0 10 * * 0",
        "label":    "Camping — Sun 10am",
    },
]


def seed_default_dispatches() -> None:
    """No-op — reserved for future onboarding flows."""
    pass


_AGENT_FIELDS = [
    "name", "label", "description", "enabled",
    "goal_template", "output_template",
    "sources", "sub_queries", "tool_bindings", "direct_sources",
    "freshness", "planner_query_count", "max_validator_loops",
    "validate_links", "fetch_content",
    "is_builtin", "color", "tags", "recall_runs", "created_at", "updated_at",
]


def get_agents(enabled_only: bool = False) -> list[dict]:
    conn = _get_conn()
    where = "WHERE enabled = 1" if enabled_only else ""
    rows = conn.execute(f"SELECT * FROM agents {where} ORDER BY name").fetchall()
    return [dict(r) for r in rows]


def get_agent(name: str) -> Optional[dict]:
    conn = _get_conn()
    row = conn.execute("SELECT * FROM agents WHERE name = ?", (name,)).fetchone()
    return dict(row) if row else None


def create_agent(data: dict) -> int:
    conn = _get_conn()
    if "color" not in data or data["color"] is None:
        data["color"] = random.randint(0, 359)
    fields = [f for f in _AGENT_FIELDS if f in data and f not in ("created_at", "updated_at")]
    placeholders = ", ".join("?" for _ in fields)
    col_list = ", ".join(fields)
    values = [data[f] for f in fields]
    cur = conn.execute(
        f"INSERT INTO agents ({col_list}) VALUES ({placeholders})",
        values,
    )
    conn.commit()
    return cur.lastrowid


def update_agent(name: str, data: dict) -> None:
    conn = _get_conn()
    fields = [f for f in _AGENT_FIELDS if f in data and f not in ("name", "created_at", "updated_at", "is_builtin")]
    if not fields:
        return
    set_clause = ", ".join(f"{f} = ?" for f in fields)
    values = [data[f] for f in fields] + [_now(), name]
    conn.execute(
        f"UPDATE agents SET {set_clause}, updated_at = ? WHERE name = ?",
        values,
    )
    conn.commit()


def delete_agent(name: str) -> None:
    conn = _get_conn()
    conn.execute("DELETE FROM agents WHERE name = ? AND is_builtin = 0", (name,))
    conn.commit()


def seed_agents_from_yaml(seeds_dir: Optional[Path] = None) -> None:
    """Upsert agent definitions from YAML files into the DB.

    Files that already exist in the DB are skipped unless the YAML has changed
    (detected by comparing goal_template + sources as a lightweight fingerprint).
    Built-in agents (is_builtin=1) are always refreshed from YAML on startup so
    edits to seed files take effect without manual DB surgery.
    """
    if not _YAML_AVAILABLE:
        return
    if seeds_dir is None:
        env_dir = os.getenv("DISPATCH_AGENTS_DIR")
        seeds_dir = Path(env_dir) if env_dir else Path(__file__).parent / "agents"
    if not seeds_dir.is_dir():
        return

    conn = _get_conn()
    for path in sorted(seeds_dir.glob("*.yaml")):
        try:
            data = _yaml.safe_load(path.read_text())
        except Exception:
            continue
        if not data or not isinstance(data, dict) or not data.get("name"):
            continue

        _normalise_agent_yaml(data)

        existing = conn.execute(
            "SELECT id, is_builtin FROM agents WHERE name = ?", (data["name"],)
        ).fetchone()

        if existing is None:
            # New agent — insert; assign a random color if not specified in YAML
            if "color" not in data or data["color"] is None:
                data["color"] = random.randint(0, 359)
            fields = [f for f in _AGENT_FIELDS if f in data and f not in ("created_at", "updated_at")]
            placeholders = ", ".join("?" for _ in fields)
            col_list = ", ".join(fields)
            conn.execute(
                f"INSERT INTO agents ({col_list}) VALUES ({placeholders})",
                [data[f] for f in fields],
            )
        elif data.get("is_builtin", 1):
            # Built-in agent that already exists — refresh from YAML but preserve user-set color
            fields = [f for f in _AGENT_FIELDS if f in data and f not in ("name", "created_at", "is_builtin", "color")]
            if fields:
                set_clause = ", ".join(f"{f} = ?" for f in fields)
                conn.execute(
                    f"UPDATE agents SET {set_clause} WHERE name = ?",
                    [data[f] for f in fields] + [data["name"]],
                )
        # User-created agents (is_builtin=0) are never overwritten by seed runs

    conn.commit()


def clean_tags(tags) -> list[str]:
    """Lowercase, trimmed, unique, in order; accepts a list or a comma-separated string."""
    if isinstance(tags, str):
        tags = tags.split(",")
    out: list[str] = []
    for t in tags or []:
        t = " ".join(str(t).strip().lower().split())
        if t and t not in out:
            out.append(t)
    return out


def _normalise_agent_yaml(data: dict) -> None:
    """Convert YAML parsed dict into DB-ready types in-place."""
    data.setdefault("is_builtin", 1)
    # JSON-encode list/dict fields
    if "tags" in data:
        data["tags"] = json.dumps(clean_tags(data["tags"]))
    for key in ("sources", "tool_bindings", "direct_sources"):
        v = data.get(key)
        if isinstance(v, list):
            data[key] = json.dumps(v)
        elif v is None and key in ("sources", "tool_bindings"):
            data[key] = "[]"
    if "sub_queries" in data and isinstance(data["sub_queries"], list):
        data["sub_queries"] = json.dumps(data["sub_queries"])
    # Booleans → ints
    for key in ("enabled", "validate_links", "fetch_content"):
        if key in data:
            data[key] = int(bool(data[key]))


def get_settings() -> dict[str, str]:
    conn = _get_conn()
    rows = conn.execute("SELECT key, value FROM settings").fetchall()
    return {r["key"]: r["value"] for r in rows}


def get_setting(key: str, default: Optional[str] = None) -> Optional[str]:
    conn = _get_conn()
    row = conn.execute("SELECT value FROM settings WHERE key = ?", (key,)).fetchone()
    return row["value"] if row else default


def set_setting(key: str, value: str) -> None:
    conn = _get_conn()
    conn.execute(
        "INSERT OR REPLACE INTO settings (key, value, updated_at) VALUES (?, ?, ?)",
        (key, value, _now()),
    )
    conn.commit()


def close() -> None:
    conn = getattr(_local, "conn", None)
    if conn:
        conn.close()
        _local.conn = None