"""Rename tables and foreign keys [5330ba58bf20].

Revision ID: 5330ba58bf20
Revises: 7280c14811d6
Create Date: 2022-11-03 16:33:15.220179

"""

from collections import defaultdict
from typing import Dict, List, Tuple

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "5330ba58bf20"
down_revision = "7280c14811d6"
branch_labels = None
depends_on = None


def _fk_constraint_name(source: str, target: str, source_column: str) -> str:
    """Defines the name of a foreign key constraint.

    See https://alembic.sqlalchemy.org/en/latest/batch.html#dropping-unnamed-or-named-foreign-key-constraints.

    Args:
        source: Source table name.
        target: Target table name.
        source_column: Source column name.

    Returns:
        Name of the foreign key constraint.
    """
    return f"fk_{source}_{source_column}_{target}"


def _drop_fk_constraint(source: str, constraint_name: str) -> None:
    """Drops a (potentially unnamed) foreign key constraint.

    This is somewhat tricky since we have not explicitly named constraints
    before and each backend names the constraints differently. SQLite even
    doesn't name them at all.

    See https://alembic.sqlalchemy.org/en/latest/batch.html#dropping-unnamed-or-named-foreign-key-constraints.

    Args:
        source: Source table name.
        constraint_name: Name of the foreign key constraint.
    """
    naming_convention = {
        "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
    }
    with op.batch_alter_table(
        source, naming_convention=naming_convention
    ) as batch_op:
        batch_op.drop_constraint(
            constraint_name=constraint_name,
            type_="foreignkey",
        )


def _create_fk_constraint(
    source: str,
    target: str,
    source_column: str,
    target_column: str,
    ondelete: str,
) -> None:
    """Create a foreign key constraint.

    Args:
        source: Source table name.
        target: Target table name.
        source_column: Source column name.
        target_column: Target column name.
        ondelete: On delete behavior.
    """
    constraint_name = _fk_constraint_name(source, target, source_column)
    with op.batch_alter_table(source) as batch_op:
        batch_op.create_foreign_key(
            constraint_name=constraint_name,
            referent_table=target,
            local_cols=[source_column],
            remote_cols=[target_column],
            ondelete=ondelete,
        )


def _get_changes() -> (
    Tuple[
        List[str],
        List[str],
        List[str],
        List[Tuple[str, str, str, str, str]],
        List[Tuple[str, str, str, str, str]],
    ]
):
    """Define the data that should be changed in the schema.

    Returns:
        A tuple of four lists:
        - old table names
        - new table names
        - new table names of tables that should have a `NOT NULL` constraint
            on the `project_id` column
        - old foreign key constraints:
            (source, target, source_column, target_column, ondelete)
        - new foreign key constraints
            (source, target, source_column, target_column, ondelete)
    """
    # Define all the tables that should be renamed
    table_name_mapping: Dict[str, str] = {
        "roleschema": "role",
        "stepinputartifactschema": "step_run_input_artifact",
        "userroleassignmentschema": "user_role_assignment",
        "steprunorderschema": "step_run_parents",
        "teamschema": "team",
        "artifactschema": "artifacts",
        "pipelinerunschema": "pipeline_run",
        "steprunschema": "step_run",
        "teamassignmentschema": "team_assignment",
        "projectschema": "workspace",
        "flavorschema": "flavor",
        "userschema": "user",
        "stackcomponentschema": "stack_component",
        "pipelineschema": "pipeline",
        "stackcompositionschema": "stack_composition",
        "teamroleassignmentschema": "team_role_assignment",
        "stackschema": "stack",
        "rolepermissionschema": "role_permission",
    }
    reversed_table_name_mapping = {v: k for k, v in table_name_mapping.items()}
    old_table_names = list(table_name_mapping.keys())
    new_table_names = list(table_name_mapping.values())

    # Define all foreign key constraints that need to be adjusted
    project_user_fk_tables = [
        "stack_component",
        "flavor",
        "pipeline",
        "pipeline_run",
        "stack",
    ]
    new_fk_constraints: List[Tuple[str, str, str, str, str]] = [
        *[
            (source, "workspace", "project_id", "id", "CASCADE")
            for source in project_user_fk_tables
        ],  # 5
        *[
            (source, "user", "user_id", "id", "SET NULL")
            for source in project_user_fk_tables
        ],  # 10
        ("stack_composition", "stack", "stack_id", "id", "CASCADE"),  # 11
        (
            "stack_composition",
            "stack_component",
            "component_id",
            "id",
            "CASCADE",
        ),  # 12
        ("pipeline_run", "pipeline", "pipeline_id", "id", "SET NULL"),  # 13
        ("pipeline_run", "stack", "stack_id", "id", "SET NULL"),  # 14
        ("step_run", "pipeline_run", "pipeline_run_id", "id", "CASCADE"),  # 15
        (
            "step_run_input_artifact",
            "step_run",
            "step_id",
            "id",
            "CASCADE",
        ),  # 16
        (
            "step_run_input_artifact",
            "artifacts",
            "artifact_id",
            "id",
            "CASCADE",
        ),  # 17
        ("step_run_parents", "step_run", "parent_id", "id", "CASCADE"),  # 18
        ("step_run_parents", "step_run", "child_id", "id", "CASCADE"),  # 19
        ("artifacts", "step_run", "producer_step_id", "id", "CASCADE"),  # 20
        ("artifacts", "step_run", "parent_step_id", "id", "CASCADE"),  # 21
        ("team_assignment", "user", "user_id", "id", "CASCADE"),  # 22
        ("team_assignment", "team", "team_id", "id", "CASCADE"),  # 23
        ("team_role_assignment", "team", "team_id", "id", "CASCADE"),  # 24
        ("team_role_assignment", "role", "role_id", "id", "CASCADE"),  # 25
        (
            "team_role_assignment",
            "workspace",
            "project_id",
            "id",
            "CASCADE",
        ),  # 26
        ("user_role_assignment", "user", "user_id", "id", "CASCADE"),  # 27
        ("user_role_assignment", "role", "role_id", "id", "CASCADE"),  # 28
        (
            "user_role_assignment",
            "workspace",
            "project_id",
            "id",
            "CASCADE",
        ),  # 29
        ("role_permission", "role", "role_id", "id", "CASCADE"),  # 30
    ]
    old_fk_constraints = [
        (
            reversed_table_name_mapping[source],
            reversed_table_name_mapping[target],
            source_col,
            target_col,
            ondelete,
        )
        for source, target, source_col, target_col, ondelete in new_fk_constraints
    ]
    return (
        old_table_names,
        new_table_names,
        project_user_fk_tables,
        old_fk_constraints,
        new_fk_constraints,
    )


