Kansho/backend/db.py

"""SQLite persistence for the local frame. PostgreSQL remains the later target."""
from __future__ import annotations

import json
import os
import sqlite3
from contextlib import contextmanager
from pathlib import Path

DATA_DIR = Path(__file__).resolve().parent / "data"
SCHEMA_PATH = Path(__file__).resolve().parent / "schema.sql"
SEED_PATH = Path(__file__).resolve().parent / "config" / "platform_seed.json"
PROMPTS_SEED_PATH = Path(__file__).resolve().parent / "config" / "prompts.seed.json"
_env_db = os.environ.get("KANSHO_DB_PATH")
DB_PATH = Path(_env_db) if _env_db else DATA_DIR / "kansho.sqlite"

_PROFILE_COLUMNS = {
    "status": "TEXT NOT NULL DEFAULT 'active'",
    "tier_id": "TEXT NOT NULL DEFAULT 'local'",
}
_PROMPT_COLUMNS = {
    "description": "TEXT NOT NULL DEFAULT ''",
    "category": "TEXT NOT NULL DEFAULT 'uncategorized'",
    "stages_json": "TEXT",
    "graph_json": "TEXT",
    "output_format": "TEXT NOT NULL DEFAULT 'text'",
    "output_schema_json": "TEXT",
    "required_feature": "TEXT NOT NULL DEFAULT 'ai_calls'",
    "default_template": "TEXT NOT NULL DEFAULT ''",
    "sort_order": "INTEGER NOT NULL DEFAULT 0",
    "updated": "TEXT NOT NULL DEFAULT (datetime('now'))",
}


def _connect() -> sqlite3.Connection:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_PATH, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON")
    return conn


@contextmanager
def get_db():
    conn = _connect()
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def row_to_dict(row: sqlite3.Row | None) -> dict | None:
    if row is None:
        return None
    return dict(row)


def _column_names(conn: sqlite3.Connection, table: str) -> set[str]:
    rows = conn.execute(f"PRAGMA table_info({table})").fetchall()
    return {row["name"] for row in rows}


def _ensure_columns(conn: sqlite3.Connection, table: str, columns: dict[str, str]) -> None:
    existing = _column_names(conn, table)
    for name, ddl in columns.items():
        if name not in existing:
            conn.execute(f"ALTER TABLE {table} ADD COLUMN {name} {ddl}")


def _mark(conn: sqlite3.Connection, migration_id: str) -> None:
    conn.execute(
        "INSERT OR IGNORE INTO schema_migrations (id) VALUES (?)",
        (migration_id,),
    )


def _seed_platform(conn: sqlite3.Connection) -> None:
    seed = json.loads(SEED_PATH.read_text(encoding="utf-8"))
    for tier in seed.get("tiers", []):
        conn.execute(
            """
            INSERT OR IGNORE INTO tiers (id, name, description, sort_order)
            VALUES (?, ?, ?, ?)
            """,
            (tier["id"], tier["name"], tier.get("description") or "", tier.get("sort_order") or 0),
        )
    for feature in seed.get("features", []):
        conn.execute(
            """
            INSERT OR IGNORE INTO features
              (id, name, description, category, limit_type, default_limit, reset_period)
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
            (
                feature["id"],
                feature["name"],
                feature.get("description") or "",
                feature.get("category") or "platform",
                feature["limit_type"],
                feature.get("default_limit"),
                feature.get("reset_period") or "none",
            ),
        )
    for limit in seed.get("tier_limits", []):
        conn.execute(
            """
            INSERT OR IGNORE INTO tier_limits (tier_id, feature_id, limit_value)
            VALUES (?, ?, ?)
            """,
            (limit["tier_id"], limit["feature_id"], limit.get("limit_value")),
        )


def _seed_prompts(conn: sqlite3.Connection) -> None:
    """Prompts come from JSON/DB, never from Python string literals."""
    items = json.loads(PROMPTS_SEED_PATH.read_text(encoding="utf-8"))
    for item in items:
        conn.execute(
            """
            INSERT OR IGNORE INTO ai_prompts
              (id, slug, name, description, category, prompt_type, template,
               required_feature, is_system_default, default_template)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, 1, ?)
            """,
            (
                item["id"],
                item["slug"],
                item["name"],
                item.get("description") or "",
                item.get("category") or "uncategorized",
                item.get("prompt_type") or "base",
                item.get("template") or "",
                item.get("required_feature") or "ai_calls",
                item.get("template") or "",
            ),
        )


def init_db() -> None:
    schema = SCHEMA_PATH.read_text(encoding="utf-8")
    with get_db() as conn:
        conn.executescript(schema)
        _seed_platform(conn)
        _seed_prompts(conn)
        _ensure_columns(conn, "profiles", _PROFILE_COLUMNS)
        _ensure_columns(conn, "ai_prompts", _PROMPT_COLUMNS)
        _mark(conn, "001_frame")
        _mark(conn, "002_platform")
        _mark(conn, "003_dialogue_memory")