"""Postgres container session helpers for integration tests."""

from collections.abc import Generator

import psycopg
from psycopg.rows import dict_row
from testcontainers.postgres import PostgresContainer

from python_repositories.config import PostgresConfig

POSTGRES_IMAGE = "postgres:16"
TEST_TABLE = "test_items"


_postgres_uri: str | None = None
_postgres_container: PostgresContainer | None = None
_raw_postgres_client: psycopg.Connection | None = None
_postgres_uri_refs = 0
_raw_postgres_client_refs = 0


def postgres_uri() -> Generator[str, None, None]:
    """Yield a session-scoped Postgres URI, starting the container once."""
    global _postgres_uri, _postgres_container, _postgres_uri_refs
    if _postgres_uri is None:
        _postgres_container = PostgresContainer(POSTGRES_IMAGE, driver=None)
        _postgres_container.start()
        _postgres_uri = _postgres_container.get_connection_url()

    _postgres_uri_refs += 1
    yield _postgres_uri
    _postgres_uri_refs -= 1

    if _postgres_uri_refs == 0 and _postgres_container is not None:
        _postgres_container.stop()
        _postgres_container = None
        _postgres_uri = None


def postgres_config_from_container(uri: str) -> PostgresConfig:
    """Build PostgresConfig from a container URI."""
    return PostgresConfig(uri=uri, table=TEST_TABLE, primary_key="id")


def raw_postgres_client_from_container(
    uri: str,
) -> Generator[psycopg.Connection, None, None]:
    """Yield a session-scoped raw Postgres client, reusing one client per session."""
    global _raw_postgres_client, _raw_postgres_client_refs
    if _raw_postgres_client is None:
        _raw_postgres_client = psycopg.connect(
            uri,
            row_factory=dict_row,  # type: ignore[arg-type]
            autocommit=True,
        )
        with _raw_postgres_client.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {TEST_TABLE} (
                    id TEXT PRIMARY KEY,
                    name TEXT NOT NULL,
                    value INTEGER,
                    email TEXT
                )
                """
            )

    _raw_postgres_client_refs += 1
    yield _raw_postgres_client
    _raw_postgres_client_refs -= 1

    if _raw_postgres_client_refs == 0 and _raw_postgres_client is not None:
        with _raw_postgres_client.cursor() as cursor:
            cursor.execute(f"TRUNCATE TABLE {TEST_TABLE}")
        _raw_postgres_client.close()
        _raw_postgres_client = None