"""Backfill composio tools

Revision ID: f895232c144a
Revises: 25fc99e97839
Create Date: 2025-01-16 14:21:33.764332

"""

from typing import Sequence, Union

from alembic import op
from letta.schemas.enums import ToolType
from letta.settings import settings

# revision identifiers, used by Alembic.
revision: str = "f895232c144a"
down_revision: Union[str, None] = "416b9d2db10b"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Skip this migration for SQLite
    if not settings.letta_pg_uri_no_default:
        return

    # ### commands auto generated by Alembic - please adjust! ###
    # Define the value for EXTERNAL_COMPOSIO
    external_composio_value = ToolType.EXTERNAL_COMPOSIO.value

    # Update tool_type to EXTERNAL_COMPOSIO if the tags field includes "composio"
    # This is super brittle and awful but no other way to do this
    op.execute(
        f"""
        UPDATE tools
        SET tool_type = '{external_composio_value}'
        WHERE tags::jsonb @> '["composio"]';
        """
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # Skip this migration for SQLite
    if not settings.letta_pg_uri_no_default:
        return

    # ### commands auto generated by Alembic - please adjust! ###
    custom_value = ToolType.CUSTOM.value

    # Update tool_type to CUSTOM if the tags field includes "composio"
    # This is super brittle and awful but no other way to do this
    op.execute(
        f"""
        UPDATE tools
        SET tool_type = '{custom_value}'
        WHERE tags::jsonb @> '["composio"]';
        """
    )
    # ### end Alembic commands ###