import asyncio
from collections import OrderedDict
from contextlib import asynccontextmanager
from uuid import uuid4

from app.infra.db import dynamic_manager
from app.infra.db.timescaledb import sync_pool
from app.native.wndb.core import connection


class RecordingAsyncPool:
    created: list[dict] = []

    @staticmethod
    async def check_connection(_conn) -> None:
        return None

    def __init__(self, **kwargs) -> None:
        self.created.append(kwargs)
        self.closed = False

    async def open(self) -> None:
        return None

    async def close(self) -> None:
        self.closed = True

    @asynccontextmanager
    async def connection(self):
        yield object()


class RecordingPool:
    created: list[dict] = []

    @staticmethod
    def check_connection(_conn) -> None:
        return None

    def __init__(self, **kwargs) -> None:
        self.created.append(kwargs)
        self.closed = False

    def close(self) -> None:
        self.closed = True


def test_dynamic_project_pools_check_connections_before_borrow(monkeypatch) -> None:
    async def create_pools() -> None:
        manager = dynamic_manager.ProjectConnectionManager()
        async with manager.pg_connection(
            uuid4(), "biz_data", "postgresql://user:password@db.example/biz", 1, 5
        ):
            pass
        async with manager.timescale_connection(
            uuid4(), "iot_data", "postgresql://user:password@db.example/ts", 1, 5
        ):
            pass

    RecordingAsyncPool.created = []
    monkeypatch.setattr(dynamic_manager, "AsyncConnectionPool", RecordingAsyncPool)
    asyncio.run(create_pools())

    assert len(RecordingAsyncPool.created) == 2
    assert all(
        options["check"] is dynamic_manager._check_async_connection
        for options in RecordingAsyncPool.created
    )


def test_synchronous_project_pools_check_connections_before_borrow(monkeypatch) -> None:
    RecordingPool.created = []
    monkeypatch.setattr(connection, "ConnectionPool", RecordingPool)
    monkeypatch.setattr(connection, "_pools", OrderedDict())
    monkeypatch.setattr(connection, "_pool_conninfo", {})
    monkeypatch.setattr(connection, "_pool_borrows", {})
    monkeypatch.setattr(connection, "_admin_pools", OrderedDict())
    monkeypatch.setattr(connection, "_admin_pool_borrows", {})
    monkeypatch.setattr(
        connection,
        "get_project_pgconn_string",
        lambda db_name: f"postgresql://user:password@db.example/{db_name}",
    )

    connection.get_project_pool("tjwater_next")
    connection.get_admin_pool()

    monkeypatch.setattr(sync_pool, "ConnectionPool", RecordingPool)
    monkeypatch.setattr(sync_pool, "_pools", OrderedDict())
    monkeypatch.setattr(sync_pool, "_pool_conninfo", {})
    monkeypatch.setattr(sync_pool, "_pool_borrows", {})
    monkeypatch.setattr(
        sync_pool,
        "get_project_timescale_pgconn_string",
        lambda db_name: f"postgresql://user:password@db.example/{db_name}",
    )
    sync_pool.get_timescale_pool("tjwater_next")

    assert len(RecordingPool.created) == 3
    assert [options["check"] for options in RecordingPool.created] == [
        connection._check_connection,
        connection._check_connection,
        sync_pool._check_connection,
    ]