# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import re
from unittest import mock

import pytest
import sqlalchemy as sa
from alembic.autogenerate import compare_metadata
from alembic.migration import MigrationContext
from sqlalchemy import MetaData

import airflow.providers.fab as provider_fab
from airflow.settings import engine
from airflow.utils.db import (
    compare_server_default,
    compare_type,
)

pytestmark = [pytest.mark.db_test]
try:
    from airflow.providers.fab.auth_manager.models.db import FABDBManager

    class TestFABDBManager:
        def setup_method(self):
            self.providers_dir: str = provider_fab.__path__[0]

        def test_version_table_name_set(self, session):
            assert FABDBManager(session=session).version_table_name == "alembic_version_fab"

        def test_migration_dir_set(self, session):
            assert FABDBManager(session=session).migration_dir == f"{self.providers_dir}/migrations"

        def test_alembic_file_set(self, session):
            assert FABDBManager(session=session).alembic_file == f"{self.providers_dir}/alembic.ini"

        def test_supports_table_dropping_set(self, session):
            assert FABDBManager(session=session).supports_table_dropping is True

        def test_database_schema_and_sqlalchemy_model_are_in_sync(self, session):
            def include_object(_, name, type_, *args):
                if type_ == "table" and name not in FABDBManager(session=session).metadata.tables:
                    return False
                return True

            all_meta_data = MetaData()
            for table_name, table in FABDBManager(session=session).metadata.tables.items():
                all_meta_data._add_table(table_name, table.schema, table)
            # create diff between database schema and SQLAlchemy model
            mctx = MigrationContext.configure(
                engine.connect(),
                opts={
                    "compare_type": compare_type,
                    "compare_server_default": compare_server_default,
                    "include_object": include_object,
                },
            )
            diff = compare_metadata(mctx, all_meta_data)
            ignores = [
                # These are conditionally added in ORM by the event listener
                lambda t: t[0] == "remove_index" and t[1].name == "idx_ab_register_user_username",
                lambda t: t[0] == "remove_index" and t[1].name == "idx_ab_user_username",
                lambda t: t[0] == "add_index" and t[1].name == "idx_ab_register_user_username",
                lambda t: t[0] == "add_index" and t[1].name == "idx_ab_user_username",
            ]
            for ignore in ignores:
                diff = [d for d in diff if not ignore(d)]

            assert not diff, "Database schema and SQLAlchemy model are not in sync: " + str(diff)

        @mock.patch("airflow.providers.fab.auth_manager.models.db._offline_migration")
        def test_downgrade_sql_no_from(self, mock_om, session, caplog):
            FABDBManager(session=session).downgrade(to_revision="abc", show_sql_only=True, from_revision=None)
            actual = mock_om.call_args.kwargs["revision"]
            assert re.match(r"[a-z0-9]+:abc", actual) is not None

        @mock.patch("airflow.providers.fab.auth_manager.models.db._offline_migration")
        def test_downgrade_sql_with_from(self, mock_om, session):
            FABDBManager(session=session).downgrade(
                to_revision="abc", show_sql_only=True, from_revision="123"
            )
            actual = mock_om.call_args.kwargs["revision"]
            assert actual == "123:abc"

        @mock.patch("alembic.command.downgrade")
        def test_downgrade_invalid_combo(self, mock_om, session):
            """can't combine `sql=False` and `from_revision`"""
            with pytest.raises(ValueError, match="can't be combined"):
                FABDBManager(session=session).downgrade(to_revision="abc", from_revision="123")

        @mock.patch("alembic.command.downgrade")
        def test_downgrade_with_from(self, mock_om, session):
            FABDBManager(session=session).downgrade(to_revision="abc")
            actual = mock_om.call_args.kwargs["revision"]
            assert actual == "abc"

        @mock.patch.object(FABDBManager, "get_current_revision")
        def test_sqlite_offline_upgrade_raises_with_revision(self, mock_gcr, session):
            with mock.patch(
                "airflow.providers.fab.auth_manager.models.db.settings.SQL_ALCHEMY_CONN",
                "sqlite:////tmp/fab-test.db",
            ):
                with pytest.raises(SystemExit, match="Offline migration not supported for SQLite"):
                    FABDBManager(session).upgradedb(from_revision=None, to_revision=None, show_sql_only=True)

        @mock.patch("alembic.command.upgrade")
        @mock.patch.object(FABDBManager, "create_db_from_orm")
        @mock.patch.object(FABDBManager, "_has_existing_manager_tables", return_value=False)
        @mock.patch.object(FABDBManager, "get_current_revision", return_value=None)
        def test_upgradedb_empty_db_without_migration_files_uses_create_db_from_orm(
            self,
            mock_get_current_revision,
            mock_has_existing_manager_tables,
            mock_create_db_from_orm,
            mock_upgrade,
            session,
        ):
            FABDBManager(session).upgradedb()

            mock_has_existing_manager_tables.assert_called_once()
            mock_create_db_from_orm.assert_called_once()
            mock_upgrade.assert_not_called()
            mock_get_current_revision.assert_called_once()

        @mock.patch("alembic.command.upgrade")
        @mock.patch("alembic.command.stamp")
        @mock.patch.object(FABDBManager, "get_script_object")
        @mock.patch.object(FABDBManager, "get_alembic_config", return_value=object())
        @mock.patch.object(FABDBManager, "create_db_from_orm")
        @mock.patch.object(FABDBManager, "_has_existing_manager_tables", return_value=True)
        @mock.patch.object(FABDBManager, "get_current_revision", return_value=None)
        def test_upgradedb_existing_tables_without_version_stamps_base_then_runs_migrations(
            self,
            mock_get_current_revision,
            mock_has_existing_manager_tables,
            mock_create_db_from_orm,
            mock_get_alembic_config,
            mock_get_script_object,
            mock_stamp,
            mock_upgrade,
            session,
        ):
            base_revision = mock.Mock(revision="base-revision", down_revision=None)
            mock_get_script_object.return_value.walk_revisions.return_value = [
                mock.Mock(revision="head-revision", down_revision="base-revision"),
                base_revision,
            ]

            FABDBManager(session).upgradedb()

            mock_has_existing_manager_tables.assert_called_once()
            mock_create_db_from_orm.assert_not_called()
            mock_stamp.assert_called_once_with(mock_get_alembic_config.return_value, "base-revision")
            mock_upgrade.assert_called_once_with(mock_get_alembic_config.return_value, revision="heads")
            mock_get_current_revision.assert_called_once()

        @mock.patch("alembic.command.upgrade")
        @mock.patch.object(FABDBManager, "create_db_from_orm")
        @mock.patch.object(FABDBManager, "get_current_revision", return_value=None)
        @mock.patch.object(FABDBManager, "get_alembic_config", return_value=object())
        def test_upgradedb_empty_db_with_use_migration_files_runs_alembic_upgrade(
            self,
            mock_get_alembic_config,
            mock_get_current_revision,
            mock_create_db_from_orm,
            mock_upgrade,
            session,
        ):
            FABDBManager(session).upgradedb(use_migration_files=True)

            mock_create_db_from_orm.assert_not_called()
            mock_get_current_revision.assert_called_once()
            mock_get_alembic_config.assert_called_once()
            mock_upgrade.assert_called_once_with(mock_get_alembic_config.return_value, revision="heads")

        @mock.patch("airflow.utils.db_manager.inspect")
        @mock.patch.object(FABDBManager, "metadata")
        @mock.patch("airflow.providers.fab.auth_manager.models.db._get_flask_db")
        def test_drop_tables(self, mock__get_flask_db, mock_metadata, mock_inspect, session):
            manager = FABDBManager(session)
            connection = mock.MagicMock()
            db = mock.MagicMock()
            mock__get_flask_db.return_value = db, mock.MagicMock()
            manager.drop_tables(connection)
            db.drop_all.assert_called_once_with()
            mock_metadata.drop_all.assert_called_once_with(connection)

        @pytest.mark.parametrize("skip_init", [True, False])
        @mock.patch.object(FABDBManager, "drop_tables")
        @mock.patch.object(FABDBManager, "initdb")
        @mock.patch("airflow.utils.db.create_global_lock", new=mock.MagicMock)
        def test_resetdb(self, mock_initdb, mock_drop_tables, session, skip_init):
            manager = FABDBManager(session)
            manager.resetdb(skip_init=skip_init)
            mock_drop_tables.assert_called_once()
            if skip_init:
                mock_initdb.assert_not_called()
            else:
                mock_initdb.assert_called_once()

        @pytest.mark.backend("mysql")
        def test_upgradedb_mysql_keeps_permission_view_role_indexes_after_fk_recreation(self, session):
            manager = FABDBManager(session=session)
            original_revision = manager.get_current_revision()

            try:
                manager.downgrade(to_revision="6709f7a774b9")
                manager.upgradedb(to_revision="02ca36b0235b")

                index_names = {
                    index["name"] for index in sa.inspect(engine).get_indexes("ab_permission_view_role")
                }

                assert "idx_permission_view_id" in index_names
                assert "idx_role_id" in index_names
            finally:
                current_revision = manager.get_current_revision()
                if original_revision and current_revision != original_revision:
                    manager.upgradedb(to_revision=original_revision)

        @pytest.mark.backend("mysql")
        def test_upgradedb_mysql_succeeds_with_implicit_fk_names(self, session):
            # Legacy MySQL DBs carry FKs auto-named like ``<table>_ibfk_N`` rather than the
            # naming-convention names ``batch_op.f()`` produces. The upgrade must resolve the
            # live FK names via introspection so the drops issue the correct identifiers.
            manager = FABDBManager(session=session)
            original_revision = manager.get_current_revision()

            try:
                manager.downgrade(to_revision="6709f7a774b9")

                replacements = {
                    "ab_permission_view_role": [
                        ("permission_view_id", "ab_permission_view", "id"),
                        ("role_id", "ab_role", "id"),
                    ],
                    "ab_user_role": [
                        ("user_id", "ab_user", "id"),
                        ("role_id", "ab_role", "id"),
                    ],
                }
                with engine.begin() as conn:
                    for table_name, fks in replacements.items():
                        for fk in sa.inspect(conn).get_foreign_keys(table_name):
                            if fk["name"]:
                                conn.execute(
                                    sa.text(f"ALTER TABLE `{table_name}` DROP FOREIGN KEY `{fk['name']}`")
                                )
                        for column, ref_table, ref_column in fks:
                            conn.execute(
                                sa.text(
                                    f"ALTER TABLE `{table_name}` "
                                    f"ADD FOREIGN KEY (`{column}`) "
                                    f"REFERENCES `{ref_table}` (`{ref_column}`) ON DELETE CASCADE"
                                )
                            )

                implicit_fk_names = {
                    fk["name"] for fk in sa.inspect(engine).get_foreign_keys("ab_permission_view_role")
                }
                assert any(name and "_ibfk_" in name for name in implicit_fk_names)

                manager.upgradedb(to_revision="02ca36b0235b")

                pvr_fk_names = {
                    fk["name"] for fk in sa.inspect(engine).get_foreign_keys("ab_permission_view_role")
                }
                user_role_fk_names = {
                    fk["name"] for fk in sa.inspect(engine).get_foreign_keys("ab_user_role")
                }
                assert "ab_permission_view_role_role_id_fkey" in pvr_fk_names
                assert "ab_permission_view_role_permission_view_id_fkey" in pvr_fk_names
                assert "ab_user_role_role_id_fkey" in user_role_fk_names
                assert "ab_user_role_user_id_fkey" in user_role_fk_names
            finally:
                current_revision = manager.get_current_revision()
                if original_revision and current_revision != original_revision:
                    manager.upgradedb(to_revision=original_revision)

except ModuleNotFoundError:
    pass
