"""Schema DDL and the schema-time migrations (tables, columns, legacy JSON rows)."""

import json

import asyncpg
import structlog

from .rows import _deserialize, _message_key

_DDL = """
CREATE TABLE IF NOT EXISTS sessions (
    id                  TEXT PRIMARY KEY,
    profile_id          TEXT NOT NULL,
    user_id             TEXT REFERENCES navi_users(id) ON DELETE SET NULL,
    messages            TEXT NOT NULL DEFAULT '[]',
    context             TEXT NOT NULL DEFAULT '',
    pinned              BOOLEAN NOT NULL DEFAULT FALSE,
    special             BOOLEAN NOT NULL DEFAULT FALSE,
    created_at          TIMESTAMPTZ NOT NULL,
    last_active         TIMESTAMPTZ NOT NULL,
    context_token_count INTEGER NOT NULL DEFAULT 0,
    name                TEXT,
    session_metadata    TEXT NOT NULL DEFAULT '{}'
)
"""

_MIGRATE = """
ALTER TABLE sessions ADD COLUMN IF NOT EXISTS name TEXT;
ALTER TABLE sessions ADD COLUMN IF NOT EXISTS planning_logs TEXT NOT NULL DEFAULT '[]';
ALTER TABLE sessions ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES navi_users(id) ON DELETE SET NULL;
ALTER TABLE sessions ADD COLUMN IF NOT EXISTS next_sequence INTEGER NOT NULL DEFAULT 0;
ALTER TABLE sessions ADD COLUMN IF NOT EXISTS archive_threshold INTEGER NOT NULL DEFAULT 0;
ALTER TABLE sessions ADD COLUMN IF NOT EXISTS session_metadata TEXT NOT NULL DEFAULT '{}';
ALTER TABLE sessions ADD COLUMN IF NOT EXISTS special BOOLEAN NOT NULL DEFAULT FALSE
"""

# Reaction threads are found by (owner, conversation key) among special sessions
# only — a tiny slice of the table, so the index is partial. The expression is the
# exact one the lookup uses; `session_metadata` is a TEXT column holding JSON.
_SYNAPSE_INDEX_DDL = """
CREATE INDEX IF NOT EXISTS idx_sessions_synapse_session_key
    ON sessions (user_id, (session_metadata::jsonb -> 'synapse' ->> 'session_key'))
    WHERE special
"""

_SESSION_MESSAGES_DDL = """
CREATE TABLE IF NOT EXISTS session_messages (
    id              SERIAL PRIMARY KEY,
    session_id      TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
    sequence_number INT NOT NULL,
    role            TEXT NOT NULL,
    content         TEXT,
    images          TEXT,                       -- JSON array of base64 strings
    tool_calls      TEXT,                       -- JSON
    tool_call_id    TEXT,
    name            TEXT,
    created_at      TIMESTAMPTZ,
    is_summary      BOOLEAN NOT NULL DEFAULT FALSE,
    thinking        TEXT,
    is_plan         BOOLEAN NOT NULL DEFAULT FALSE,
    is_compression  BOOLEAN NOT NULL DEFAULT FALSE,
    is_context      BOOLEAN NOT NULL DEFAULT TRUE,
    is_display      BOOLEAN NOT NULL DEFAULT TRUE,
    elapsed_seconds FLOAT,
    tool_call_count INT,
    token_count     INT,
    files           TEXT,                       -- JSON
    metadata        TEXT,                       -- JSON
    is_recall       BOOLEAN NOT NULL DEFAULT FALSE,
    UNIQUE(session_id, sequence_number)
);

CREATE INDEX IF NOT EXISTS idx_session_messages_session_seq ON session_messages(session_id, sequence_number);
CREATE INDEX IF NOT EXISTS idx_session_messages_context ON session_messages(session_id, is_context, sequence_number);
"""

_SESSION_IMAGES_DDL = """
CREATE TABLE IF NOT EXISTS session_images (
    id          SERIAL PRIMARY KEY,
    session_id  TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
    message_id  INT REFERENCES session_messages(id) ON DELETE CASCADE,
    base64      TEXT NOT NULL,
    filename    TEXT
);

CREATE INDEX IF NOT EXISTS idx_session_images_session ON session_images(session_id);
CREATE INDEX IF NOT EXISTS idx_session_images_message ON session_images(message_id);
"""

