"""Add api_tokens table and agent_type column

Revision ID: 002
Revises: 001
Create Date: 2026-03-31

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


revision: str = '002'
down_revision: Union[str, None] = '001'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Add agent_type column to documents table
    op.add_column('documents', sa.Column('agent_type', sa.String(20), sa.TEXT(), nullable=True))

    # Set default value for existing documents
    op.execute("UPDATE documents SET agent_type = 'general' WHERE agent_type IS NULL")

    # Make the column NOT NULL after setting defaults
    op.alter_column('documents', 'agent_type', nullable=False)

    # Add CHECK constraint for agent_type
    op.execute("""
        CREATE INDEX IF NOT EXISTS idx_documents_agent_type ON documents(agent_type)
    """)

    # Create api_tokens table
    op.execute("""
        CREATE TABLE IF NOT EXISTS api_tokens (
            id TEXT PRIMARY KEY,
            name TEXT NOT NULL,
            token_hash TEXT NOT NULL,
            role TEXT NOT NULL CHECK (role IN ('researcher', 'developer', 'viewer')),
            agent_id TEXT NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
            created_at TIMESTAMP NOT NULL DEFAULT (datetime('now')),
            last_used_at TIMESTAMP NULL
        )
    """)

    # Create index for api_tokens
    op.execute("CREATE INDEX IF NOT EXISTS idx_api_tokens_agent ON api_tokens(agent_id)")


def downgrade() -> None:
    op.drop_table('api_tokens')
    op.drop_column('documents', 'agent_type')