from __future__ import annotations

import os
import threading
import unittest
import uuid
from concurrent.futures import ThreadPoolExecutor

from sqlalchemy import create_engine, text
from sqlalchemy.orm import Session

from govoplan_core.db.base import Base
from govoplan_datasources.backend.db.models import (
    DatasourceMaterializationRecord,
    DatasourcePayloadRecord,
    DatasourcePayloadRowRecord,
    DatasourceRecord,
)
from govoplan_datasources.backend.service import _append_materialization
from govoplan_datasources.backend.tabular import (
    encoded_size,
    field_payload,
    fingerprint_rows,
    infer_schema,
)


@unittest.skipUnless(
    os.environ.get("GOVOPLAN_DATASOURCES_TEST_POSTGRES_URL"),
    "set GOVOPLAN_DATASOURCES_TEST_POSTGRES_URL for PostgreSQL concurrency checks",
)
class DatasourceMaterializationPostgresTests(unittest.TestCase):
    def setUp(self) -> None:
        database_url = os.environ["GOVOPLAN_DATASOURCES_TEST_POSTGRES_URL"]
        self.schema = f"datasource_revision_race_{uuid.uuid4().hex}"
        self.admin_engine = create_engine(database_url)
        with self.admin_engine.begin() as connection:
            connection.execute(text(f'CREATE SCHEMA "{self.schema}"'))
        self.engine = create_engine(
            database_url,
            connect_args={"options": f"-c search_path={self.schema}"},
        )
        self.tables = [
            DatasourceRecord.__table__,
            DatasourcePayloadRecord.__table__,
            DatasourcePayloadRowRecord.__table__,
            DatasourceMaterializationRecord.__table__,
        ]
        Base.metadata.create_all(self.engine, tables=self.tables)
        with Session(self.engine) as session:
            session.add(
                DatasourceRecord(
                    id="datasource-1",
                    tenant_id="tenant-1",
                    source_name="concurrent",
                    name="Concurrent",
                    kind="custom",
                    mode="static",
                    shape="tabular",
                )
            )
            session.commit()

    def tearDown(self) -> None:
        try:
            Base.metadata.drop_all(
                self.engine,
                tables=list(reversed(self.tables)),
            )
        finally:
            self.engine.dispose()
            with self.admin_engine.begin() as connection:
                connection.execute(text(f'DROP SCHEMA "{self.schema}"'))
            self.admin_engine.dispose()

    def test_competing_publishers_receive_distinct_monotonic_revisions(self) -> None:
        barrier = threading.Barrier(2)

        def publish(value: int) -> int:
            rows = ({"id": value},)
            schema = infer_schema(rows)
            with Session(self.engine) as session:
                datasource = session.get(DatasourceRecord, "datasource-1")
                self.assertIsNotNone(datasource)
                barrier.wait(timeout=10)
                materialization = _append_materialization(
                    session,
                    datasource=datasource,
                    rows=rows,
                    schema=[field_payload(field) for field in schema],
                    fingerprint=fingerprint_rows(rows, schema),
                    byte_count=encoded_size(rows),
                    actor_id=f"worker-{value}",
                    set_current=True,
                )
                revision = materialization.revision
                session.commit()
                return revision

        with ThreadPoolExecutor(max_workers=2) as executor:
            revisions = tuple(executor.map(publish, (1, 2)))

        self.assertEqual((1, 2), tuple(sorted(revisions)))
        with Session(self.engine) as session:
            self.assertEqual(
                [1, 2],
                list(
                    session.scalars(
                        DatasourceMaterializationRecord.__table__.select()
                        .with_only_columns(
                            DatasourceMaterializationRecord.revision
                        )
                        .order_by(
                            DatasourceMaterializationRecord.revision.asc()
                        )
                    )
                ),
            )


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