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

unique const changes for prompt and resource

Revision ID: e5a59c16e041
Revises: 8a2934be50c0
Create Date: 2025-10-15 11:20:53.888488
"""

# Standard
from typing import Sequence, Union

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

# revision identifiers, used by Alembic.
revision: str = "e5a59c16e041"
down_revision: Union[str, Sequence[str], None] = "8a2934be50c0"
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 prompts, resources and a2a agents.
    This migration recreates tables with updated unique constraints and preserves data.
    Compatible with SQLite and PostgreSQL.
    """
    bind = op.get_bind()
    inspector = sa.inspect(bind)

    # ### commands auto generated by Alembic - please adjust! ###
    for tbl, constraints in {
        "prompts": [("name", "uq_team_owner_name_prompts")],
        "resources": [("uri", "uq_team_owner_uri_resources")],
        "a2a_agents": [("slug", "uq_team_owner_slug_a2a_agents")],
    }.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)

            # For PostgreSQL, find and drop incoming foreign keys from other tables
            incoming_fks = []
            if bind.dialect.name == "postgresql":
                # Query PostgreSQL system catalogs to find foreign keys pointing to this table
                fk_query = sa.text("""
                    SELECT
                        tc.table_name,
                        tc.constraint_name,
                        kcu.column_name,
                        ccu.column_name AS foreign_column_name
                    FROM information_schema.table_constraints AS tc
                    JOIN information_schema.key_column_usage AS kcu
                        ON tc.constraint_name = kcu.constraint_name
                        AND tc.table_schema = kcu.table_schema
                    JOIN information_schema.constraint_column_usage AS ccu
                        ON ccu.constraint_name = tc.constraint_name
                        AND ccu.table_schema = tc.table_schema
                    WHERE tc.constraint_type = 'FOREIGN KEY'
                        AND ccu.table_name = :table_name
                """)
                result = bind.execute(fk_query, {"table_name": tbl})
                for row in result:
                    incoming_fks.append(
                        {
                            "table": row[0],
                            "constraint": row[1],
                            "column": row[2],
                            "foreign_column": row[3],
                        }
                    )
                    print(f"  Found incoming FK: {row[0]}.{row[1]} -> {tbl}.{row[3]}")

                # Drop incoming foreign keys
                for fk in incoming_fks:
                    print(f"  Dropping FK: {fk['table']}.{fk['constraint']}")
                    op.drop_constraint(fk["constraint"], fk["table"], type_="foreignkey")

            # 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())
            uqs_to_copy = []
            # # # Copy unique constraints that we're not replacing, and skip any unique constraint only on 'name'
            if tbl == "prompts":
                uqs_to_copy = []
                for uq in table.constraints:
                    if isinstance(uq, sa.UniqueConstraint) and set([col.name for col in uq.columns]) != {"name"} and not any(uq.name == c[1] if uq.name else False for c in constraints):
                        uqs_to_copy.append(uq)
            # Copy unique constraints that we're not replacing, and skip any unique constraint only on 'name'
            if tbl == "resources":
                uqs_to_copy = [
                    uq
                    for uq in table.constraints
                    if isinstance(uq, sa.UniqueConstraint) and set([col.name for col in uq.columns]) != {"uri"} and not any(uq.name == c[1] if uq.name else False for c in constraints)
                ]

            # For a2a_agents, also drop any unique constraint on just 'name'
            if tbl == "a2a_agents":
                uqs_to_copy = [
                    uq
                    for uq in table.constraints
                    if isinstance(uq, sa.UniqueConstraint)
                    and set([col.name for col in uq.columns]) != {"name"}
                    and set([col.name for col in uq.columns]) != {"slug"}
                    and not any(uq.name == c[1] if uq.name else False for c in constraints)
                ]
            for uq in uqs_to_copy:
                if uq is not None:
                    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)

            # For PostgreSQL, recreate the incoming foreign keys
            if bind.dialect.name == "postgresql":
                for fk in incoming_fks:
                    print(f"  Recreating FK: {fk['table']}.{fk['constraint']} -> {tbl}.{fk['foreign_column']}")
                    op.create_foreign_key(fk["constraint"], fk["table"], tbl, [fk["column"]], [fk["foreign_column"]])

        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 prompts, resources and a2a_agents.
    This migration recreates tables with the original unique constraints and preserves data.
    Compatible with SQLite and PostgreSQL.
    """
    bind = op.get_bind()
    inspector = sa.inspect(bind)

    for tbl, constraints in {
        "prompts": [("name", "uq_team_owner_name_prompts")],
        "resources": [("uri", "uq_team_owner_uri_resources")],
        "a2a_agents": [("slug", "uq_team_owner_slug_a2a_agents")],
    }.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)

            # For PostgreSQL, find and drop incoming foreign keys from other tables
            incoming_fks = []
            if bind.dialect.name == "postgresql":
                # Query PostgreSQL system catalogs to find foreign keys pointing to this table
                fk_query = sa.text("""
                    SELECT
                        tc.table_name,
                        tc.constraint_name,
                        kcu.column_name,
                        ccu.column_name AS foreign_column_name
                    FROM information_schema.table_constraints AS tc
                    JOIN information_schema.key_column_usage AS kcu
                        ON tc.constraint_name = kcu.constraint_name
                        AND tc.table_schema = kcu.table_schema
                    JOIN information_schema.constraint_column_usage AS ccu
                        ON ccu.constraint_name = tc.constraint_name
                        AND ccu.table_schema = tc.table_schema
                    WHERE tc.constraint_type = 'FOREIGN KEY'
                        AND ccu.table_name = :table_name
                """)
                result = bind.execute(fk_query, {"table_name": tbl})
                for row in result:
                    incoming_fks.append(
                        {
                            "table": row[0],
                            "constraint": row[1],
                            "column": row[2],
                            "foreign_column": row[3],
                        }
                    )
                    print(f"  Found incoming FK: {row[0]}.{row[1]} -> {tbl}.{row[3]}")

                # Drop incoming foreign keys
                for fk in incoming_fks:
                    print(f"  Dropping FK: {fk['table']}.{fk['constraint']}")
                    op.drop_constraint(fk["constraint"], fk["table"], type_="foreignkey")

            # 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
            uqs_to_copy = [uq for uq in table.constraints if isinstance(uq, sa.UniqueConstraint) and not any(uq.name == c[1] if uq.name else False for c in constraints)]
            for uq in uqs_to_copy:
                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))
            if tbl == "a2a_agents":
                # Also re-add unique constraint on 'name' for a2a_agents
                new_table.append_constraint(sa.UniqueConstraint("name"))
            # 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)

            # For PostgreSQL, recreate the incoming foreign keys
            if bind.dialect.name == "postgresql":
                for fk in incoming_fks:
                    print(f"  Recreating FK: {fk['table']}.{fk['constraint']} -> {tbl}.{fk['foreign_column']}")
                    op.create_foreign_key(fk["constraint"], fk["table"], tbl, [fk["column"]], [fk["foreign_column"]])

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