# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/alembic/versions/e182847d89e6_unique_constraints_changes_for_gateways_.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Alembic migration for unique constraints on gateways, tools, and servers.

Revision ID: e182847d89e6
Revises: cfc3d6aa0fb2
Create Date: 2025-09-12 12:55:17.537176
"""

# Standard
from typing import Sequence, Union

# Third-Party
from alembic import op
import sqlalchemy as sa

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


def upgrade() -> None:
    """
    Apply schema changes to add or update unique constraints for gateways, tools, and servers.
    This migration recreates tables with updated unique constraints and preserves data.
    Compatible with SQLite, MySQL, and PostgreSQL.
    """
    bind = op.get_bind()
    inspector = sa.inspect(bind)

    # ### commands auto generated by Alembic - please adjust! ###
    for tbl, constraints in {
        "servers": [("name", "uq_team_owner_name_servers")],
        "tools": [("name", "uq_team_owner_name_tools")],
        "gateways": [("slug", "uq_team_owner_slug_gateway"), ("url", "uq_team_owner_url_gateway")],
    }.items():
        try:
            print(f"Processing {tbl} for unique constraint update...")

            # Get table metadata using SQLAlchemy
            metadata = sa.MetaData()
            table = sa.Table(tbl, metadata, autoload_with=bind)

            # Create temporary table name
            tmp_table = f"{tbl}_tmp_nounique"

            # Drop temp table if it exists
            if inspector.has_table(tmp_table):
                op.drop_table(tmp_table)

            # Create new table structure with same columns but no old unique constraints
            new_table = sa.Table(tmp_table, metadata)

            for column in table.columns:
                # Copy column with same properties
                new_column = column.copy()
                new_table.append_column(new_column)

            # Copy foreign key constraints
            for fk in table.foreign_keys:
                new_table.append_constraint(fk.constraint.copy())

            # Copy unique constraints that we're not replacing
            for uq in table.constraints:
                if isinstance(uq, sa.UniqueConstraint):
                    # Only keep constraints that we're not replacing
                    if not any(uq.name == c[1] if uq.name else False for c in constraints):
                        new_table.append_constraint(uq.copy())

            # Create the temporary table
            new_table.create(bind)

            # Copy data
            column_names = [c.name for c in table.columns]
            insert_stmt = new_table.insert().from_select(column_names, sa.select(*[table.c[name] for name in column_names]))
            bind.execute(insert_stmt)

            # Add new unique constraints using batch operations for SQLite compatibility
            with op.batch_alter_table(tmp_table, schema=None) as batch_op:
                for col, constraint_name in constraints:
                    cols = ["team_id", "owner_email", col]
                    batch_op.create_unique_constraint(constraint_name, cols)

            # Drop original table and rename temp table
            op.drop_table(tbl)
            op.rename_table(tmp_table, tbl)

        except Exception as e:
            print(f"Warning: Could not update unique constraint on {tbl} table: {e}")
        # ### end Alembic commands ###


def downgrade() -> None:
    """
    Revert schema changes, restoring previous unique constraints for gateways, tools, and servers.
    This migration recreates tables with the original unique constraints and preserves data.
    Compatible with SQLite, MySQL, and PostgreSQL.
    """
    bind = op.get_bind()
    inspector = sa.inspect(bind)

    for tbl, constraints in {
        "servers": [("name", "uq_team_owner_name_servers")],
        "tools": [("name", "uq_team_owner_name_tools")],
        "gateways": [("slug", "uq_team_owner_slug_gateway"), ("url", "uq_team_owner_url_gateway")],
    }.items():
        try:
            print(f"Processing {tbl} for unique constraint revert...")

            # Get table metadata using SQLAlchemy
            metadata = sa.MetaData()
            table = sa.Table(tbl, metadata, autoload_with=bind)

            # Create temporary table name
            tmp_table = f"{tbl}_tmp_revert"

            # Drop temp table if it exists
            if inspector.has_table(tmp_table):
                op.drop_table(tmp_table)

            # Create new table structure with same columns but original unique constraints
            new_table = sa.Table(tmp_table, metadata)

            for column in table.columns:
                # Copy column with same properties
                new_column = column.copy()
                new_table.append_column(new_column)

            # Copy foreign key constraints
            for fk in table.foreign_keys:
                new_table.append_constraint(fk.constraint.copy())

            # Copy unique constraints that we're not reverting
            for uq in table.constraints:
                if isinstance(uq, sa.UniqueConstraint):
                    # Only keep constraints that we're not reverting
                    if not any(uq.name == c[1] if uq.name else False for c in constraints):
                        new_table.append_constraint(uq.copy())

            # Add back the original single-column unique constraints
            for col, _ in constraints:
                if col in [c.name for c in table.columns]:
                    new_table.append_constraint(sa.UniqueConstraint(col))

            # Create the temporary table
            new_table.create(bind)

            # Copy data
            column_names = [c.name for c in table.columns]
            insert_stmt = new_table.insert().from_select(column_names, sa.select(*[table.c[name] for name in column_names]))
            bind.execute(insert_stmt)

            # Drop original table and rename temp table
            op.drop_table(tbl)
            op.rename_table(tmp_table, tbl)

        except Exception as e:
            print(f"Warning: Could not revert unique constraint on {tbl} table: {e}")
    # ### end Alembic commands ###
