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):
    def test_csv_evidence_upgrade_preserves_legacy_rows_and_fingerprints(self) -> None:
        with tempfile.TemporaryDirectory(prefix="govoplan-datasources-csv-migration-") as directory:
            url = f"sqlite:///{Path(directory) / 'datasources.db'}"
            config = self._config(url)
            command.upgrade(config, "d1a7c3e9f5b2")
            engine = create_engine(url)
            try:
                metadata = MetaData()
                tables = {name: Table(name, metadata, autoload_with=engine) for name in (
                    "datasource_catalogue", "datasource_stages", "datasource_materializations",
                )}
                now = datetime.now(UTC)
                common = dict(tenant_id="tenant-1", schema=[], fingerprint="a" * 64,
                              metadata={"retained": True}, provenance={}, created_at=now, updated_at=now)
                with engine.begin() as connection:
                    connection.execute(tables["datasource_catalogue"].insert().values(
                        **common, id="legacy-source", source_name="legacy", name="Legacy", kind="upload",
                        mode="static", shape="tabular", status="active", schema_version=1,
                    ))
                    connection.execute(tables["datasource_stages"].insert().values(
                        **common, id="legacy-stage", source_name="staged", name="Staged", kind="upload",
                        mode="static", shape="tabular", state="ready", rows=[{"id": 1}],
                        row_count=1, byte_count=10, validation={}, governance={}, approval={},
                    ))
                    connection.execute(tables["datasource_materializations"].insert().values(
                        **common, id="legacy-materialization", datasource_id="legacy-source", revision=1,
                        state="published", schema_version=1, rows=[{"id": 1}], row_count=1,
                        byte_count=10, governance_snapshot={},
                    ))
                    before = {name: dict(connection.execute(select(table)).mappings().one()) for name, table in tables.items()}
                command.upgrade(config, "e2b8d4a0f6c3")
                with engine.connect() as connection:
                    for name in tables:
                        upgraded = Table(name, MetaData(), autoload_with=connection)
                        after = dict(connection.execute(select(upgraded)).mappings().one())
                        if name != "datasource_catalogue":
                            self.assertIsNone(after.pop("csv_source"))
                        self.assertEqual(before[name], after)
            finally:
                engine.dispose()

    @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(
                        "e2b8d4a0f6c3",
                        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",
                            "approval_policy",
                            "retention_policy",
                            "dependency_refs",
                        }.issubset(catalogue_columns)
                    )
                    stage_columns = {
                        item["name"]
                        for item in inspect(connection).get_columns(
                            "datasource_stages"
                        )
                    }
                    materialization_columns = {
                        item["name"]
                        for item in inspect(connection).get_columns(
                            "datasource_materializations"
                        )
                    }
                    self.assertIn("approval", stage_columns)
                    self.assertTrue(
                        {"disposed_at", "disposition"}.issubset(
                            materialization_columns
                        )
                    )
                    self.assertEqual(
                        {
                            "datasource_catalogue",
                            "datasource_governance_references",
                            "datasource_lifecycle_evidence",
                            "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()