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

Alembic migration to fix constraints for resources and prompts to allow gateway namespacing,
using team/owner/gateway composite constraints, plus partial indexes for local uniqueness.

Revision ID: b1b2b3b4b5b6
Revises: 4e6273136e56
Create Date: 2026-01-26 13:01:00.000000
"""

# Standard
from typing import Sequence, Union

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

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


def constraint_exists(inspector: sa.Inspector, table_name: str, constraint_name: str) -> bool:
    """Check if a unique constraint exists on a table.

    Args:
        inspector: SQLAlchemy inspector instance.
        table_name: Name of the table to inspect.
        constraint_name: Name of the constraint to check.

    Returns:
        True if constraint exists, False if not or if check failed.
    """
    try:
        unique_constraints = inspector.get_unique_constraints(table_name)
        return any(uc["name"] == constraint_name for uc in unique_constraints)
    except Exception:
        # If introspection fails, return False so creates are attempted
        return False


def index_exists(inspector: sa.Inspector, table_name: str, index_name: str) -> bool:
    """Check if an index exists on a table.

    Args:
        inspector: SQLAlchemy inspector instance.
        table_name: Name of the table to inspect.
        index_name: Name of the index to check.

    Returns:
        True if index exists, False if not or if check failed.
    """
    try:
        indexes = inspector.get_indexes(table_name)
        return any(idx["name"] == index_name for idx in indexes)
    except Exception:
        # If introspection fails, return False so creates are attempted
        return False


def _upgrade_resources(inspector: sa.Inspector) -> None:
    """Update constraints for resources table.

    Args:
        inspector: SQLAlchemy inspector instance for database introspection.
    """
    if "resources" not in inspector.get_table_names():
        print("Resources table not found. Skipping resources migration.")
        return

    print("Processing resources table constraints...")
    with op.batch_alter_table("resources", schema=None) as batch_op:
        # Drop old constraint if exists
        if constraint_exists(inspector, "resources", "uq_team_owner_uri_resource"):
            try:
                batch_op.drop_constraint("uq_team_owner_uri_resource", type_="unique")
                print("Dropped constraint uq_team_owner_uri_resource.")
            except Exception as e:
                print(f"Could not drop constraint uq_team_owner_uri_resource: {e}")

        # Add new composite constraint if not exists
        if not constraint_exists(inspector, "resources", "uq_team_owner_gateway_uri_resource"):
            try:
                batch_op.create_unique_constraint("uq_team_owner_gateway_uri_resource", ["team_id", "owner_email", "gateway_id", "uri"])
                print("Created constraint uq_team_owner_gateway_uri_resource.")
            except Exception as e:
                print(f"Could not create constraint uq_team_owner_gateway_uri_resource (may already exist): {e}")
        else:
            print("Constraint uq_team_owner_gateway_uri_resource already exists, skipping create.")

        # Add partial index for local resources if not exists
        if not index_exists(inspector, "resources", "uq_team_owner_uri_resource_local"):
            try:
                batch_op.create_index(
                    "uq_team_owner_uri_resource_local", ["team_id", "owner_email", "uri"], unique=True, postgresql_where=text("gateway_id IS NULL"), sqlite_where=text("gateway_id IS NULL")
                )
                print("Created index uq_team_owner_uri_resource_local.")
            except Exception as e:
                print(f"Could not create index uq_team_owner_uri_resource_local (may already exist): {e}")
        else:
            print("Index uq_team_owner_uri_resource_local already exists, skipping create.")


def _upgrade_prompts(inspector: sa.Inspector) -> None:
    """Update constraints for prompts table.

    Args:
        inspector: SQLAlchemy inspector instance for database introspection.
    """
    if "prompts" not in inspector.get_table_names():
        print("Prompts table not found. Skipping prompts migration.")
        return

    print("Processing prompts table constraints...")
    with op.batch_alter_table("prompts", schema=None) as batch_op:
        # Drop old constraint if exists
        if constraint_exists(inspector, "prompts", "uq_team_owner_name_prompt"):
            try:
                batch_op.drop_constraint("uq_team_owner_name_prompt", type_="unique")
                print("Dropped constraint uq_team_owner_name_prompt.")
            except Exception as e:
                print(f"Could not drop constraint uq_team_owner_name_prompt: {e}")

        # Add new composite constraint if not exists
        if not constraint_exists(inspector, "prompts", "uq_team_owner_gateway_name_prompt"):
            try:
                batch_op.create_unique_constraint("uq_team_owner_gateway_name_prompt", ["team_id", "owner_email", "gateway_id", "name"])
                print("Created constraint uq_team_owner_gateway_name_prompt.")
            except Exception as e:
                print(f"Could not create constraint uq_team_owner_gateway_name_prompt (may already exist): {e}")
        else:
            print("Constraint uq_team_owner_gateway_name_prompt already exists, skipping create.")

        # Add partial index for local prompts if not exists
        if not index_exists(inspector, "prompts", "uq_team_owner_name_prompt_local"):
            try:
                batch_op.create_index(
                    "uq_team_owner_name_prompt_local", ["team_id", "owner_email", "name"], unique=True, postgresql_where=text("gateway_id IS NULL"), sqlite_where=text("gateway_id IS NULL")
                )
                print("Created index uq_team_owner_name_prompt_local.")
            except Exception as e:
                print(f"Could not create index uq_team_owner_name_prompt_local (may already exist): {e}")
        else:
            print("Index uq_team_owner_name_prompt_local already exists, skipping create.")


def upgrade() -> None:
    """Update unique constraints for Resources and Prompts.

    1. Drop restrictive old constraints (if they exist).
    2. Add new team-aware composite constraints (team, owner, gateway, id).
    3. Add partial unique indexes for local items (gateway_id IS NULL).

    Each table is processed independently - failure on one does not skip the other.
    """
    bind = op.get_bind()
    inspector = sa.inspect(bind)

    _upgrade_resources(inspector)
    _upgrade_prompts(inspector)


def _downgrade_resources(inspector: sa.Inspector) -> None:
    """Revert constraints for resources table.

    Args:
        inspector: SQLAlchemy inspector instance for database introspection.
    """
    if "resources" not in inspector.get_table_names():
        print("Resources table not found. Skipping resources downgrade.")
        return

    print("Reverting resources table constraints...")
    with op.batch_alter_table("resources", schema=None) as batch_op:
        # Drop new index if exists
        if index_exists(inspector, "resources", "uq_team_owner_uri_resource_local"):
            try:
                batch_op.drop_index("uq_team_owner_uri_resource_local")
                print("Dropped index uq_team_owner_uri_resource_local.")
            except Exception as e:
                print(f"Could not drop index uq_team_owner_uri_resource_local: {e}")

        # Drop new constraint if exists
        if constraint_exists(inspector, "resources", "uq_team_owner_gateway_uri_resource"):
            try:
                batch_op.drop_constraint("uq_team_owner_gateway_uri_resource", type_="unique")
                print("Dropped constraint uq_team_owner_gateway_uri_resource.")
            except Exception as e:
                print(f"Could not drop constraint uq_team_owner_gateway_uri_resource: {e}")

        # Recreate old constraint if not exists
        if not constraint_exists(inspector, "resources", "uq_team_owner_uri_resource"):
            try:
                batch_op.create_unique_constraint("uq_team_owner_uri_resource", ["team_id", "owner_email", "uri"])
                print("Created constraint uq_team_owner_uri_resource.")
            except Exception as e:
                print(f"Could not create constraint uq_team_owner_uri_resource (may already exist): {e}")


def _downgrade_prompts(inspector: sa.Inspector) -> None:
    """Revert constraints for prompts table.

    Args:
        inspector: SQLAlchemy inspector instance for database introspection.
    """
    if "prompts" not in inspector.get_table_names():
        print("Prompts table not found. Skipping prompts downgrade.")
        return

    print("Reverting prompts table constraints...")
    with op.batch_alter_table("prompts", schema=None) as batch_op:
        # Drop new index if exists
        if index_exists(inspector, "prompts", "uq_team_owner_name_prompt_local"):
            try:
                batch_op.drop_index("uq_team_owner_name_prompt_local")
                print("Dropped index uq_team_owner_name_prompt_local.")
            except Exception as e:
                print(f"Could not drop index uq_team_owner_name_prompt_local: {e}")

        # Drop new constraint if exists
        if constraint_exists(inspector, "prompts", "uq_team_owner_gateway_name_prompt"):
            try:
                batch_op.drop_constraint("uq_team_owner_gateway_name_prompt", type_="unique")
                print("Dropped constraint uq_team_owner_gateway_name_prompt.")
            except Exception as e:
                print(f"Could not drop constraint uq_team_owner_gateway_name_prompt: {e}")

        # Recreate old constraint if not exists
        if not constraint_exists(inspector, "prompts", "uq_team_owner_name_prompt"):
            try:
                batch_op.create_unique_constraint("uq_team_owner_name_prompt", ["team_id", "owner_email", "name"])
                print("Created constraint uq_team_owner_name_prompt.")
            except Exception as e:
                print(f"Could not create constraint uq_team_owner_name_prompt (may already exist): {e}")


def downgrade() -> None:
    """Revert constraints to original state.

    Each table is processed independently - failure on one does not skip the other.
    """
    bind = op.get_bind()
    inspector = sa.inspect(bind)

    _downgrade_resources(inspector)
    _downgrade_prompts(inspector)
