from __future__ import annotations

import tempfile
import unittest
from datetime import UTC, datetime
from pathlib import Path

from alembic import command
from alembic.runtime.migration import MigrationContext
from sqlalchemy import MetaData, Table, create_engine, inspect, select

from govoplan_core.db.migrations import alembic_config, migrate_database
from govoplan_datasources.backend.manifest import get_manifest


class DatasourceMigrationTests(unittest.TestCase):
    @staticmethod
    def _config(url: str):
        return alembic_config(
            database_url=url,
            enabled_modules=("datasources",),
            manifest_factories=(get_manifest,),
        )

    def test_baseline_creates_datasource_tables_and_head(self) -> None:
        with tempfile.TemporaryDirectory(prefix="govoplan-datasources-migration-") as directory:
            url = f"sqlite:///{Path(directory) / 'datasources.db'}"
            migrate_database(
                database_url=url,
                enabled_modules=("datasources",),
                manifest_factories=(get_manifest,),
            )
            engine = create_engine(url)
            try:
                with engine.connect() as connection:
                    self.assertIn(
                        "c9e3a6f1d4b8",
                        set(MigrationContext.configure(connection).get_current_heads()),
                    )
                    catalogue_columns = {
                        item["name"]
                        for item in inspect(connection).get_columns(
                            "datasource_catalogue"
                        )
                    }
                    self.assertTrue(
                        {
                            "authority_mode",
                            "classification",
                            "publication_state",
                            "owner_ref",
                            "quality_policy",
                            "access_policy_ref",
                            "visibility_policy",
                            "dependency_refs",
                        }.issubset(catalogue_columns)
                    )
                    self.assertEqual(
                        {
                            "datasource_catalogue",
                            "datasource_governance_references",
                            "datasource_materializations",
                            "datasource_payload_rows",
                            "datasource_payloads",
                            "datasource_publications",
                            "datasource_stages",
                        },
                        {
                            name
                            for name in inspect(connection).get_table_names()
                            if name.startswith("datasource_")
                        },
                    )
            finally:
                engine.dispose()

    def test_governance_reference_index_backfills_existing_catalogue_rows(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory(
            prefix="govoplan-datasources-governance-migration-"
        ) as directory:
            url = f"sqlite:///{Path(directory) / 'datasources.db'}"
            config = self._config(url)
            command.upgrade(config, "a7c1e4d9b2f6")
            engine = create_engine(url)
            try:
                metadata = MetaData()
                catalogue = Table(
                    "datasource_catalogue", metadata, autoload_with=engine
                )
                now = datetime.now(UTC)
                with engine.begin() as connection:
                    connection.execute(
                        catalogue.insert().values(
                            id="datasource-1",
                            tenant_id="tenant-1",
                            source_name="governed",
                            name="Governed",
                            kind="upload",
                            mode="static",
                            shape="tabular",
                            status="active",
                            schema_version=1,
                            schema=[],
                            fingerprint="",
                            provenance={},
                            metadata={},
                            affected_refs=["service:permit", "service:permit"],
                            dependency_refs=["dataflow:monthly"],
                            created_at=now,
                            updated_at=now,
                        )
                    )

                command.upgrade(config, "b8d2f5a0c3e7")

                references = Table(
                    "datasource_governance_references",
                    MetaData(),
                    autoload_with=engine,
                )
                with engine.connect() as connection:
                    rows = connection.execute(
                        select(
                            references.c.relation,
                            references.c.reference,
                        ).order_by(
                            references.c.relation,
                            references.c.reference,
                        )
                    ).all()
                self.assertEqual(
                    [
                        ("affected", "service:permit"),
                        ("depends_on", "dataflow:monthly"),
                    ],
                    rows,
                )
            finally:
                engine.dispose()


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