"""Add approvals table and waiting_for_input task status.

Revision ID: 0015
Revises: 0014
"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

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

_NEW_STATUS_CHECK = (
    "status IN ('grooming', 'todo', 'in_progress', 'done', 'cancelled',"
    " 'timeout', 'ready-to-deploy', 'merged', 'closed', 'waiting_for_input')"
)

_OLD_STATUS_CHECK = (
    "status IN ('grooming', 'todo', 'in_progress', 'done', 'cancelled',"
    " 'timeout', 'ready-to-deploy', 'merged', 'closed')"
)


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

    # 1. Add waiting_for_input to task status check constraint
    rows = conn.execute(
        sa.text(
            "SELECT conname FROM pg_constraint"
            " WHERE conrelid = 'tasks'::regclass AND contype = 'c'"
            " AND pg_get_constraintdef(oid) ILIKE '%status%'"
        )
    ).fetchall()

    for (conname,) in rows:
        conn.execute(sa.text(f'ALTER TABLE tasks DROP CONSTRAINT IF EXISTS "{conname}"'))

    conn.execute(
        sa.text(
            f"ALTER TABLE tasks ADD CONSTRAINT tasks_status_check CHECK ({_NEW_STATUS_CHECK})"
        )
    )

    # 2. Create approvals table
    conn.execute(
        sa.text(
            """
            CREATE TABLE IF NOT EXISTS approvals (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                agent_session_id INTEGER REFERENCES agent_sessions(id) ON DELETE SET NULL,
                task_id INTEGER REFERENCES tasks(id) ON DELETE SET NULL,
                title TEXT NOT NULL,
                context TEXT NOT NULL DEFAULT '',
                options JSONB NOT NULL DEFAULT '[]',
                priority TEXT NOT NULL DEFAULT 'medium'
                    CHECK (priority IN ('critical', 'medium', 'low')),
                status TEXT NOT NULL DEFAULT 'pending'
                    CHECK (status IN ('pending', 'approved', 'rejected', 'deferred', 'expired')),
                response TEXT,
                expires_at TIMESTAMPTZ,
                created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
            )
            """
        )
    )

    conn.execute(
        sa.text(
            "CREATE INDEX IF NOT EXISTS idx_approvals_status ON approvals(status)"
        )
    )
    conn.execute(
        sa.text(
            "CREATE INDEX IF NOT EXISTS idx_approvals_agent_session_id ON approvals(agent_session_id)"
        )
    )
    conn.execute(
        sa.text(
            "CREATE INDEX IF NOT EXISTS idx_approvals_task_id ON approvals(task_id)"
        )
    )


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

    conn.execute(sa.text("DROP TABLE IF EXISTS approvals"))

    # Restore old task status check constraint
    rows = conn.execute(
        sa.text(
            "SELECT conname FROM pg_constraint"
            " WHERE conrelid = 'tasks'::regclass AND contype = 'c'"
            " AND pg_get_constraintdef(oid) ILIKE '%status%'"
        )
    ).fetchall()

    for (conname,) in rows:
        conn.execute(sa.text(f'ALTER TABLE tasks DROP CONSTRAINT IF EXISTS "{conname}"'))

    conn.execute(
        sa.text(
            f"ALTER TABLE tasks ADD CONSTRAINT tasks_status_check CHECK ({_OLD_STATUS_CHECK})"
        )
    )
