from __future__ import annotations

import importlib
import unittest

from alembic.migration import MigrationContext
from alembic.operations import Operations
from sqlalchemy import create_engine, inspect


class PostboxMigrationTests(unittest.TestCase):
    def test_baseline_creates_and_drops_owned_tables(self) -> None:
        migration = importlib.import_module(
            "govoplan_postbox.backend.migrations.versions."
            "c7d2e5f8a1b4_v010_postbox_baseline"
        )
        route_migration = importlib.import_module(
            "govoplan_postbox.backend.migrations.versions."
            "e4b7c9d2a6f1_v011_hierarchy_routes"
        )
        occ_migration = importlib.import_module(
            "govoplan_postbox.backend.migrations.versions."
            "f5c8d0e3b7a2_v012_authoring_and_occ"
        )
        envelope_migration = importlib.import_module(
            "govoplan_postbox.backend.migrations.versions."
            "a6d9e1f4c8b3_v013_external_recipient_tokens"
        )
        protection_migration = importlib.import_module(
            "govoplan_postbox.backend.migrations.versions."
            "d8e3f6a9b2c5_postbox_content_protection"
        )
        engine = create_engine("sqlite:///:memory:")
        try:
            with engine.begin() as connection:
                operations = Operations(MigrationContext.configure(connection))
                original = migration.op
                route_original = route_migration.op
                occ_original = occ_migration.op
                envelope_original = envelope_migration.op
                protection_original = protection_migration.op
                migration.op = operations
                route_migration.op = operations
                occ_migration.op = operations
                envelope_migration.op = operations
                protection_migration.op = operations
                try:
                    migration.upgrade()
                    route_migration.upgrade()
                    occ_migration.upgrade()
                    envelope_migration.upgrade()
                    protection_migration.upgrade()
                    tables = set(inspect(connection).get_table_names())
                    self.assertIn("postboxes", tables)
                    self.assertIn("postbox_messages", tables)
                    self.assertIn("postbox_deliveries", tables)
                    self.assertIn("postbox_access_events", tables)
                    message_columns = {
                        column["name"]
                        for column in inspect(connection).get_columns(
                            "postbox_messages"
                        )
                    }
                    self.assertTrue(
                        {
                            "ciphertext_ref",
                            "signed_manifest_ref",
                            "wrapped_keys",
                            "external_recipient_tokens",
                            "key_epoch",
                            "expires_at",
                            "withdrawn_at",
                            "body_ciphertext",
                            "encryption_envelope_id",
                            "encryption_resource_id",
                        }.issubset(message_columns)
                    )
                    self.assertIn(
                        "encryption_vault_id",
                        {
                            column["name"]
                            for column in inspect(connection).get_columns(
                                "postbox_template_revisions"
                            )
                        },
                    )
                    self.assertIn("authoring_key", message_columns)
                    for table_name in (
                        "postbox_templates",
                        "postboxes",
                        "postbox_groupings",
                    ):
                        self.assertIn(
                            "resource_revision",
                            {
                                column["name"]
                                for column in inspect(connection).get_columns(
                                    table_name
                                )
                            },
                        )
                    route_columns = {
                        column["name"]
                        for column in inspect(connection).get_columns(
                            "postbox_routes"
                        )
                    }
                    self.assertTrue(
                        {"execute_after", "processed_at"}.issubset(
                            route_columns
                        )
                    )
                    protection_migration.downgrade()
                    envelope_migration.downgrade()
                    occ_migration.downgrade()
                    route_migration.downgrade()
                    migration.downgrade()
                    self.assertFalse(
                        {
                            table
                            for table in inspect(connection).get_table_names()
                            if table.startswith("postbox")
                        }
                    )
                finally:
                    migration.op = original
                    route_migration.op = route_original
                    occ_migration.op = occ_original
                    envelope_migration.op = envelope_original
                    protection_migration.op = protection_original
        finally:
            engine.dispose()


if __name__ == "__main__":
    unittest.main()