world/database.py

"""SQLite database connection, schema initialization, and migrations."""

from __future__ import annotations

import os
import sqlite3
from contextlib import contextmanager
from pathlib import Path
from typing import Generator

DB_PATH = Path(os.environ.get("TTR_DB_PATH", "data/world.db"))

SCHEMA_SQL = """
PRAGMA journal_mode=WAL;
PRAGMA foreign_keys=ON;
PRAGMA synchronous=NORMAL;

CREATE TABLE IF NOT EXISTS world_state (
    id                 INTEGER PRIMARY KEY DEFAULT 1,
    world_name         TEXT    NOT NULL DEFAULT 'Aether Garden',
    founding_date      TEXT    NOT NULL,
    current_day        INTEGER NOT NULL DEFAULT 1,
    total_entities     INTEGER NOT NULL DEFAULT 0,
    total_interactions INTEGER NOT NULL DEFAULT 0,
    total_events       INTEGER NOT NULL DEFAULT 0,
    last_simulation_run TEXT,
    active_world_event TEXT,
    updated_at         TEXT    NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS locations (
    id                   INTEGER PRIMARY KEY AUTOINCREMENT,
    slug                 TEXT    NOT NULL UNIQUE,
    name                 TEXT    NOT NULL,
    short_description    TEXT    NOT NULL,
    full_lore            TEXT    NOT NULL,
    vibe_tags            TEXT    NOT NULL,
    special_property     TEXT    NOT NULL,
    aesthetic_description TEXT   NOT NULL,
    map_x                REAL    NOT NULL,
    map_y                REAL    NOT NULL,
    glow_color           TEXT    NOT NULL,
    entity_count         INTEGER NOT NULL DEFAULT 0,
    interaction_multiplier REAL  NOT NULL DEFAULT 1.0
);

CREATE TABLE IF NOT EXISTS entities (
    id                  TEXT    NOT NULL PRIMARY KEY,
    name                TEXT    NOT NULL,
    display_name        TEXT    NOT NULL,
    type                TEXT    NOT NULL CHECK(type IN ('character','creature','object','place')),
    input_description   TEXT    NOT NULL,
    appearance          TEXT    NOT NULL,
    backstory           TEXT    NOT NULL,
    personality_traits  TEXT    NOT NULL,
    primary_goal        TEXT    NOT NULL,
    secondary_goal      TEXT    NOT NULL,
    primary_fear        TEXT    NOT NULL,
    speech_style        TEXT    NOT NULL,
    greeting            TEXT    NOT NULL,
    location_id         INTEGER NOT NULL REFERENCES locations(id),
    created_at          TEXT    NOT NULL DEFAULT (datetime('now')),
    created_by_session  TEXT,
    days_in_realm       INTEGER NOT NULL DEFAULT 0,
    memory_summary      TEXT    NOT NULL DEFAULT '',
    relationship_count  INTEGER NOT NULL DEFAULT 0,
    interaction_count   INTEGER NOT NULL DEFAULT 0,
    last_active         TEXT    NOT NULL DEFAULT (datetime('now')),
    status              TEXT    NOT NULL DEFAULT 'active'
                        CHECK(status IN ('active','dormant','legendary')),
    tags                TEXT    NOT NULL DEFAULT '[]',
    secret_name         TEXT,
    arrival_note        TEXT    NOT NULL,
    wisdom_unlocked     INTEGER NOT NULL DEFAULT 0
);

CREATE INDEX IF NOT EXISTS idx_entities_location ON entities(location_id);
CREATE INDEX IF NOT EXISTS idx_entities_status   ON entities(status);
CREATE INDEX IF NOT EXISTS idx_entities_created  ON entities(created_at);
CREATE INDEX IF NOT EXISTS idx_entities_active   ON entities(last_active);

CREATE TABLE IF NOT EXISTS relationships (
    id                   INTEGER PRIMARY KEY AUTOINCREMENT,
    entity_a_id          TEXT    NOT NULL REFERENCES entities(id),
    entity_b_id          TEXT    NOT NULL REFERENCES entities(id),
    relationship_type    TEXT    NOT NULL,
    description          TEXT    NOT NULL,
    strength             INTEGER NOT NULL DEFAULT 1 CHECK(strength BETWEEN 1 AND 5),
    formed_at            TEXT    NOT NULL DEFAULT (datetime('now')),
    formed_on_day        INTEGER NOT NULL DEFAULT 1,
    last_interaction     TEXT    NOT NULL DEFAULT (datetime('now')),
    UNIQUE(entity_a_id, entity_b_id),
    CHECK(entity_a_id < entity_b_id)
);

CREATE INDEX IF NOT EXISTS idx_rel_a ON relationships(entity_a_id);
CREATE INDEX IF NOT EXISTS idx_rel_b ON relationships(entity_b_id);

CREATE TABLE IF NOT EXISTS interactions (
    id               INTEGER PRIMARY KEY AUTOINCREMENT,
    entity_a_id      TEXT    NOT NULL REFERENCES entities(id),
    entity_b_id      TEXT    NOT NULL REFERENCES entities(id),
    location_id      INTEGER NOT NULL REFERENCES locations(id),
    interaction_type TEXT    NOT NULL,
    description      TEXT    NOT NULL,
    notable_outcome  TEXT,
    book_of_ages_entry TEXT  NOT NULL,
    timestamp        TEXT    NOT NULL DEFAULT (datetime('now')),
    world_day        INTEGER NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_int_timestamp ON interactions(timestamp);
CREATE INDEX IF NOT EXISTS idx_int_a         ON interactions(entity_a_id);
CREATE INDEX IF NOT EXISTS idx_int_b         ON interactions(entity_b_id);
CREATE INDEX IF NOT EXISTS idx_int_day       ON interactions(world_day);

CREATE TABLE IF NOT EXISTS world_events (
    id                INTEGER PRIMARY KEY AUTOINCREMENT,
    event_type        TEXT NOT NULL,
    title             TEXT NOT NULL,
    description       TEXT NOT NULL,
    affected_locations TEXT NOT NULL,
    entity_effect     TEXT NOT NULL,
    book_of_ages_entry TEXT NOT NULL,
    mystery_hook      TEXT,
    timestamp         TEXT NOT NULL DEFAULT (datetime('now')),
    world_day         INTEGER NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_events_day ON world_events(world_day);

CREATE TABLE IF NOT EXISTS book_of_ages (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    world_day   INTEGER NOT NULL,
    entry_type  TEXT    NOT NULL CHECK(entry_type IN
                ('arrival','interaction','world_event','milestone','dream_fragment')),
    title       TEXT,
    content     TEXT    NOT NULL,
    entity_ids  TEXT    NOT NULL DEFAULT '[]',
    location_id INTEGER REFERENCES locations(id),
    timestamp   TEXT    NOT NULL DEFAULT (datetime('now')),
    is_milestone INTEGER NOT NULL DEFAULT 0
);

CREATE INDEX IF NOT EXISTS idx_boa_day  ON book_of_ages(world_day);
CREATE INDEX IF NOT EXISTS idx_boa_type ON book_of_ages(entry_type);

CREATE TABLE IF NOT EXISTS sessions (
    session_id     TEXT NOT NULL PRIMARY KEY,
    last_creation  TEXT,
    creation_count INTEGER NOT NULL DEFAULT 0,
    first_seen     TEXT    NOT NULL DEFAULT (datetime('now')),
    last_seen      TEXT    NOT NULL DEFAULT (datetime('now'))
);
"""


