# SPDX-License-Identifier: Apache-2.0
# Copyright 2022 The HuggingFace Authors.

from collections.abc import Iterator
from typing import Optional

import pytest

from mongodb_migration.database_migrations import (
    DatabaseMigration,
    _clean_maintenance_database,
)
from mongodb_migration.migration import IrreversibleMigrationError, Migration
from mongodb_migration.plan import Plan, SavedMigrationsError
from mongodb_migration.resources import MigrationsMongoResource


@pytest.fixture(autouse=True)
def migrations_mongo_resource(mongo_host: str) -> Iterator[MigrationsMongoResource]:
    database = "datasets_server_migrations_test"
    if "test" not in database:
        raise ValueError("Test must be launched on a test mongo database")
    with MigrationsMongoResource(database=database, host=mongo_host) as resource:
        yield resource
        _clean_maintenance_database()


class MigrationOK(Migration):
    def up(self) -> None:
        pass

    def down(self) -> None:
        pass

    def validate(self) -> None:
        pass


class MigrationErrorInUp(Migration):
    def up(self) -> None:
        raise RuntimeError("Error in up")

    def down(self) -> None:
        pass

    def validate(self) -> None:
        pass


class MigrationErrorInValidate(Migration):
    def up(self) -> None:
        pass

    def down(self) -> None:
        pass

    def validate(self) -> None:
        raise RuntimeError("Error in validation")


class MigrationErrorInUpAndDown(Migration):
    def up(self) -> None:
        raise RuntimeError("Error in up")

    def down(self) -> None:
        raise RuntimeError("Error in down")

    def validate(self) -> None:
        pass


class MigrationErrorIrreversible(Migration):
    def up(self) -> None:
        raise RuntimeError("Error in up")

    def down(self) -> None:
        raise IrreversibleMigrationError("Error in down")

    def validate(self) -> None:
        pass


def test_empty_plan() -> None:
    plan = Plan(collected_migrations=[])

    assert plan.collected_migrations == []
    plan.execute()
    assert plan.executed_migrations == []


migration_ok_a = MigrationOK(version="20221110230400", description="ok a")
migration_ok_b = MigrationOK(version="20221110230401", description="ok b")
migration_error_in_up = MigrationErrorInUp(version="20221110230402", description="error in up")
migration_error_in_validate = MigrationErrorInValidate(version="20221110230403", description="error in validate")
migration_error_in_up_and_down = MigrationErrorInUpAndDown(
    version="20221110230404", description="error in up and down"
)
migration_error_irreversible = MigrationErrorIrreversible(
    version="20221110230405", description="error because migration is irreversible"
)


@pytest.mark.parametrize(
    "collected_migrations",
    (
        [migration_ok_a, migration_ok_b],
        [migration_ok_b, migration_ok_a],
    ),
)
def test_collected_migrations_order_dont_matter(collected_migrations: list[Migration]) -> None:
    assert DatabaseMigration.objects.distinct("version") == []
    plan = Plan(collected_migrations=collected_migrations)
    assert plan.executed_migrations == []
    plan.execute()
    sorted_migrations = sorted(collected_migrations, key=lambda migration: migration.version)
    assert plan.executed_migrations == sorted_migrations
    assert DatabaseMigration.objects.distinct("version") == [migration.version for migration in sorted_migrations]


@pytest.mark.parametrize(
    "collected_migrations,executed_migrations,exception",
    [
        ([migration_error_in_up], [], RuntimeError),
        ([migration_error_in_validate], [], RuntimeError),
        ([migration_error_in_up_and_down], [migration_error_in_up_and_down], RuntimeError),
        ([migration_error_irreversible], [migration_error_irreversible], IrreversibleMigrationError),
        ([migration_ok_a, migration_error_in_up], [], RuntimeError),
        (
            [migration_ok_a, migration_error_in_up_and_down],
            [migration_ok_a, migration_error_in_up_and_down],
            RuntimeError,
        ),
    ],
)
def test_errors_in_migration_steps(
    collected_migrations: list[Migration], executed_migrations: list[Migration], exception: Optional[type[Exception]]
) -> None:
    assert DatabaseMigration.objects.distinct("version") == []
    plan = Plan(collected_migrations=collected_migrations)
    assert plan.executed_migrations == []
    if exception is None:
        # rollback worked
        plan.execute()
    else:
        # rollback failed
        with pytest.raises(exception):
            plan.execute()
    assert plan.executed_migrations == executed_migrations
    assert DatabaseMigration.objects.distinct("version") == [migration.version for migration in executed_migrations]


@pytest.mark.parametrize(
    "previous_migrations,collected_migrations,executed_migrations,exception",
    [
        ([], [], [], None),
        ([], [migration_ok_a], [migration_ok_a], None),
        ([migration_ok_a], [migration_ok_a, migration_ok_b], [migration_ok_b], None),
        # the previous migrations must be in the collected migrations
        ([migration_ok_a], [], [], SavedMigrationsError),
        ([migration_ok_a], [migration_ok_b], [], SavedMigrationsError),
        # error with the versions order
        ([migration_ok_b], [migration_ok_a, migration_ok_b], [], SavedMigrationsError),
    ],
)
def test_get_planned_migrations(
    previous_migrations: list[Migration],
    collected_migrations: list[Migration],
    executed_migrations: list[Migration],
    exception: Optional[type[Exception]],
) -> None:
    for migration in previous_migrations:
        DatabaseMigration(version=migration.version, description=migration.description).save()
    assert DatabaseMigration.objects.distinct("version") == [migration.version for migration in previous_migrations]
    plan = Plan(collected_migrations=collected_migrations)
    assert plan.executed_migrations == []
    if exception is None:
        # up worked
        plan.apply()
    else:
        # up failed
        with pytest.raises(exception):
            plan.apply()
    assert plan.executed_migrations == executed_migrations
    assert DatabaseMigration.objects.distinct("version") == [
        migration.version for migration in (previous_migrations + executed_migrations)
    ]


def test_internal_operations_are_idempotent() -> None:
    plan = Plan(collected_migrations=[migration_ok_a, migration_ok_b])
    plan.rollback()
    plan.rollback()
    plan.rollback()
    plan.apply()
    plan.apply()
    plan.apply()
    plan.apply()
    plan.rollback()
    plan.apply()
    plan.rollback()


def test_execute_is_idempotent() -> None:
    plan = Plan(collected_migrations=[migration_ok_a, migration_ok_b])
    plan.execute()
    plan.execute()
    Plan(collected_migrations=[migration_ok_a, migration_ok_b]).execute()
