"""Initial schema — memory, directive, session, LLM config, rate limit tables

Revision ID: 001
Revises:
Create Date: 2026-04-12
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision = '001'
down_revision = None
branch_labels = None
depends_on = None


def upgrade() -> None:
    # -----------------------------------------------------------------------
    # Tier 1: Conversation memory per user
    # Hard cap: 200 rows per user — enforced in application layer before insert
    # -----------------------------------------------------------------------
    op.execute("""
        CREATE TABLE IF NOT EXISTS ab_conversation_memory (
            id SERIAL PRIMARY KEY,
            user_id INTEGER NOT NULL,
            role VARCHAR(16) NOT NULL,
            content TEXT NOT NULL,
            directive_id VARCHAR(64),
            is_summary BOOLEAN DEFAULT FALSE,
            created_at TIMESTAMP DEFAULT NOW()
        )
    """)
    op.execute("""
        CREATE INDEX IF NOT EXISTS idx_conv_user
        ON ab_conversation_memory(user_id, created_at DESC)
    """)

    # -----------------------------------------------------------------------
    # Tier 2: Operational memory — 90-day TTL
    # -----------------------------------------------------------------------
    op.execute("""
        CREATE TABLE IF NOT EXISTS ab_operational_memory (
            id SERIAL PRIMARY KEY,
            scope VARCHAR(64) NOT NULL,
            summary TEXT NOT NULL,
            raw_data JSONB,
            source_directive_id VARCHAR(64),
            expires_at TIMESTAMP,
            created_at TIMESTAMP DEFAULT NOW()
        )
    """)
    op.execute("""
        CREATE INDEX IF NOT EXISTS idx_op_scope
        ON ab_operational_memory(scope, created_at DESC)
    """)

    # -----------------------------------------------------------------------
    # Tier 3: Long-term knowledge — permanent
    # -----------------------------------------------------------------------
    op.execute("""
        CREATE TABLE IF NOT EXISTS ab_knowledge_store (
            id SERIAL PRIMARY KEY,
            entity_type VARCHAR(32) NOT NULL,
            entity_key VARCHAR(128) NOT NULL,
            facts JSONB NOT NULL,
            updated_at TIMESTAMP DEFAULT NOW(),
            UNIQUE(entity_type, entity_key)
        )
    """)

    # -----------------------------------------------------------------------
    # Directive execution graph
    # -----------------------------------------------------------------------
    op.execute("""
        CREATE TABLE IF NOT EXISTS ab_directive_log (
            id SERIAL PRIMARY KEY,
            directive_id VARCHAR(64) UNIQUE NOT NULL,
            user_id INTEGER NOT NULL,
            channel_id INTEGER NOT NULL,
            raw_message TEXT NOT NULL,
            intent_summary TEXT,
            agents_involved JSONB DEFAULT '[]',
            peer_calls JSONB DEFAULT '[]',
            actions_taken JSONB DEFAULT '[]',
            escalations JSONB DEFAULT '[]',
            final_response TEXT,
            status VARCHAR(32) DEFAULT 'pending',
            started_at TIMESTAMP DEFAULT NOW(),
            completed_at TIMESTAMP,
            error TEXT
        )
    """)
    op.execute("""
        CREATE INDEX IF NOT EXISTS idx_dir_user
        ON ab_directive_log(user_id, started_at DESC)
    """)
    op.execute("""
        CREATE INDEX IF NOT EXISTS idx_dir_status
        ON ab_directive_log(status)
    """)

    # -----------------------------------------------------------------------
    # Per-agent session state during execution
    # -----------------------------------------------------------------------
    op.execute("""
        CREATE TABLE IF NOT EXISTS ab_agent_session (
            id SERIAL PRIMARY KEY,
            directive_id VARCHAR(64) NOT NULL,
            agent_name VARCHAR(64) NOT NULL,
            messages JSONB NOT NULL DEFAULT '[]',
            status VARCHAR(32) DEFAULT 'running',
            started_at TIMESTAMP DEFAULT NOW(),
            completed_at TIMESTAMP
        )
    """)

    # -----------------------------------------------------------------------
    # LLM backend config — runtime per-agent model routing
    # caller='__system__' stores the global privacy mode
    # -----------------------------------------------------------------------
    op.execute("""
        CREATE TABLE IF NOT EXISTS ab_llm_config (
            caller VARCHAR(64) PRIMARY KEY,
            backend VARCHAR(16) NOT NULL,
            set_by INTEGER,
            set_at TIMESTAMP DEFAULT NOW(),
            note TEXT
        )
    """)

    # Seed default privacy mode = 'local'
    op.execute("""
        INSERT INTO ab_llm_config (caller, backend, note)
        VALUES ('__system__', 'local', 'Default privacy mode — local Ollama only')
        ON CONFLICT (caller) DO NOTHING
    """)

    # -----------------------------------------------------------------------
    # Rate limiting
    # -----------------------------------------------------------------------
    op.execute("""
        CREATE TABLE IF NOT EXISTS ab_rate_limit (
            user_id INTEGER NOT NULL,
            window_start TIMESTAMP NOT NULL,
            request_count INTEGER DEFAULT 0,
            PRIMARY KEY (user_id, window_start)
        )
    """)


def downgrade() -> None:
    op.execute("DROP TABLE IF EXISTS ab_rate_limit")
    op.execute("DROP TABLE IF EXISTS ab_llm_config")
    op.execute("DROP TABLE IF EXISTS ab_agent_session")
    op.execute("DROP INDEX IF EXISTS idx_dir_status")
    op.execute("DROP INDEX IF EXISTS idx_dir_user")
    op.execute("DROP TABLE IF EXISTS ab_directive_log")
    op.execute("DROP TABLE IF EXISTS ab_knowledge_store")
    op.execute("DROP INDEX IF EXISTS idx_op_scope")
    op.execute("DROP TABLE IF EXISTS ab_operational_memory")
    op.execute("DROP INDEX IF EXISTS idx_conv_user")
    op.execute("DROP TABLE IF EXISTS ab_conversation_memory")