def get_connection() -> sqlite3.Connection:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


@contextmanager
def db_session() -> Generator[sqlite3.Connection, None, None]:
    conn = get_connection()
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def init_database() -> None:
    with db_session() as conn:
        conn.executescript(SCHEMA_SQL)
        row = conn.execute("SELECT id FROM world_state WHERE id = 1").fetchone()
        if row is None:
            conn.execute(
                """
                INSERT INTO world_state (id, founding_date, current_day)
                VALUES (1, datetime('now'), 1)
                """
            )
        # Safe migrations (idempotent)
        _safe_migrate(conn)

    # Init extension tables
    from world.quests import init_quests_table
    from world.presence import init_presence_table
    init_quests_table()
    init_presence_table()


def ensure_database_ready() -> None:
    """Restore (if configured), create schema, and seed locations. Idempotent."""
    from persistence.backup import database_is_valid, restore_database
    from world.seed_data import seed_locations
    from world.locations import update_entity_counts

    DB_PATH.parent.mkdir(parents=True, exist_ok=True)

    if not database_is_valid():
        try:
            restore_database()
        except Exception as exc:
            print(f"[db] restore failed: {exc}")
        if DB_PATH.exists() and not database_is_valid():
            DB_PATH.unlink(missing_ok=True)

    init_database()
    with db_session() as conn:
        seed_locations(conn)
    update_entity_counts()

    if not database_is_valid():
        raise RuntimeError(f"Database failed to initialize at {DB_PATH.resolve()}")


def _safe_migrate(conn) -> None:
    """Add new columns/indexes that may not exist in older DB files."""
    migrations = [
        "ALTER TABLE entities ADD COLUMN portrait_url TEXT",
    ]
    for sql in migrations:
        try:
            conn.execute(sql)
        except Exception:
            pass  # Column already exists


def get_world_state() -> dict:
    with db_session() as conn:
        row = conn.execute("SELECT * FROM world_state WHERE id = 1").fetchone()
        return dict(row) if row else {}