_SESSION_ARCHIVE_DDL = """
CREATE TABLE IF NOT EXISTS session_messages_archive (
    id              SERIAL PRIMARY KEY,
    session_id      TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
    sequence_number INT NOT NULL,
    role            TEXT NOT NULL,
    content         TEXT,
    images          TEXT,
    tool_calls      TEXT,
    tool_call_id    TEXT,
    name            TEXT,
    created_at      TIMESTAMPTZ,
    is_summary      BOOLEAN NOT NULL DEFAULT FALSE,
    thinking        TEXT,
    is_plan         BOOLEAN NOT NULL DEFAULT FALSE,
    is_compression  BOOLEAN NOT NULL DEFAULT FALSE,
    is_context      BOOLEAN NOT NULL DEFAULT TRUE,
    is_display      BOOLEAN NOT NULL DEFAULT TRUE,
    elapsed_seconds FLOAT,
    tool_call_count INT,
    token_count     INT,
    files           TEXT,
    metadata        TEXT,
    is_recall       BOOLEAN NOT NULL DEFAULT FALSE,
    UNIQUE(session_id, sequence_number)
);
CREATE INDEX IF NOT EXISTS idx_session_messages_archive_session_seq ON session_messages_archive(session_id, sequence_number);
"""

# Session search runs `content ILIKE '%q%'` over both message tables. Without a
# trigram index that is a sequential scan of every message (measured ~350 ms on
# a 29k-row dev clone); with one PostgreSQL can use the GIN index. Like the
# memory tables, the index is created only when pg_trgm is already installed:
# CREATE EXTENSION needs privileges the service user may not have.
_SEARCH_INDEX_DDL = """
DO $$
BEGIN
    IF EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm') THEN
        CREATE INDEX IF NOT EXISTS idx_session_messages_content_trgm
            ON session_messages USING gin (content gin_trgm_ops) WHERE is_display;
        CREATE INDEX IF NOT EXISTS idx_session_messages_archive_content_trgm
            ON session_messages_archive USING gin (content gin_trgm_ops) WHERE is_display;
    END IF;
END $$;
"""

_SESSION_SHARES_DDL = """
CREATE TABLE IF NOT EXISTS session_shares (
    session_id  TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE,
    token       TEXT NOT NULL UNIQUE,
    owner_id    TEXT,
    created_at  TIMESTAMPTZ NOT NULL DEFAULT now()
);
"""



async def _ensure_normalized_tables(conn: asyncpg.Connection) -> None:
    await conn.execute(_SESSION_MESSAGES_DDL)
    await conn.execute(_SESSION_IMAGES_DDL)
    await conn.execute(_SESSION_ARCHIVE_DDL)
    await conn.execute(_SESSION_SHARES_DDL)
    # Best-effort for the same reason as the index below: search still works
    # without it, only slower, and a missing extension/privilege must not take
    # startup down.
    try:
        await conn.execute(_SEARCH_INDEX_DDL)
    except Exception as e:  # noqa: BLE001 — never block startup over an index
        structlog.get_logger().warning(
            "pg_session_store.search_index_failed", error=str(e)
        )
    # Best-effort: a row whose session_metadata is not valid JSON would make the
    # expression index fail and take the whole schema init down with it. The
    # lookup still works without the index, so a failure is a slowdown, not an
    # outage — log it and carry on.
    try:
        await conn.execute(_SYNAPSE_INDEX_DDL)
    except Exception as e:  # noqa: BLE001 — never block startup over an index
        structlog.get_logger().warning(
            "pg_session_store.synapse_index_failed", error=str(e)
        )


async def _migrate_to_normalized(conn: asyncpg.Connection) -> None:
    """One-shot migration: copy existing JSON messages/context into session_messages."""
    migrated = await conn.fetchval("SELECT COUNT(*) FROM session_messages")
    if migrated:
        return

    rows = await conn.fetch("SELECT id, messages, context FROM sessions")
    for row in rows:
        session_id = row["id"]
        messages = _deserialize(row["messages"] or "[]")
        context = _deserialize(row["context"] or "")
        if not context and messages:
            context = list(messages)

        context_set = {_message_key(m) for m in context}

        for seq, m in enumerate(messages):
            is_ctx = _message_key(m) in context_set
            await conn.execute(
                """
                INSERT INTO session_messages
                (session_id, sequence_number, role, content, images, tool_calls, tool_call_id, name,
                 created_at, is_summary, thinking, is_plan, is_compression, is_context, is_display,
                 elapsed_seconds, tool_call_count, token_count, files, metadata, is_recall)
                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21)
                ON CONFLICT (session_id, sequence_number) DO NOTHING
                """,
                session_id,
                seq,
                m.role,
                m.content,
                json.dumps(m.images, ensure_ascii=False) if m.images else None,
                json.dumps([tc.model_dump(mode="json") for tc in m.tool_calls], ensure_ascii=False) if m.tool_calls else None,
                m.tool_call_id,
                m.name,
                m.created_at,
                m.is_summary,
                m.thinking,
                m.is_plan,
                m.is_compression,
                is_ctx,
                True,
                m.elapsed_seconds,
                m.tool_call_count,
                m.token_count,
                json.dumps(m.files, ensure_ascii=False) if m.files else None,
                json.dumps(m.metadata, ensure_ascii=False) if m.metadata else None,
                m.is_recall,
            )