"""custom mcp providers

Revision ID: f4695a7728a8
Revises: d2eb2dcbb369
Create Date: 2025-12-03 17:18:10.625213

"""

import re
import uuid
from collections.abc import Sequence

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

# revision identifiers, used by Alembic.
revision: str = "f4695a7728a8"
down_revision: str | None = "d2eb2dcbb369"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

# Mapping of OAuth provider_id to MCP server URIs
MCP_PROVIDER_SERVER_URIS = {
    "sentry_mcp": "https://mcp.sentry.dev/mcp",
    "runreveal_mcp": "https://api.runreveal.com/mcp",
    "notion_mcp": "https://mcp.notion.com/mcp",
    "linear_mcp": "https://mcp.linear.app/mcp",
    "github_mcp": "https://api.githubcopilot.com/mcp",
    "secureannex_mcp": "https://mcp.secureannex.com/mcp",
}

# Mapping of OAuth provider_id to display names
MCP_PROVIDER_NAMES = {
    "sentry_mcp": "Sentry",
    "runreveal_mcp": "RunReveal",
    "notion_mcp": "Notion",
    "linear_mcp": "Linear",
    "github_mcp": "GitHub Copilot",
    "secureannex_mcp": "Secure Annex",
}


def slugify(text: str) -> str:
    """Convert text to a URL-friendly slug."""
    text = text.lower().strip()
    text = re.sub(r"[^\w\s-]", "", text)
    text = re.sub(r"[-\s]+", "-", text)
    return text


def generate_unique_slug(connection, owner_id: uuid.UUID, base_slug: str) -> str:
    """Generate a unique slug for an MCP integration."""
    candidate = base_slug
    suffix = 1

    while True:
        result = connection.execute(
            sa.text("""
                SELECT COUNT(*)
                FROM mcp_integration
                WHERE owner_id = :owner_id AND slug = :slug
            """),
            {"owner_id": owner_id, "slug": candidate},
        ).scalar()

        if result == 0:
            return candidate

        candidate = f"{base_slug}-{suffix}"
        suffix += 1


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    sa.Enum("OAUTH2", "CUSTOM", "NONE", name="mcpauthtype").create(op.get_bind())
    op.create_table(
        "mcp_integration",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("owner_id", sa.UUID(), nullable=False),
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("description", sa.String(), nullable=True),
        sa.Column("slug", sa.String(), nullable=False),
        sa.Column("server_uri", sa.String(), nullable=False),
        sa.Column(
            "auth_type",
            postgresql.ENUM(
                "OAUTH2", "CUSTOM", "NONE", name="mcpauthtype", create_type=False
            ),
            nullable=False,
        ),
        sa.Column("oauth_integration_id", sa.UUID(), nullable=True),
        sa.Column("encrypted_headers", sa.LargeBinary(), nullable=True),
        sa.Column(
            "created_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["oauth_integration_id"],
            ["oauth_integration.id"],
            name=op.f("fk_mcp_integration_oauth_integration_id_oauth_integration"),
            ondelete="SET NULL",
        ),
        sa.ForeignKeyConstraint(
            ["owner_id"],
            ["workspace.id"],
            name=op.f("fk_mcp_integration_owner_id_workspace"),
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_mcp_integration")),
        sa.UniqueConstraint("owner_id", "slug", name="uq_mcp_integration_owner_slug"),
    )
    op.create_index(
        op.f("ix_mcp_integration_id"), "mcp_integration", ["id"], unique=True
    )

    connection = op.get_bind()

    oauth_mcp_integrations = connection.execute(
        sa.text("""
            SELECT id, owner_id, provider_id
            FROM oauth_integration
            WHERE provider_id LIKE '%_mcp'
        """)
    ).fetchall()

    for id, owner_id, provider_id in oauth_mcp_integrations:
        # Skip if we don't have a mapping for this provider
        if provider_id not in MCP_PROVIDER_SERVER_URIS:
            continue

        server_uri = MCP_PROVIDER_SERVER_URIS[provider_id]
        name = MCP_PROVIDER_NAMES[provider_id]
        slug = generate_unique_slug(connection, owner_id, slugify(name))
        mcp_integration_id = uuid.uuid4()
        connection.execute(
            sa.text("""
                INSERT INTO mcp_integration (
                    id, owner_id, name, description, slug, server_uri,
                    auth_type, oauth_integration_id,
                    created_at, updated_at
                )
                VALUES (
                    :id, :owner_id, :name, :description, :slug, :server_uri,
                    'OAUTH2', :oauth_integration_id,
                    now(), now()
                )
            """),
            {
                "id": mcp_integration_id,
                "owner_id": owner_id,
                "name": name,
                "description": None,
                "slug": slug,
                "server_uri": server_uri,
                "oauth_integration_id": id,
            },
        )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f("ix_mcp_integration_id"), table_name="mcp_integration")
    op.drop_table("mcp_integration")
    sa.Enum("OAUTH2", "CUSTOM", "NONE", name="mcpauthtype").drop(op.get_bind())
    # ### end Alembic commands ###
