from __future__ import annotations

import os
import sqlite3
from contextlib import contextmanager
from dataclasses import dataclass
import hashlib
import secrets
from typing import Any, Iterator


def _utc_now_iso() -> str:
    # Keep it simple; wall-clock ordering only.
    import datetime as _dt

    return _dt.datetime.utcnow().replace(microsecond=0).isoformat() + "Z"


@dataclass(frozen=True)
class Db:
    db_path: str

    def _hash_token(self, token: str) -> str:
        # Internal tokens only; store hash to avoid plaintext at rest.
        return hashlib.sha256(token.encode("utf-8")).hexdigest()

    def _connect(self) -> sqlite3.Connection:
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        conn = sqlite3.connect(self.db_path, timeout=30, isolation_level=None)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL;")
        conn.execute("PRAGMA foreign_keys=ON;")
        return conn

    def init(self) -> None:
        with self._connect() as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS users (
                  user_id TEXT PRIMARY KEY,
                  display_name TEXT,
                  state TEXT NOT NULL,
                  created_at TEXT NOT NULL,
                  updated_at TEXT NOT NULL
                )
                """
            )
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS api_tokens (
                  token_hash TEXT PRIMARY KEY,
                  user_id TEXT NOT NULL,
                  token_plain TEXT NOT NULL,
                  created_at TEXT NOT NULL,
                  last_used_at TEXT,
                  FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
                )
                """
            )
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS tasks (
                  task_id TEXT PRIMARY KEY,
                  workload TEXT NOT NULL,
                  state TEXT NOT NULL,
                  jobspec_yaml TEXT NOT NULL,
                  nnodes INTEGER NOT NULL,
                  n_gpus_per_node INTEGER NOT NULL,
                  created_at TEXT NOT NULL,
                  updated_at TEXT NOT NULL,
                  next_run_at TEXT,
                  error_summary TEXT,
                  latest_attempt_no INTEGER NOT NULL DEFAULT 0
                )
                """
            )
            # Best-effort: add user_id column for forward compatibility (v2.5+).
            # v2.0 tables may exist without it; SQLite supports ADD COLUMN.
            try:
                conn.execute("ALTER TABLE tasks ADD COLUMN user_id TEXT")
            except sqlite3.OperationalError:
                pass
            # Best-effort: add missing columns (forward compatibility).
            try:
                conn.execute("ALTER TABLE users ADD COLUMN updated_at TEXT")
            except sqlite3.OperationalError:
                pass
            try:
                conn.execute("ALTER TABLE api_tokens ADD COLUMN token_plain TEXT")
            except sqlite3.OperationalError:
                pass
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS attempts (
                  task_id TEXT NOT NULL,
                  attempt_no INTEGER NOT NULL,
                  ray_submission_id TEXT NOT NULL UNIQUE,
                  ray_status TEXT,
                  failure_kind TEXT,
                  message TEXT,
                  start_time TEXT NOT NULL,
                  end_time TEXT,
                  PRIMARY KEY (task_id, attempt_no),
                  FOREIGN KEY (task_id) REFERENCES tasks(task_id) ON DELETE CASCADE
                )
                """
            )
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS events (
                  id INTEGER PRIMARY KEY AUTOINCREMENT,
                  task_id TEXT,
                  ts TEXT NOT NULL,
                  event_type TEXT NOT NULL,
                  payload_json TEXT,
                  FOREIGN KEY (task_id) REFERENCES tasks(task_id) ON DELETE CASCADE
                )
                """
            )
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS serve_models (
                  model_key TEXT PRIMARY KEY,
                  user_id TEXT NOT NULL,
                  model_id_suffix TEXT NOT NULL,
                  model_id_prefix TEXT NOT NULL,
                  model_id TEXT NOT NULL,
                  model_source TEXT NOT NULL,
                  num_replicas INTEGER NOT NULL,
                  gpus_per_replica INTEGER NOT NULL,
                  engine_kwargs_json TEXT,
                  state TEXT NOT NULL,
                  spec_yaml TEXT NOT NULL,
                  resolved_spec_yaml TEXT NOT NULL,
                  error_summary TEXT,
                  created_at TEXT NOT NULL,
                  updated_at TEXT NOT NULL,
                  deleted_at TEXT
                )
                """
            )
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS serve_events (
                  id INTEGER PRIMARY KEY AUTOINCREMENT,
                  model_key TEXT NOT NULL,
                  ts TEXT NOT NULL,
                  event_type TEXT NOT NULL,
                  payload_json TEXT,
                  FOREIGN KEY (model_key) REFERENCES serve_models(model_key) ON DELETE CASCADE
                )
                """
            )
            conn.execute("CREATE INDEX IF NOT EXISTS idx_serve_models_user ON serve_models(user_id)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_serve_models_state ON serve_models(state)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_serve_events_model ON serve_events(model_key)")

    @contextmanager
    def tx(self) -> Iterator[sqlite3.Connection]:
        conn = self._connect()
        try:
            conn.execute("BEGIN IMMEDIATE;")
            yield conn
            conn.execute("COMMIT;")
        except Exception:
            conn.execute("ROLLBACK;")
            raise
        finally:
            conn.close()

    def create_task(self, *, task_id: str, workload: str, jobspec_yaml: str, nnodes: int, n_gpus_per_node: int) -> dict[str, Any]:
        now = _utc_now_iso()
        with self.tx() as conn:
            conn.execute(
                """
                INSERT INTO tasks (task_id, workload, state, jobspec_yaml, nnodes, n_gpus_per_node, created_at, updated_at, user_id)
                VALUES (?, ?, 'QUEUED', ?, ?, ?, ?, ?, ?)
                """,
                (task_id, workload, jobspec_yaml, nnodes, n_gpus_per_node, now, now, None),
            )
            conn.execute(
                "INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, 'TASK_CREATED', ?)",
                (task_id, now, None),
            )
            row = conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
            return dict(row) if row else {}

    def create_task_v25(
        self,
        *,
        task_id: str,
        user_id: str,
        workload: str,
        jobspec_yaml: str,
        nnodes: int,
        n_gpus_per_node: int,
    ) -> dict[str, Any]:
        now = _utc_now_iso()
        with self.tx() as conn:
            conn.execute(
                """
                INSERT INTO tasks (task_id, user_id, workload, state, jobspec_yaml, nnodes, n_gpus_per_node, created_at, updated_at)
                VALUES (?, ?, ?, 'QUEUED', ?, ?, ?, ?, ?)
                """,
                (task_id, user_id, workload, jobspec_yaml, nnodes, n_gpus_per_node, now, now),
            )
            conn.execute(
                "INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, 'TASK_CREATED', ?)",
                (task_id, now, None),
            )
            row = conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
            return dict(row) if row else {}

    def create_user(self, *, user_id: str, display_name: str | None = None) -> dict[str, Any]:
        now = _utc_now_iso()
        with self.tx() as conn:
            conn.execute(
                """
                INSERT INTO users (user_id, display_name, state, created_at, updated_at)
                VALUES (?, ?, 'ACTIVE', ?, ?)
                """,
                (user_id, display_name, now, now),
            )
            conn.execute(
                "INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (NULL, ?, 'USER_CREATED', ?)",
                (now, user_id),
            )
            row = conn.execute("SELECT * FROM users WHERE user_id = ?", (user_id,)).fetchone()
            return dict(row) if row else {}

    def disable_user(self, *, user_id: str) -> None:
        now = _utc_now_iso()
        with self.tx() as conn:
            conn.execute("UPDATE users SET state = 'DISABLED', updated_at = ? WHERE user_id = ?", (now, user_id))
            conn.execute(
                "INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (NULL, ?, 'USER_DISABLED', ?)",
                (now, user_id),
            )

    def get_user(self, user_id: str) -> dict[str, Any] | None:
        with self._connect() as conn:
            row = conn.execute("SELECT * FROM users WHERE user_id = ?", (user_id,)).fetchone()
            return dict(row) if row else None

    def issue_token(self, *, user_id: str) -> str:
        # Returns plaintext token once.
        # Note: For v3.0 WebUI admin convenience we persist the plaintext token so admins
        # can re-copy it later. This is acceptable only for internal/dev deployments.
        now = _utc_now_iso()
        token = f"mvp_u_{user_id}_{secrets.token_urlsafe(18)}"
        token_hash = self._hash_token(token)
        with self.tx() as conn:
            conn.execute(
                "INSERT INTO api_tokens (token_hash, user_id, token_plain, created_at) VALUES (?, ?, ?, ?)",
                (token_hash, user_id, token, now),
            )
            conn.execute("UPDATE users SET updated_at = ? WHERE user_id = ?", (now, user_id))
            conn.execute(
                "INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (NULL, ?, 'TOKEN_ISSUED', ?)",
                (now, user_id),
            )
        return token

    def resolve_token(self, token: str) -> dict[str, Any] | None:
        token_hash = self._hash_token(token)
        with self.tx() as conn:
            row = conn.execute(
                """
                SELECT t.user_id, u.state
                FROM api_tokens t
                JOIN users u ON u.user_id = t.user_id
                WHERE t.token_hash = ?
                """,
                (token_hash,),
            ).fetchone()
            if not row:
                return None
            now = _utc_now_iso()
            conn.execute("UPDATE api_tokens SET last_used_at = ? WHERE token_hash = ?", (now, token_hash))
            conn.execute("UPDATE users SET updated_at = ? WHERE user_id = ?", (now, str(row["user_id"])))
            return {"user_id": row["user_id"], "state": row["state"]}

    def list_users(self, limit: int = 200) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT
                  u.user_id,
                  u.display_name,
                  u.state,
                  u.created_at,
                  u.updated_at,
                  (
                    SELECT t.token_plain
                    FROM api_tokens t
                    WHERE t.user_id = u.user_id
                    ORDER BY t.created_at DESC
                    LIMIT 1
                  ) AS token,
                  (
                    SELECT t.created_at
                    FROM api_tokens t
                    WHERE t.user_id = u.user_id
                    ORDER BY t.created_at DESC
                    LIMIT 1
                  ) AS token_created_at,
                  (
                    SELECT t.last_used_at
                    FROM api_tokens t
                    WHERE t.user_id = u.user_id
                    ORDER BY t.created_at DESC
                    LIMIT 1
                  ) AS token_last_used_at
                FROM users u
                ORDER BY u.created_at DESC
                LIMIT ?
                """,
                (int(limit),),
            ).fetchall()
            return [dict(r) for r in rows]

    def get_task(self, task_id: str) -> dict[str, Any] | None:
        with self._connect() as conn:
            row = conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
            return dict(row) if row else None

    def list_attempts(self, task_id: str) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT * FROM attempts WHERE task_id = ? ORDER BY attempt_no ASC", (task_id,)
            ).fetchall()
            return [dict(r) for r in rows]

    def list_queue(self, *, user_id: str | None = None) -> dict[str, list[dict[str, Any]]]:
        with self._connect() as conn:
            params: list[Any] = []
            user_filter_sql = ""
            if user_id is not None:
                user_filter_sql = " AND user_id = ?"
                params = [user_id]

            pending_sql = (
                "SELECT task_id, workload, state, nnodes, n_gpus_per_node, next_run_at, created_at, updated_at "
                "FROM tasks "
                "WHERE state IN ('QUEUED','PENDING_RESOURCES')"
                f"{user_filter_sql} "
                "ORDER BY created_at ASC "
                "LIMIT 200"
            )
            running_sql = (
                "SELECT task_id, workload, state, nnodes, n_gpus_per_node, latest_attempt_no, created_at, updated_at "
                "FROM tasks "
                "WHERE state IN ('SUBMITTING','SUBMITTED','RUNNING')"
                f"{user_filter_sql} "
                "ORDER BY updated_at ASC "
                "LIMIT 200"
            )

            pending = conn.execute(pending_sql, tuple(params)).fetchall()
            running = conn.execute(running_sql, tuple(params)).fetchall()
            return {"pending": [dict(r) for r in pending], "running": [dict(r) for r in running]}

    def list_tasks(
        self,
        *,
        user_id: str | None = None,
        states: list[str] | None = None,
        limit: int = 200,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        """
        Returns recent tasks (including terminal tasks).
        """
        with self._connect() as conn:
            params: list[Any] = []
            where_clauses: list[str] = []
            if user_id is not None:
                where_clauses.append("user_id = ?")
                params.append(user_id)
            if states:
                placeholders = ",".join(["?"] * len(states))
                where_clauses.append(f"state IN ({placeholders})")
                params.extend(states)
            where_sql = f" WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
            sql = (
                "SELECT task_id, user_id, workload, state, nnodes, n_gpus_per_node, latest_attempt_no, created_at, updated_at, error_summary "
                "FROM tasks"
                f"{where_sql} "
                "ORDER BY created_at DESC "
                "LIMIT ? OFFSET ?"
            )
            params.append(int(limit))
            params.append(max(0, int(offset)))
            rows = conn.execute(sql, tuple(params)).fetchall()
            return [dict(r) for r in rows]

    def count_running(self) -> int:
        with self._connect() as conn:
            row = conn.execute(
                "SELECT COUNT(1) AS n FROM tasks WHERE state IN ('SUBMITTING','SUBMITTED','RUNNING')"
            ).fetchone()
            return int(row["n"]) if row else 0

    def list_active_tasks(self, limit: int = 50) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT * FROM tasks WHERE state IN ('SUBMITTING','SUBMITTED','RUNNING') ORDER BY updated_at ASC LIMIT ?",
                (int(limit),),
            ).fetchall()
            return [dict(r) for r in rows]

    def pick_next_runnable_task(self) -> dict[str, Any] | None:
        now = _utc_now_iso()
        with self._connect() as conn:
            row = conn.execute(
                """
                SELECT *
                FROM tasks
                WHERE state IN ('QUEUED','PENDING_RESOURCES')
                  AND (next_run_at IS NULL OR next_run_at <= ?)
                ORDER BY created_at ASC
                LIMIT 1
                """,
                (now,),
            ).fetchone()
            return dict(row) if row else None

    def set_task_state(
        self,
        *,
        task_id: str,
        state: str,
        error_summary: str | None = None,
        next_run_at: str | None = None,
        latest_attempt_no: int | None = None,
        event_type: str = "STATE_UPDATE",
        payload_json: str | None = None,
    ) -> None:
        now = _utc_now_iso()
        with self.tx() as conn:
            sets = ["state = ?", "updated_at = ?"]
            params: list[Any] = [state, now]
            if error_summary is not None:
                sets.append("error_summary = ?")
                params.append(error_summary)
            if next_run_at is not None:
                sets.append("next_run_at = ?")
                params.append(next_run_at)
            if latest_attempt_no is not None:
                sets.append("latest_attempt_no = ?")
                params.append(int(latest_attempt_no))
            params.append(task_id)
            conn.execute(f"UPDATE tasks SET {', '.join(sets)} WHERE task_id = ?", tuple(params))
            conn.execute(
                "INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, ?, ?)",
                (task_id, now, event_type, payload_json),
            )

    def create_attempt(self, *, task_id: str, attempt_no: int, ray_submission_id: str) -> None:
        now = _utc_now_iso()
        with self.tx() as conn:
            conn.execute(
                """
                INSERT INTO attempts (task_id, attempt_no, ray_submission_id, ray_status, start_time)
                VALUES (?, ?, ?, ?, ?)
                """,
                (task_id, attempt_no, ray_submission_id, "SUBMITTING", now),
            )
            conn.execute(
                "INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, 'ATTEMPT_CREATED', ?)",
                (task_id, now, None),
            )

    def update_attempt(
        self,
        *,
        task_id: str,
        attempt_no: int,
        ray_status: str | None = None,
        failure_kind: str | None = None,
        message: str | None = None,
        end_time: str | None = None,
    ) -> None:
        now = _utc_now_iso()
        with self.tx() as conn:
            sets = []
            params: list[Any] = []
            if ray_status is not None:
                sets.append("ray_status = ?")
                params.append(ray_status)
            if failure_kind is not None:
                sets.append("failure_kind = ?")
                params.append(failure_kind)
            if message is not None:
                sets.append("message = ?")
                params.append(message)
            if end_time is not None:
                sets.append("end_time = ?")
                params.append(end_time)
            if not sets:
                return
            params.extend([task_id, attempt_no])
            conn.execute(
                f"UPDATE attempts SET {', '.join(sets)} WHERE task_id = ? AND attempt_no = ?",
                tuple(params),
            )
            conn.execute(
                "INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, 'ATTEMPT_UPDATE', ?)",
                (task_id, now, None),
            )

    def list_ended_attempts_before(self, *, end_time_le: str, limit: int = 2000) -> list[dict[str, Any]]:
        """
        Used by the jobs janitor:
        - Only considers tasks with a non-null user_id (v2.5+).
        - Returns attempts that have end_time <= end_time_le.
        """
        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT t.task_id, t.user_id, a.ray_submission_id, a.end_time
                FROM attempts a
                JOIN tasks t ON t.task_id = a.task_id
                WHERE t.user_id IS NOT NULL
                  AND a.end_time IS NOT NULL
                  AND a.end_time <= ?
                ORDER BY a.end_time ASC
                LIMIT ?
                """,
                (str(end_time_le), int(limit)),
            ).fetchall()
            return [dict(r) for r in rows]

    def create_serve_model(
        self,
        *,
        model_key: str,
        user_id: str,
        model_id_suffix: str,
        model_id_prefix: str,
        model_id: str,
        model_source: str,
        num_replicas: int,
        gpus_per_replica: int,
        spec_yaml: str,
        resolved_spec_yaml: str,
        engine_kwargs_json: str | None = None,
    ) -> dict[str, Any]:
        now = _utc_now_iso()
        with self.tx() as conn:
            conn.execute(
                """
                INSERT INTO serve_models (
                  model_key,
                  user_id,
                  model_id_suffix,
                  model_id_prefix,
                  model_id,
                  model_source,
                  num_replicas,
                  gpus_per_replica,
                  engine_kwargs_json,
                  state,
                  spec_yaml,
                  resolved_spec_yaml,
                  created_at,
                  updated_at
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'QUEUED', ?, ?, ?, ?)
                """,
                (
                    model_key,
                    user_id,
                    model_id_suffix,
                    model_id_prefix,
                    model_id,
                    model_source,
                    int(num_replicas),
                    int(gpus_per_replica),
                    engine_kwargs_json,
                    spec_yaml,
                    resolved_spec_yaml,
                    now,
                    now,
                ),
            )
            conn.execute(
                "INSERT INTO serve_events (model_key, ts, event_type, payload_json) VALUES (?, ?, 'SERVE_MODEL_CREATED', ?)",
                (model_key, now, None),
            )
            row = conn.execute("SELECT * FROM serve_models WHERE model_key = ?", (model_key,)).fetchone()
            return dict(row) if row else {}

    def list_serve_models(
        self,
        *,
        user_id: str,
        include_deleted: bool = False,
        limit: int = 200,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        with self._connect() as conn:
            where_sql = "WHERE user_id = ?"
            params: list[Any] = [user_id]
            if not include_deleted:
                where_sql += " AND deleted_at IS NULL"
            params.append(int(limit))
            params.append(max(0, int(offset)))
            rows = conn.execute(
                f"""
                SELECT
                  model_key,
                  user_id,
                  model_id_suffix,
                  model_id_prefix,
                  model_id,
                  model_source,
                  num_replicas,
                  gpus_per_replica,
                  engine_kwargs_json,
                  state,
                  error_summary,
                  created_at,
                  updated_at,
                  deleted_at
                FROM serve_models
                {where_sql}
                ORDER BY created_at DESC, model_key DESC
                LIMIT ? OFFSET ?
                """,
                tuple(params),
            ).fetchall()
            return [dict(r) for r in rows]

    def list_all_serve_models(
        self,
        *,
        include_deleted: bool = False,
        limit: int = 2000,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        with self._connect() as conn:
            where_sql = ""
            if not include_deleted:
                where_sql = "WHERE deleted_at IS NULL"
            rows = conn.execute(
                f"""
                SELECT
                  model_key,
                  user_id,
                  model_id_suffix,
                  model_id_prefix,
                  model_id,
                  model_source,
                  num_replicas,
                  gpus_per_replica,
                  engine_kwargs_json,
                  state,
                  error_summary,
                  spec_yaml,
                  resolved_spec_yaml,
                  created_at,
                  updated_at,
                  deleted_at
                FROM serve_models
                {where_sql}
                ORDER BY created_at ASC, model_key ASC
                LIMIT ? OFFSET ?
                """,
                (int(limit), max(0, int(offset))),
            ).fetchall()
            return [dict(r) for r in rows]

    def get_serve_model(self, model_key: str) -> dict[str, Any] | None:
        with self._connect() as conn:
            row = conn.execute("SELECT * FROM serve_models WHERE model_key = ?", (model_key,)).fetchone()
            return dict(row) if row else None

    def list_serve_events(self, model_key: str, limit: int = 200, offset: int = 0) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT id, model_key, ts, event_type, payload_json
                FROM serve_events
                WHERE model_key = ?
                ORDER BY id DESC
                LIMIT ? OFFSET ?
                """,
                (model_key, int(limit), max(0, int(offset))),
            ).fetchall()
            return [dict(r) for r in rows]

    def append_serve_event(self, *, model_key: str, event_type: str, payload_json: str | None = None) -> None:
        now = _utc_now_iso()
        with self.tx() as conn:
            conn.execute(
                "INSERT INTO serve_events (model_key, ts, event_type, payload_json) VALUES (?, ?, ?, ?)",
                (model_key, now, event_type, payload_json),
            )

    def set_serve_model_state(
        self,
        *,
        model_key: str,
        state: str,
        error_summary: str | None = None,
        event_type: str = "SERVE_STATE_UPDATE",
        payload_json: str | None = None,
    ) -> None:
        now = _utc_now_iso()
        with self.tx() as conn:
            sets = ["state = ?", "updated_at = ?"]
            params: list[Any] = [state, now]
            if error_summary is not None:
                sets.append("error_summary = ?")
                params.append(error_summary)
            if state == "DELETED":
                sets.append("deleted_at = ?")
                params.append(now)
            params.append(model_key)
            conn.execute(f"UPDATE serve_models SET {', '.join(sets)} WHERE model_key = ?", tuple(params))
            conn.execute(
                "INSERT INTO serve_events (model_key, ts, event_type, payload_json) VALUES (?, ?, ?, ?)",
                (model_key, now, event_type, payload_json),
            )

    def update_serve_model_num_replicas(self, *, model_key: str, num_replicas: int) -> None:
        if not isinstance(num_replicas, int) or num_replicas < 1:
            raise ValueError("num_replicas must be an integer >= 1")
        now = _utc_now_iso()
        with self.tx() as conn:
            conn.execute(
                """
                UPDATE serve_models
                SET num_replicas = ?, state = 'QUEUED', error_summary = NULL, updated_at = ?
                WHERE model_key = ?
                """,
                (int(num_replicas), now, model_key),
            )
            conn.execute(
                "INSERT INTO serve_events (model_key, ts, event_type, payload_json) VALUES (?, ?, 'SERVE_PATCH_NUM_REPLICAS', ?)",
                (model_key, now, str(num_replicas)),
            )

    def pick_next_runnable_serve_change(self) -> dict[str, Any] | None:
        """
        Returns the next serve model that needs reconciliation.

        Minimal state machine for now:
        - QUEUED: needs (re)apply
        - DELETING: needs removal
        """
        with self._connect() as conn:
            row = conn.execute(
                """
                SELECT *
                FROM serve_models
                WHERE deleted_at IS NULL
                  AND state IN ('QUEUED','DELETING')
                ORDER BY updated_at ASC
                LIMIT 1
                """
            ).fetchone()
            return dict(row) if row else None

    # Backward compatible naming (v3.8 docs originally used "upsert").
    def upsert_serve_model(self, **kwargs: Any) -> dict[str, Any]:
        return self.create_serve_model(**kwargs)