def upgrade() -> None:
    """Upgrade database schema and/or data, creating a new revision.

    Raises:
        NotImplementedError: If the database engine is not SQLite or MySQL.
    """
    (
        old_table_names,
        new_table_names,
        project_not_null_tables,
        old_fk_constraints,
        new_fk_constraints,
    ) = _get_changes()

    engine_name = op.get_bind().engine.name

    # Under MySQL, we need to sort the old foreign keys by source table and
    # source column first since the default foreign key names contain the
    # foreign key number.
    if engine_name == "mysql":
        old_fk_constraints.sort(key=lambda x: (x[0], x[2]))
    source_table_fk_constraint_counts: Dict[str, int] = defaultdict(int)

    # Drop old foreign key constraints.
    for source, target, source_column, _, _ in old_fk_constraints:
        if engine_name == "sqlite":
            constraint_name = _fk_constraint_name(
                source, target, source_column
            )
        elif engine_name == "mysql":
            source_table_fk_constraint_counts[source] += 1
            fk_num = source_table_fk_constraint_counts[source]
            constraint_name = f"{source}_ibfk_{fk_num}"
        else:
            raise NotImplementedError(f"Unsupported engine: {engine_name}")
        _drop_fk_constraint(source, constraint_name)

    # Rename tables
    for old_table_name, new_table_name in zip(
        old_table_names, new_table_names
    ):
        op.rename_table(old_table_name, new_table_name)

    # Set `project_id` to `NOT NULL` where appropriate.
    for table_name in project_not_null_tables:
        with op.batch_alter_table(table_name, schema=None) as batch_op:
            batch_op.alter_column(
                "project_id", existing_type=sa.CHAR(length=32), nullable=False
            )

    # Create new foreign key constraints
    for (
        source,
        target,
        source_column,
        target_column,
        ondelete,
    ) in new_fk_constraints:
        _create_fk_constraint(
            source, target, source_column, target_column, ondelete
        )


def downgrade() -> None:
    """Downgrade database schema and/or data back to the previous revision."""
    (
        old_table_names,
        new_table_names,
        project_not_null_tables,
        old_fk_constraints,
        new_fk_constraints,
    ) = _get_changes()

    # Drop new foreign key constraints
    for source, target, source_column, _, _ in new_fk_constraints:
        constraint_name = _fk_constraint_name(source, target, source_column)
        _drop_fk_constraint(source, constraint_name)

    # Remove `project_id NOT NULL` where appropriate.
    for table_name in project_not_null_tables:
        with op.batch_alter_table(table_name, schema=None) as batch_op:
            batch_op.alter_column(
                "project_id", existing_type=sa.CHAR(length=32), nullable=True
            )

    # Rename tables
    for old_table_name, new_table_name in zip(
        old_table_names, new_table_names
    ):
        op.rename_table(new_table_name, old_table_name)

    # Create old foreign key constraints
    for (
        source,
        target,
        source_column,
        target_column,
        ondelete,
    ) in old_fk_constraints:
        _create_fk_constraint(
            source, target, source_column, target_column, ondelete
        )
