"""SQLite-backed implementation of DobbyStore."""

from __future__ import annotations

import json
import logging
from datetime import UTC

import aiosqlite

from dobby.models import (
    Priority,
    Project,
    Result,
    Spoke,
    Task,
    TaskStatus,
    Todo,
    validate_transition,
)

logger = logging.getLogger(__name__)

# Current schema version — bump when adding migrations.
SCHEMA_VERSION = 11

# Priority weight map for ORDER BY in SQL
_PRIORITY_WEIGHTS = {
    Priority.URGENT: 40,
    Priority.HIGH: 30,
    Priority.NORMAL: 20,
    Priority.LOW: 10,
}

_WEIGHT_CASE = (
    "CASE json_extract(data, '$.priority') "
    "WHEN 'urgent' THEN 40 "
    "WHEN 'high' THEN 30 "
    "WHEN 'normal' THEN 20 "
    "WHEN 'low' THEN 10 "
    "ELSE 20 END"
)


class SqliteStore:
    """Implements DobbyStore using SQLite with JSON blob storage."""

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

    async def initialize(self) -> None:
        self._db = await aiosqlite.connect(self._db_path)
        await self._db.execute("PRAGMA journal_mode=WAL")

        # Create the schema_version table first (always safe, idempotent)
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS schema_version (
                id INTEGER PRIMARY KEY CHECK (id = 1),
                version INTEGER NOT NULL
            )
            """
        )

        # Determine current DB version
        async with self._db.execute(
            "SELECT version FROM schema_version WHERE id = 1"
        ) as cur:
            row = await cur.fetchone()
        current_version = row[0] if row else 0

        if current_version == 0:
            # Fresh database — create all tables at latest schema
            await self._create_tables_v1()
            await self._db.execute(
                "INSERT OR REPLACE INTO schema_version (id, version) VALUES (1, ?)",
                (SCHEMA_VERSION,),
            )
        elif current_version < SCHEMA_VERSION:
            # Run incremental migrations
            await self._run_migrations(current_version)

        await self._db.commit()

    async def _create_tables_v1(self) -> None:
        """Create the base schema (v1 tables)."""
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS tasks (
                id TEXT PRIMARY KEY,
                data TEXT NOT NULL
            )
            """
        )
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS results (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id TEXT NOT NULL,
                data TEXT NOT NULL
            )
            """
        )
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS projects (
                name TEXT PRIMARY KEY,
                data TEXT NOT NULL
            )
            """
        )
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS spokes (
                id TEXT PRIMARY KEY,
                data TEXT NOT NULL
            )
            """
        )

        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS utilization_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp TEXT NOT NULL,
                five_hour REAL,
                seven_day REAL
            )
            """
        )

        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS memories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                key TEXT NOT NULL UNIQUE,
                value TEXT NOT NULL,
                type TEXT DEFAULT 'general',
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            )
            """
        )

        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS todos (
                id TEXT PRIMARY KEY,
                data TEXT NOT NULL
            )
            """
        )

        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS daily_plans (
                id TEXT PRIMARY KEY,
                data TEXT NOT NULL
            )
            """
        )

        # Trace system tables (v11)
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS traces (
                trace_id TEXT PRIMARY KEY,
                source TEXT NOT NULL,
                source_id TEXT,
                user_id TEXT,
                project TEXT,
                input_text TEXT,
                created_at TEXT NOT NULL,
                completed_at TEXT,
                status TEXT DEFAULT 'active',
                total_input_tokens INTEGER DEFAULT 0,
                total_output_tokens INTEGER DEFAULT 0,
                total_cost_usd REAL DEFAULT 0
            )
            """
        )
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS spans (
                span_id TEXT PRIMARY KEY,
                trace_id TEXT NOT NULL,
                parent_span_id TEXT,
                kind TEXT NOT NULL,
                task_id TEXT,
                session_id TEXT,
                log_path TEXT,
                started_at TEXT NOT NULL,
                ended_at TEXT,
                duration_ms INTEGER,
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                cost_usd REAL DEFAULT 0,
                metadata JSON,
                FOREIGN KEY (trace_id) REFERENCES traces(trace_id)
            )
            """
        )
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS execution_turns (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                span_id TEXT NOT NULL,
                turn_number INTEGER,
                role TEXT,
                content_type TEXT,
                tool_name TEXT,
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                cache_read_tokens INTEGER DEFAULT 0,
                cost_usd REAL DEFAULT 0,
                content_preview TEXT,
                timestamp TEXT,
                FOREIGN KEY (span_id) REFERENCES spans(span_id)
            )
            """
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_spans_trace ON spans(trace_id)"
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_spans_task ON spans(task_id)"
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_turns_span ON execution_turns(span_id)"
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_traces_project ON traces(project)"
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_traces_created ON traces(created_at)"
        )

    async def _run_migrations(self, from_version: int) -> None:
        """Run migrations from from_version up to SCHEMA_VERSION."""
        migrations = {
            2: self._migrate_v1_to_v2,
            3: self._migrate_v2_to_v3,
            4: self._migrate_v3_to_v4,
            5: self._migrate_v4_to_v5,
            6: self._migrate_v5_to_v6,
            7: self._migrate_v6_to_v7,
            8: self._migrate_v7_to_v8,
            9: self._migrate_v8_to_v9,
            10: self._migrate_v9_to_v10,
            11: self._migrate_v10_to_v11,
        }
        for ver in range(from_version + 1, SCHEMA_VERSION + 1):
            migrate_fn = migrations.get(ver)
            if migrate_fn is not None:
                logger.info("Running DB migration to v%d", ver)
                await migrate_fn()
        await self._db.execute(
            "UPDATE schema_version SET version = ? WHERE id = 1",
            (SCHEMA_VERSION,),
        )

    async def _migrate_v1_to_v2(self) -> None:
        """v1→v2: backfill events and attempt_number into existing task/result JSON blobs."""
        async with self._db.execute("SELECT id, data FROM tasks") as cur:
            rows = await cur.fetchall()
        for task_id, raw in rows:
            d = json.loads(raw)
            if "events" not in d:
                d["events"] = [
                    {"status": d["status"], "timestamp": d["created_at"], "detail": "created"}
                ]
            if "attempt_number" not in d:
                d["attempt_number"] = 0
            await self._db.execute(
                "UPDATE tasks SET data = ? WHERE id = ?",
                (json.dumps(d), task_id),
            )
        async with self._db.execute("SELECT id, data FROM results") as cur:
            rows = await cur.fetchall()
        for row_id, raw in rows:
            d = json.loads(raw)
            if "attempt_number" not in d:
                d["attempt_number"] = 0
            await self._db.execute(
                "UPDATE results SET data = ? WHERE id = ?",
                (json.dumps(d), row_id),
            )

    async def _migrate_v2_to_v3(self) -> None:
        """v2→v3: backfill next_run_at, max_failures, consecutive_failures into task JSON blobs."""
        async with self._db.execute("SELECT id, data FROM tasks") as cur:
            rows = await cur.fetchall()
        for task_id, raw in rows:
            d = json.loads(raw)
            changed = False
            if "next_run_at" not in d:
                d["next_run_at"] = None
                changed = True
            if "max_failures" not in d:
                d["max_failures"] = 3
                changed = True
            if "consecutive_failures" not in d:
                d["consecutive_failures"] = 0
                changed = True
            if changed:
                await self._db.execute(
                    "UPDATE tasks SET data = ? WHERE id = ?",
                    (json.dumps(d), task_id),
                )

    async def _migrate_v3_to_v4(self) -> None:
        """v3→v4: backfill batch_id into task JSON blobs."""
        async with self._db.execute("SELECT id, data FROM tasks") as cur:
            rows = await cur.fetchall()
        for task_id, raw in rows:
            d = json.loads(raw)
            if "batch_id" not in d:
                d["batch_id"] = None
                await self._db.execute(
                    "UPDATE tasks SET data = ? WHERE id = ?",
                    (json.dumps(d), task_id),
                )

    async def _migrate_v5_to_v6(self) -> None:
        """v5→v6: add spokes table and target field to tasks."""
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS spokes (
                id TEXT PRIMARY KEY,
                data TEXT NOT NULL
            )
            """
        )
        # Backfill target field into existing task JSON blobs
        async with self._db.execute("SELECT id, data FROM tasks") as cur:
            rows = await cur.fetchall()
        for tid, raw in rows:
            d = json.loads(raw)
            if "target" not in d:
                d["target"] = "local"
                await self._db.execute(
                    "UPDATE tasks SET data = ? WHERE id = ?",
                    (json.dumps(d), tid),
                )

    async def _migrate_v6_to_v7(self) -> None:
        """v6→v7: add utilization_history table."""
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS utilization_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp TEXT NOT NULL,
                five_hour REAL,
                seven_day REAL
            )
            """
        )

    async def _migrate_v7_to_v8(self) -> None:
        """v7→v8: add memories table."""
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS memories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                key TEXT NOT NULL UNIQUE,
                value TEXT NOT NULL,
                type TEXT DEFAULT 'general',
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            )
            """
        )

    async def _migrate_v8_to_v9(self) -> None:
        """v8→v9: add todos table."""
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS todos (
                id TEXT PRIMARY KEY,
                data TEXT NOT NULL
            )
            """
        )

    async def _migrate_v9_to_v10(self) -> None:
        """v9→v10: add daily_plans table."""
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS daily_plans (
                id TEXT PRIMARY KEY,
                data TEXT NOT NULL
            )
            """
        )

    async def _migrate_v10_to_v11(self) -> None:
        """v10→v11: add traces, spans, and execution_turns tables."""
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS traces (
                trace_id TEXT PRIMARY KEY,
                source TEXT NOT NULL,
                source_id TEXT,
                user_id TEXT,
                project TEXT,
                input_text TEXT,
                created_at TEXT NOT NULL,
                completed_at TEXT,
                status TEXT DEFAULT 'active',
                total_input_tokens INTEGER DEFAULT 0,
                total_output_tokens INTEGER DEFAULT 0,
                total_cost_usd REAL DEFAULT 0
            )
            """
        )
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS spans (
                span_id TEXT PRIMARY KEY,
                trace_id TEXT NOT NULL,
                parent_span_id TEXT,
                kind TEXT NOT NULL,
                task_id TEXT,
                session_id TEXT,
                log_path TEXT,
                started_at TEXT NOT NULL,
                ended_at TEXT,
                duration_ms INTEGER,
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                cost_usd REAL DEFAULT 0,
                metadata JSON,
                FOREIGN KEY (trace_id) REFERENCES traces(trace_id)
            )
            """
        )
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS execution_turns (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                span_id TEXT NOT NULL,
                turn_number INTEGER,
                role TEXT,
                content_type TEXT,
                tool_name TEXT,
                input_tokens INTEGER DEFAULT 0,
                output_tokens INTEGER DEFAULT 0,
                cache_read_tokens INTEGER DEFAULT 0,
                cost_usd REAL DEFAULT 0,
                content_preview TEXT,
                timestamp TEXT,
                FOREIGN KEY (span_id) REFERENCES spans(span_id)
            )
            """
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_spans_trace ON spans(trace_id)"
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_spans_task ON spans(task_id)"
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_turns_span ON execution_turns(span_id)"
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_traces_project ON traces(project)"
        )
        await self._db.execute(
            "CREATE INDEX IF NOT EXISTS idx_traces_created ON traces(created_at)"
        )
        # Backfill trace_id into existing task JSON blobs
        async with self._db.execute("SELECT id, data FROM tasks") as cur:
            rows = await cur.fetchall()
        for task_id, raw in rows:
            d = json.loads(raw)
            if "trace_id" not in d:
                d["trace_id"] = None
                await self._db.execute(
                    "UPDATE tasks SET data = ? WHERE id = ?",
                    (json.dumps(d), task_id),
                )

    async def _migrate_v4_to_v5(self) -> None:
        """v4→v5: backfill policies into project JSON blobs."""
        async with self._db.execute("SELECT name, data FROM projects") as cur:
            rows = await cur.fetchall()
        for name, raw in rows:
            d = json.loads(raw)
            if "policies" not in d:
                d["policies"] = {
                    "default_backend": None,
                    "default_timeout": None,
                    "max_concurrent": None,
                    "default_priority": None,
                    "auto_chain": False,
                }
                await self._db.execute(
                    "UPDATE projects SET data = ? WHERE name = ?",
                    (json.dumps(d), name),
                )

    async def get_schema_version(self) -> int:
        """Return the current schema version of the database."""
        async with self._db.execute(
            "SELECT version FROM schema_version WHERE id = 1"
        ) as cur:
            row = await cur.fetchone()
        return row[0] if row else 0

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

    # ------------------------------------------------------------------
    # Tasks
    # ------------------------------------------------------------------

    async def create_task(self, task: Task) -> str:
        data = json.dumps(task.to_dict())
        await self._db.execute(
            "INSERT INTO tasks (id, data) VALUES (?, ?)",
            (task.id, data),
        )
        await self._db.commit()
        return task.id

    async def get_task(self, task_id: str) -> Task:
        async with self._db.execute(
            "SELECT data FROM tasks WHERE id = ?", (task_id,)
        ) as cursor:
            row = await cursor.fetchone()
        if row is None:
            raise KeyError(f"Task not found: {task_id}")
        return Task.from_dict(json.loads(row[0]))

    async def update_task(self, task_id: str, **fields) -> None:
        task = await self.get_task(task_id)

        # Enforce state transition guardrails
        if "status" in fields:
            new_status = TaskStatus(fields["status"])
            validate_transition(task.status, new_status)
            # Record event in the task's journal
            task.record_event(new_status)

        d = task.to_dict()
        d.update(fields)
        # Restore the updated task via from_dict to validate fields
        updated = Task.from_dict(d)
        await self._db.execute(
            "UPDATE tasks SET data = ? WHERE id = ?",
            (json.dumps(updated.to_dict()), task_id),
        )
        await self._db.commit()

    async def list_tasks(
        self,
        project: str | None = None,
        status: str | None = None,
        search: str | None = None,
        sort: str = "created_at",
        limit: int | None = None,
    ) -> list[Task]:
        clauses: list[str] = []
        params: list = []

        if project is not None:
            clauses.append("json_extract(data, '$.project') = ?")
            params.append(project)
        if status is not None:
            clauses.append("json_extract(data, '$.status') = ?")
            params.append(status)
        if search is not None:
            clauses.append(
                "(json_extract(data, '$.description') LIKE ?"
                " OR json_extract(data, '$.exit_criteria') LIKE ?)"
            )
            pattern = f"%{search}%"
            params.extend([pattern, pattern])

        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
        # Determine sort clause
        _sort_map = {
            "created_at": "json_extract(data, '$.created_at') ASC",
            "-created_at": "json_extract(data, '$.created_at') DESC",
            "priority": f"{_WEIGHT_CASE} ASC",
            "-priority": f"{_WEIGHT_CASE} DESC",
        }
        order = _sort_map.get(sort, "json_extract(data, '$.created_at') ASC")
        query = f"SELECT data FROM tasks {where} ORDER BY {order}"
        if limit is not None:
            query += " LIMIT ?"
            params.append(limit)
        async with self._db.execute(query, params) as cursor:
            rows = await cursor.fetchall()
        return [Task.from_dict(json.loads(r[0])) for r in rows]

    async def get_pending_tasks(self, limit: int = 10) -> list[Task]:
        query = (
            f"SELECT data FROM tasks "
            f"WHERE json_extract(data, '$.status') IN ('pending', 'drafting') "
            f"ORDER BY {_WEIGHT_CASE} DESC, json_extract(data, '$.created_at') ASC "
            f"LIMIT ?"
        )
        async with self._db.execute(query, (limit,)) as cursor:
            rows = await cursor.fetchall()
        return [Task.from_dict(json.loads(r[0])) for r in rows]

    async def get_scheduled_tasks(self) -> list[Task]:
        """Return tasks with non-once schedules and status = 'active'."""
        query = (
            "SELECT data FROM tasks "
            "WHERE json_extract(data, '$.schedule.type') != 'once' "
            "AND json_extract(data, '$.status') = 'active'"
        )
        async with self._db.execute(query) as cursor:
            rows = await cursor.fetchall()
        return [Task.from_dict(json.loads(r[0])) for r in rows]

    async def get_dependent_tasks(self, task_id: str) -> list[Task]:
        """Return all AFTER tasks that depend on the given task_id."""
        query = (
            "SELECT data FROM tasks "
            "WHERE json_extract(data, '$.schedule.type') = 'after' "
            "AND json_extract(data, '$.schedule.depends_on') = ?"
        )
        async with self._db.execute(query, (task_id,)) as cursor:
            rows = await cursor.fetchall()
        return [Task.from_dict(json.loads(r[0])) for r in rows]

    # ------------------------------------------------------------------
    # Batches
    # ------------------------------------------------------------------

    async def list_tasks_by_batch(self, batch_id: str) -> list[Task]:
        """Return all tasks belonging to a batch."""
        query = (
            "SELECT data FROM tasks "
            "WHERE json_extract(data, '$.batch_id') = ? "
            "ORDER BY json_extract(data, '$.created_at') ASC"
        )
        async with self._db.execute(query, (batch_id,)) as cursor:
            rows = await cursor.fetchall()
        return [Task.from_dict(json.loads(r[0])) for r in rows]

    async def list_batches(self) -> list[dict]:
        """Return aggregate info for every batch_id that has tasks."""
        query = (
            "SELECT json_extract(data, '$.batch_id') as bid, "
            "COUNT(*) as cnt, "
            "SUM(CASE WHEN json_extract(data, '$.status')"
            " IN ('completed', 'failed', 'cancelled') THEN 1 ELSE 0 END) as done, "
            "SUM(CASE WHEN json_extract(data, '$.status') = 'completed' THEN 1 ELSE 0 END) as ok, "
            "SUM(CASE WHEN json_extract(data, '$.status') = 'failed' THEN 1 ELSE 0 END) as fail, "
            "SUM(CASE WHEN json_extract(data, '$.status') = 'running'"
            " THEN 1 ELSE 0 END) as running, "
            "MIN(json_extract(data, '$.created_at')) as created "
            "FROM tasks "
            "WHERE json_extract(data, '$.batch_id') IS NOT NULL "
            "GROUP BY bid ORDER BY created DESC"
        )
        async with self._db.execute(query) as cursor:
            rows = await cursor.fetchall()
        return [
            {
                "batch_id": r[0],
                "total": r[1],
                "done": r[2],
                "completed": r[3],
                "failed": r[4],
                "running": r[5],
                "created_at": r[6],
            }
            for r in rows
        ]

    # ------------------------------------------------------------------
    # Results
    # ------------------------------------------------------------------

    async def store_result(self, result: Result) -> None:
        data = json.dumps(result.to_dict())
        await self._db.execute(
            "INSERT INTO results (task_id, data) VALUES (?, ?)",
            (result.task_id, data),
        )
        await self._db.commit()

    async def get_result(self, task_id: str) -> Result | None:
        async with self._db.execute(
            "SELECT data FROM results WHERE task_id = ? ORDER BY id DESC LIMIT 1",
            (task_id,),
        ) as cursor:
            row = await cursor.fetchone()
        if row is None:
            return None
        return Result.from_dict(json.loads(row[0]))

    async def list_results(
        self,
        project: str | None = None,
        task_id: str | None = None,
        last_n: int | None = None,
    ) -> list[Result]:
        # Build base query selecting id (for ordering) and data
        if project is not None:
            inner = (
                "SELECT r.id, r.data FROM results r "
                "JOIN tasks t ON r.task_id = t.id "
                "WHERE json_extract(t.data, '$.project') = ?"
            )
            params: list = [project]
        elif task_id is not None:
            inner = "SELECT id, data FROM results WHERE task_id = ?"
            params = [task_id]
        else:
            inner = "SELECT id, data FROM results"
            params = []

        if last_n is not None:
            # Get the last N by insertion id, then reverse for chronological order
            query = f"SELECT id, data FROM ({inner}) ORDER BY id DESC LIMIT ?"
            params = list(params) + [last_n]
            async with self._db.execute(query, params) as cursor:
                rows = await cursor.fetchall()
            # rows are in DESC order; reverse to chronological
            rows = list(reversed(rows))
        else:
            query = f"SELECT id, data FROM ({inner}) ORDER BY id ASC"
            async with self._db.execute(query, params) as cursor:
                rows = await cursor.fetchall()

        return [Result.from_dict(json.loads(r[1])) for r in rows]

    # ------------------------------------------------------------------
    # Projects
    # ------------------------------------------------------------------

    async def register_project(self, project: Project) -> None:
        data = json.dumps(project.to_dict())
        await self._db.execute(
            "INSERT OR REPLACE INTO projects (name, data) VALUES (?, ?)",
            (project.name, data),
        )
        await self._db.commit()

    async def get_project(self, name: str) -> Project | None:
        async with self._db.execute(
            "SELECT data FROM projects WHERE name = ?", (name,)
        ) as cursor:
            row = await cursor.fetchone()
        if row is None:
            return None
        return Project.from_dict(json.loads(row[0]))

    async def list_projects(self) -> list[Project]:
        async with self._db.execute(
            "SELECT data FROM projects ORDER BY name ASC"
        ) as cursor:
            rows = await cursor.fetchall()
        return [Project.from_dict(json.loads(r[0])) for r in rows]

    # ------------------------------------------------------------------
    # Spokes
    # ------------------------------------------------------------------

    async def register_spoke(self, spoke: Spoke) -> None:
        await self._db.execute(
            "INSERT OR REPLACE INTO spokes (id, data) VALUES (?, ?)",
            (spoke.id, json.dumps(spoke.to_dict())),
        )
        await self._db.commit()

    async def get_spoke(self, spoke_id: str) -> Spoke | None:
        async with self._db.execute(
            "SELECT data FROM spokes WHERE id = ?", (spoke_id,)
        ) as cur:
            row = await cur.fetchone()
        if row is None:
            return None
        return Spoke.from_dict(json.loads(row[0]))

    async def update_spoke(self, spoke_id: str, **updates) -> None:
        async with self._db.execute(
            "SELECT data FROM spokes WHERE id = ?", (spoke_id,)
        ) as cur:
            row = await cur.fetchone()
        if row is None:
            return
        d = json.loads(row[0])
        d.update(updates)
        await self._db.execute(
            "UPDATE spokes SET data = ? WHERE id = ?",
            (json.dumps(d), spoke_id),
        )
        await self._db.commit()

    async def list_spokes(self) -> list[Spoke]:
        async with self._db.execute(
            "SELECT data FROM spokes ORDER BY id ASC"
        ) as cur:
            rows = await cur.fetchall()
        return [Spoke.from_dict(json.loads(r[0])) for r in rows]

    async def claim_task_for_spoke(self, spoke_id: str) -> Task | None:
        """Find the next pending task targeted at this spoke and mark it running."""
        async with self._db.execute(
            "SELECT id, data FROM tasks"
        ) as cur:
            rows = await cur.fetchall()
        for tid, raw in rows:
            d = json.loads(raw)
            if d.get("status") not in ("pending", "drafting"):
                continue
            target = d.get("target", "local")
            if target != spoke_id and target != "any":
                continue
            # Claim it
            d["status"] = "running"
            from datetime import datetime
            d["started_at"] = datetime.now(UTC).isoformat()
            await self._db.execute(
                "UPDATE tasks SET data = ? WHERE id = ?",
                (json.dumps(d), tid),
            )
            await self._db.commit()
            return Task.from_dict(d)
        return None

    # ------------------------------------------------------------------
    # Utilization History
    # ------------------------------------------------------------------

    async def record_utilization(
        self, timestamp: str, five_hour: float | None, seven_day: float | None
    ) -> None:
        """Record a utilization snapshot."""
        await self._db.execute(
            "INSERT INTO utilization_history (timestamp, five_hour, seven_day) VALUES (?, ?, ?)",
            (timestamp, five_hour, seven_day),
        )
        await self._db.commit()
        # Prune old entries (keep last 30 days = ~8640 at 5min intervals)
        await self._db.execute(
            "DELETE FROM utilization_history WHERE id NOT IN"
            " (SELECT id FROM utilization_history ORDER BY id DESC LIMIT 8640)"
        )
        await self._db.commit()

    async def get_utilization_history(self, limit: int = 2000) -> list[dict]:
        """Return utilization history, newest first."""
        async with self._db.execute(
            "SELECT timestamp, five_hour, seven_day FROM utilization_history"
            " ORDER BY id DESC LIMIT ?",
            (limit,),
        ) as cur:
            rows = await cur.fetchall()
        # Return in chronological order (oldest first)
        return [{"t": r[0], "5h": r[1], "7d": r[2]} for r in reversed(rows)]

    # ------------------------------------------------------------------
    # Memories
    # ------------------------------------------------------------------

    async def save_memory(self, key: str, value: str, type: str = "general") -> dict:
        """Upsert a memory — insert or update if key exists."""
        from datetime import datetime
        now = datetime.now(UTC).isoformat()
        async with self._db.execute(
            "SELECT id, created_at FROM memories WHERE key = ?", (key,)
        ) as cur:
            existing = await cur.fetchone()
        if existing:
            await self._db.execute(
                "UPDATE memories SET value = ?, type = ?, updated_at = ? WHERE key = ?",
                (value, type, now, key),
            )
            created_at = existing[1]
        else:
            await self._db.execute(
                "INSERT INTO memories (key, value, type, created_at, updated_at)"
                " VALUES (?, ?, ?, ?, ?)",
                (key, value, type, now, now),
            )
            created_at = now
        await self._db.commit()
        return {
            "key": key, "value": value, "type": type, "created_at": created_at, "updated_at": now,
        }

    async def get_memory(self, key: str) -> dict | None:
        """Return a single memory by exact key."""
        async with self._db.execute(
            "SELECT key, value, type, created_at, updated_at FROM memories WHERE key = ?",
            (key,),
        ) as cur:
            row = await cur.fetchone()
        if row is None:
            return None
        return {
            "key": row[0], "value": row[1], "type": row[2],
            "created_at": row[3], "updated_at": row[4],
        }

    async def search_memories(self, query: str | None = None, type: str | None = None) -> list[dict]:
        """Search memories by key/value substring and/or type."""
        clauses: list[str] = []
        params: list = []
        if query:
            clauses.append("(key LIKE ? OR value LIKE ?)")
            pattern = f"%{query}%"
            params.extend([pattern, pattern])
        if type:
            clauses.append("type = ?")
            params.append(type)
        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
        sql = (
            f"SELECT key, value, type, created_at, updated_at FROM memories {where}"
            " ORDER BY updated_at DESC"
        )
        async with self._db.execute(sql, params) as cur:
            rows = await cur.fetchall()
        return [
            {"key": r[0], "value": r[1], "type": r[2], "created_at": r[3], "updated_at": r[4]}
            for r in rows
        ]

    async def list_memories(self, type: str | None = None) -> list[dict]:
        """List all memories, optionally filtered by type."""
        return await self.search_memories(query=None, type=type)

    async def delete_memory(self, key: str) -> bool:
        """Delete a memory by key. Returns True if deleted, False if not found."""
        async with self._db.execute(
            "SELECT id FROM memories WHERE key = ?", (key,)
        ) as cur:
            exists = await cur.fetchone()
        if not exists:
            return False
        await self._db.execute("DELETE FROM memories WHERE key = ?", (key,))
        await self._db.commit()
        return True

    # ------------------------------------------------------------------
    # Todos
    # ------------------------------------------------------------------

    async def create_todo(self, todo: Todo) -> str:
        data = json.dumps(todo.to_dict())
        await self._db.execute(
            "INSERT INTO todos (id, data) VALUES (?, ?)",
            (todo.id, data),
        )
        await self._db.commit()
        return todo.id

    async def get_todo(self, todo_id: str) -> Todo:
        async with self._db.execute(
            "SELECT data FROM todos WHERE id = ?", (todo_id,)
        ) as cursor:
            row = await cursor.fetchone()
        if row is None:
            raise KeyError(f"Todo not found: {todo_id}")
        return Todo.from_dict(json.loads(row[0]))

    async def update_todo(self, todo_id: str, **fields) -> Todo:
        todo = await self.get_todo(todo_id)
        d = todo.to_dict()
        d.update(fields)
        # Stamp updated_at
        from datetime import datetime
        d["updated_at"] = datetime.now(UTC).isoformat()
        # If marking done, stamp completed_at
        if fields.get("status") == "done" and not d.get("completed_at"):
            d["completed_at"] = d["updated_at"]
        updated = Todo.from_dict(d)
        await self._db.execute(
            "UPDATE todos SET data = ? WHERE id = ?",
            (json.dumps(updated.to_dict()), todo_id),
        )
        await self._db.commit()
        return updated

    async def delete_todo(self, todo_id: str) -> bool:
        async with self._db.execute(
            "SELECT id FROM todos WHERE id = ?", (todo_id,)
        ) as cur:
            exists = await cur.fetchone()
        if not exists:
            return False
        await self._db.execute("DELETE FROM todos WHERE id = ?", (todo_id,))
        await self._db.commit()
        return True

    async def list_todos(
        self,
        project: str | None = None,
        status: str | None = None,
        due_date_before: str | None = None,
        source_task_id: str | None = None,
    ) -> list[Todo]:
        clauses: list[str] = []
        params: list = []

        if project is not None:
            clauses.append("json_extract(data, '$.project') = ?")
            params.append(project)
        if status is not None:
            clauses.append("json_extract(data, '$.status') = ?")
            params.append(status)
        if due_date_before is not None:
            clauses.append(
                "json_extract(data, '$.due_date') IS NOT NULL"
                " AND json_extract(data, '$.due_date') <= ?"
            )
            params.append(due_date_before)
        if source_task_id is not None:
            clauses.append("json_extract(data, '$.source_task_id') = ?")
            params.append(source_task_id)

        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
        query = f"SELECT data FROM todos {where} ORDER BY json_extract(data, '$.created_at') DESC"
        async with self._db.execute(query, params) as cursor:
            rows = await cursor.fetchall()
        return [Todo.from_dict(json.loads(r[0])) for r in rows]

    # ------------------------------------------------------------------
    # Daily Plans
    # ------------------------------------------------------------------

    async def store_daily_plan(self, plan) -> None:
        """Persist a daily plan."""
        await self._db.execute(
            "INSERT OR REPLACE INTO daily_plans (id, data) VALUES (?, ?)",
            (plan.id, json.dumps(plan.to_dict())),
        )
        await self._db.commit()

    async def get_daily_plan(self, plan_id: str):
        """Retrieve a daily plan by ID."""
        from dobby.models import DailyPlan
        async with self._db.execute(
            "SELECT data FROM daily_plans WHERE id = ?", (plan_id,)
        ) as cursor:
            row = await cursor.fetchone()
        if row:
            return DailyPlan.from_dict(json.loads(row[0]))
        return None

    async def get_plan_for_date(self, date: str):
        """Retrieve the most recent plan for a given date (YYYY-MM-DD)."""
        from dobby.models import DailyPlan
        async with self._db.execute(
            "SELECT data FROM daily_plans WHERE json_extract(data, '$.date') = ? "
            "ORDER BY json_extract(data, '$.created_at') DESC LIMIT 1",
            (date,),
        ) as cursor:
            row = await cursor.fetchone()
        if row:
            return DailyPlan.from_dict(json.loads(row[0]))
        return None

    # ------------------------------------------------------------------
    # Traces (trajectory tracking)
    # ------------------------------------------------------------------

    async def create_trace(
        self,
        trace_id: str,
        source: str,
        created_at: str,
        source_id: str | None = None,
        user_id: str | None = None,
        project: str | None = None,
        input_text: str | None = None,
    ) -> None:
        await self._db.execute(
            "INSERT INTO traces"
            " (trace_id, source, source_id, user_id, project, input_text, created_at)"
            " VALUES (?, ?, ?, ?, ?, ?, ?)",
            (trace_id, source, source_id, user_id, project, input_text, created_at),
        )
        await self._db.commit()

    async def get_trace(self, trace_id: str) -> dict | None:
        async with self._db.execute(
            "SELECT trace_id, source, source_id, user_id, project, input_text,"
            " created_at, completed_at, status, total_input_tokens,"
            " total_output_tokens, total_cost_usd FROM traces WHERE trace_id = ?",
            (trace_id,),
        ) as cur:
            row = await cur.fetchone()
        if row is None:
            return None
        return {
            "trace_id": row[0], "source": row[1], "source_id": row[2],
            "user_id": row[3], "project": row[4], "input_text": row[5],
            "created_at": row[6], "completed_at": row[7], "status": row[8],
            "total_input_tokens": row[9], "total_output_tokens": row[10],
            "total_cost_usd": row[11],
        }

    async def list_traces(
        self,
        project: str | None = None,
        source: str | None = None,
        status: str | None = None,
        days: int | None = None,
        limit: int = 50,
    ) -> list[dict]:
        clauses: list[str] = []
        params: list = []
        if project:
            clauses.append("project = ?")
            params.append(project)
        if source:
            clauses.append("source = ?")
            params.append(source)
        if status:
            clauses.append("status = ?")
            params.append(status)
        if days:
            from datetime import datetime, timedelta
            cutoff = (datetime.now(UTC) - timedelta(days=days)).isoformat()
            clauses.append("created_at >= ?")
            params.append(cutoff)
        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
        query = (
            f"SELECT trace_id, source, source_id, user_id, project, input_text,"
            f" created_at, completed_at, status, total_input_tokens,"
            f" total_output_tokens, total_cost_usd FROM traces {where}"
            f" ORDER BY created_at DESC LIMIT ?"
        )
        params.append(limit)
        async with self._db.execute(query, params) as cur:
            rows = await cur.fetchall()
        return [
            {
                "trace_id": r[0], "source": r[1], "source_id": r[2],
                "user_id": r[3], "project": r[4], "input_text": r[5],
                "created_at": r[6], "completed_at": r[7], "status": r[8],
                "total_input_tokens": r[9], "total_output_tokens": r[10],
                "total_cost_usd": r[11],
            }
            for r in rows
        ]

    async def finalize_trace(
        self,
        trace_id: str,
        status: str,
        completed_at: str,
        total_input_tokens: int = 0,
        total_output_tokens: int = 0,
        total_cost_usd: float = 0.0,
    ) -> None:
        await self._db.execute(
            "UPDATE traces SET status = ?, completed_at = ?,"
            " total_input_tokens = ?, total_output_tokens = ?, total_cost_usd = ?"
            " WHERE trace_id = ?",
            (status, completed_at, total_input_tokens, total_output_tokens,
             total_cost_usd, trace_id),
        )
        await self._db.commit()

    async def get_trace_stats(
        self, project: str | None = None, days: int = 30,
    ) -> dict:
        from datetime import datetime, timedelta
        cutoff = (datetime.now(UTC) - timedelta(days=days)).isoformat()
        project_clause = "AND project = ?" if project else ""
        params: list = [cutoff]
        if project:
            params.append(project)
        async with self._db.execute(
            f"SELECT COUNT(*), COALESCE(SUM(total_cost_usd), 0),"
            f" COALESCE(SUM(total_input_tokens), 0), COALESCE(SUM(total_output_tokens), 0)"
            f" FROM traces WHERE created_at >= ? {project_clause}",
            params,
        ) as cur:
            row = await cur.fetchone()
        return {
            "total_traces": row[0],
            "total_cost_usd": row[1],
            "total_input_tokens": row[2],
            "total_output_tokens": row[3],
        }

    # ------------------------------------------------------------------
    # Spans
    # ------------------------------------------------------------------

    async def create_span(
        self,
        span_id: str,
        trace_id: str,
        kind: str,
        started_at: str,
        parent_span_id: str | None = None,
        task_id: str | None = None,
        session_id: str | None = None,
        log_path: str | None = None,
        metadata: dict | None = None,
    ) -> None:
        meta_json = json.dumps(metadata) if metadata else None
        await self._db.execute(
            "INSERT INTO spans (span_id, trace_id, parent_span_id, kind, task_id,"
            " session_id, log_path, started_at, metadata)"
            " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (span_id, trace_id, parent_span_id, kind, task_id,
             session_id, log_path, started_at, meta_json),
        )
        await self._db.commit()

    async def finalize_span(
        self,
        span_id: str,
        ended_at: str,
        duration_ms: int | None = None,
        input_tokens: int = 0,
        output_tokens: int = 0,
        cost_usd: float = 0.0,
    ) -> None:
        await self._db.execute(
            "UPDATE spans SET ended_at = ?, duration_ms = ?,"
            " input_tokens = ?, output_tokens = ?, cost_usd = ?"
            " WHERE span_id = ?",
            (ended_at, duration_ms, input_tokens, output_tokens, cost_usd, span_id),
        )
        await self._db.commit()

    async def list_spans(self, trace_id: str) -> list[dict]:
        async with self._db.execute(
            "SELECT span_id, trace_id, parent_span_id, kind, task_id, session_id,"
            " log_path, started_at, ended_at, duration_ms, input_tokens,"
            " output_tokens, cost_usd, metadata FROM spans"
            " WHERE trace_id = ? ORDER BY started_at ASC",
            (trace_id,),
        ) as cur:
            rows = await cur.fetchall()
        return [
            {
                "span_id": r[0], "trace_id": r[1], "parent_span_id": r[2],
                "kind": r[3], "task_id": r[4], "session_id": r[5],
                "log_path": r[6], "started_at": r[7], "ended_at": r[8],
                "duration_ms": r[9], "input_tokens": r[10], "output_tokens": r[11],
                "cost_usd": r[12], "metadata": json.loads(r[13]) if r[13] else None,
            }
            for r in rows
        ]

    async def get_execution_span_for_task(self, task_id: str) -> dict | None:
        """Return the execution span for a task (if any)."""
        async with self._db.execute(
            "SELECT span_id, trace_id FROM spans"
            " WHERE task_id = ? AND kind = 'execution'"
            " ORDER BY started_at DESC LIMIT 1",
            (task_id,),
        ) as cur:
            row = await cur.fetchone()
        if row is None:
            return None
        return {"span_id": row[0], "trace_id": row[1]}

    # ------------------------------------------------------------------
    # Execution Turns
    # ------------------------------------------------------------------

    async def insert_execution_turn(
        self,
        span_id: str,
        turn_number: int | None = None,
        role: str | None = None,
        content_type: str | None = None,
        tool_name: str | None = None,
        input_tokens: int = 0,
        output_tokens: int = 0,
        cache_read_tokens: int = 0,
        cost_usd: float = 0.0,
        content_preview: str | None = None,
        timestamp: str | None = None,
    ) -> None:
        await self._db.execute(
            "INSERT INTO execution_turns (span_id, turn_number, role, content_type,"
            " tool_name, input_tokens, output_tokens, cache_read_tokens, cost_usd,"
            " content_preview, timestamp) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (span_id, turn_number, role, content_type, tool_name,
             input_tokens, output_tokens, cache_read_tokens, cost_usd,
             content_preview, timestamp),
        )
        # No commit per turn — batched for performance

    async def list_execution_turns(
        self, span_id: str | None = None, trace_id: str | None = None,
    ) -> list[dict]:
        if span_id:
            query = (
                "SELECT id, span_id, turn_number, role, content_type, tool_name,"
                " input_tokens, output_tokens, cache_read_tokens, cost_usd,"
                " content_preview, timestamp FROM execution_turns"
                " WHERE span_id = ? ORDER BY id ASC"
            )
            params = [span_id]
        elif trace_id:
            query = (
                "SELECT et.id, et.span_id, et.turn_number, et.role, et.content_type,"
                " et.tool_name, et.input_tokens, et.output_tokens, et.cache_read_tokens,"
                " et.cost_usd, et.content_preview, et.timestamp"
                " FROM execution_turns et"
                " JOIN spans s ON et.span_id = s.span_id"
                " WHERE s.trace_id = ? ORDER BY et.id ASC"
            )
            params = [trace_id]
        else:
            return []
        async with self._db.execute(query, params) as cur:
            rows = await cur.fetchall()
        return [
            {
                "id": r[0], "span_id": r[1], "turn_number": r[2], "role": r[3],
                "content_type": r[4], "tool_name": r[5], "input_tokens": r[6],
                "output_tokens": r[7], "cache_read_tokens": r[8], "cost_usd": r[9],
                "content_preview": r[10], "timestamp": r[11],
            }
            for r in rows
        ]

    async def commit(self) -> None:
        """Explicit commit — used after batched inserts (e.g. execution turns)."""
        await self._db.commit()
