"""Add loop_directives and grukk_messages tables.

Revision ID: 0014
Revises: 0013
"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

revision: str = "0014"
down_revision: Union[str, None] = "82b950a0c033"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    conn = op.get_bind()

    # loop_directives table
    conn.execute(
        sa.text(
            """
            CREATE TABLE IF NOT EXISTS loop_directives (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                loop_id VARCHAR NOT NULL,
                directive_type VARCHAR NOT NULL CHECK (directive_type IN ('pause', 'resume', 'stop', 'redirect', 'answer')),
                payload JSONB NOT NULL DEFAULT '{}',
                status VARCHAR NOT NULL DEFAULT 'pending' CHECK (status IN ('pending', 'acknowledged', 'completed')),
                created_by VARCHAR NOT NULL,
                created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                acknowledged_at TIMESTAMPTZ
            )
            """
        )
    )
    conn.execute(
        sa.text(
            "CREATE INDEX IF NOT EXISTS idx_loop_directives_loop_id_status ON loop_directives(loop_id, status)"
        )
    )

    # grukk_messages table
    conn.execute(
        sa.text(
            """
            CREATE TABLE IF NOT EXISTS grukk_messages (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                user_id INTEGER REFERENCES users(id),
                role VARCHAR NOT NULL CHECK (role IN ('king', 'grukk')),
                content TEXT NOT NULL,
                metadata JSONB NOT NULL DEFAULT '{}',
                created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
            )
            """
        )
    )
    conn.execute(
        sa.text(
            "CREATE INDEX IF NOT EXISTS idx_grukk_messages_role_created ON grukk_messages(role, created_at)"
        )
    )


def downgrade() -> None:
    conn = op.get_bind()
    conn.execute(sa.text("DROP TABLE IF EXISTS grukk_messages"))
    conn.execute(sa.text("DROP TABLE IF EXISTS loop_directives"))
