"""PostgreSQL connection helpers."""
from __future__ import annotations

import dataclasses as dc
from typing import Any
from typing import ClassVar

import pydantic.dataclasses as pydc
from pydantic import validate_call

from .base import _DBDatasetArgs
from .base import DatasetBase
from .filter import FilterMixin
from tempora.utils.dataclass import pydc_config
from tempora.utils.db import close_db_connection
from tempora.utils.db import ConnectionParams
from tempora.utils.db import open_db_connection
from tempora.utils.db import resolve_connection_params


@dc.dataclass(frozen=True)
class _PostgreSQLDatasetArgs(_DBDatasetArgs):
    ...


@pydc.dataclass(frozen=True, eq=False, config=pydc_config)
class PostgreSQLDataset(DatasetBase, _PostgreSQLDatasetArgs):
    """Class representing a PostgreSQL-backed dataset.

    Args:
        database: The name of the PostgreSQL database.
        table: The name of the table under the given database.
        db_schema: Optional schema name.
    """
    filtered_dataset_cls: ClassVar[type[FilteredPostgreSQLDataset]]

    def __post_init__(self, _op_dict: dict[str, Any] | None) -> None:
        super().__post_init__(_op_dict=_op_dict)


@pydc.dataclass(frozen=True, eq=False, config=pydc_config)
class FilteredPostgreSQLDataset(FilterMixin, PostgreSQLDataset):
    """Class representing a filtered PostgreSQL-based dataset."""

    def __post_init__(self, _op_dict: dict[str, Any] | None) -> None:
        super().__post_init__(_op_dict=_op_dict)


PostgreSQLDataset.filtered_dataset_cls = FilteredPostgreSQLDataset


@validate_call
def connect_postgresql(
    host: str | None = None,
    user: str | None = None,
    password: str | None = None,
    port: str | None = None,
    role: str | None = None,
) -> None:
    """Open a connection to PostgreSQL.

    Args:
        host: PostgreSQL server host. If omitted, this must be specified using the
            POSTGRESQL_HOST env var.
        user: Username for PostgreSQL authentication. If omitted, this must be
            specified using the POSTGRESQL_USER env var.
        password: Password for PostgreSQL authentication. If omitted, this must be
            specified using the POSTGRESQL_PASSWORD env var.
        port: Server port. If omitted, uses POSTGRESQL_PORT or defaults to '5432'.
        role: Optional session role. If omitted, uses POSTGRESQL_ROLE.

    Returns:
        None
    """
    params = _get_postgresql_connection_params(
        host=host,
        user=user,
        password=password,
        port=port,
        role=role,
    )
    open_db_connection(dbms_name='PostgreSQL', params=params)
    return None


def _get_postgresql_connection_params(
    host: str | None = None,
    user: str | None = None,
    password: str | None = None,
    port: str | None = None,
    role: str | None = None,
) -> ConnectionParams:
    params = resolve_connection_params(
        args={
            'host': host,
            'user': user,
            'password': password,
            'port': port,
            'role': role,
        },
        env_vars={
            'host': 'POSTGRESQL_HOST',
            'user': 'POSTGRESQL_USER',
            'password': 'POSTGRESQL_PASSWORD',
            'port': 'POSTGRESQL_PORT',
            'role': 'POSTGRESQL_ROLE',
        },
        required_args={'host', 'user', 'password'},
    )
    if not params.get('port'):
        params['port'] = '5432'
    return params


def disconnect_postgresql() -> None:
    """Close the currently open PostgreSQL connection."""
    return close_db_connection(dbms_name='PostgreSQL')
