"""empty message

Revision ID: 812b93ec54d5
Revises: 
Create Date: 2026-03-07 20:53:22.618174

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from db.models.provider import LlmModelCapability
from db.models.tasks.shared import TaskUsage
from db.models.utils import DataClassJSON
from db.models.utils import PydanticJSON

# revision identifiers, used by Alembic.
revision: str = '812b93ec54d5'
down_revision: Union[str, Sequence[str], None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('providers',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(), nullable=False),
    sa.Column('type', sa.Enum('OPENAI', name='llmproviders', native_enum=False), nullable=False),
    sa.Column('base_url', sa.String(), nullable=False),
    sa.Column('api_key', sa.String(), nullable=False),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('toolsets',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(), nullable=False),
    sa.Column('internal_key', sa.String(), nullable=False),
    sa.Column('type', sa.Enum('BUILT_IN', 'MCP_LOCAL', 'MCP_REMOTE', name='toolsettype', native_enum=False), nullable=False),
    sa.Column('params', PydanticJSON(None), nullable=True),
    sa.Column('is_enabled', sa.Boolean(), nullable=False),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('internal_key')
    )
    op.create_table('workspaces',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(), nullable=False),
    sa.Column('directory', sa.String(), nullable=False),
    sa.Column('instruction', sa.String(), nullable=False),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('llm_models',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(), nullable=False),
    sa.Column('context_size', sa.Integer(), nullable=False),
    sa.Column('capability', DataClassJSON(LlmModelCapability), nullable=False),
    sa.Column('_provider_id', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['_provider_id'], ['providers.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('tools',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(), nullable=False),
    sa.Column('description', sa.String(), nullable=False),
    sa.Column('internal_key', sa.String(), nullable=False),
    sa.Column('is_enabled', sa.Boolean(), nullable=False),
    sa.Column('auto_approve', sa.Boolean(), nullable=False),
    sa.Column('_toolset_id', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['_toolset_id'], ['toolsets.id'], ),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('internal_key')
    )
    op.create_table('agents',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(), nullable=False),
    sa.Column('icon_name', sa.String(), nullable=False),
    sa.Column('instruction', sa.String(), nullable=False),
    sa.Column('model_id', sa.Integer(), nullable=True),
    sa.ForeignKeyConstraint(['model_id'], ['llm_models.id'], ondelete='SET NULL'),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('workspace_tool_association',
    sa.Column('workspace_id', sa.Integer(), nullable=False),
    sa.Column('tool_id', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['tool_id'], ['tools.id'], ondelete='CASCADE'),
    sa.ForeignKeyConstraint(['workspace_id'], ['workspaces.id'], ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('workspace_id', 'tool_id')
    )
    op.create_table('agent_tool_association',
    sa.Column('agent_id', sa.Integer(), nullable=False),
    sa.Column('tool_id', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['agent_id'], ['agents.id'], ondelete='CASCADE'),
    sa.ForeignKeyConstraint(['tool_id'], ['tools.id'], ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('agent_id', 'tool_id')
    )
    op.create_table('tasks',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('title', sa.String(), nullable=False),
    sa.Column('usage', DataClassJSON(TaskUsage), nullable=False),
    sa.Column('messages', PydanticJSON(None), nullable=False),
    sa.Column('last_run_at', sa.Integer(), nullable=False),
    sa.Column('agent_id', sa.Integer(), nullable=True),
    sa.Column('_workspace_id', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['_workspace_id'], ['workspaces.id'], ),
    sa.ForeignKeyConstraint(['agent_id'], ['agents.id'], ondelete='SET NULL'),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('workspace_agent_association',
    sa.Column('workspace_id', sa.Integer(), nullable=False),
    sa.Column('agent_id', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['agent_id'], ['agents.id'], ondelete='CASCADE'),
    sa.ForeignKeyConstraint(['workspace_id'], ['workspaces.id'], ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('workspace_id', 'agent_id')
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('workspace_agent_association')
    op.drop_table('tasks')
    op.drop_table('agent_tool_association')
    op.drop_table('workspace_tool_association')
    op.drop_table('agents')
    op.drop_table('tools')
    op.drop_table('llm_models')
    op.drop_table('workspaces')
    op.drop_table('toolsets')
    op.drop_table('providers')
    # ### end Alembic commands ###
