"""MotherDuck 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 _MotherDuckDatasetArgs(_DBDatasetArgs):
    ...


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

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

    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 FilteredMotherDuckDataset(FilterMixin, MotherDuckDataset):
    """Class representing a filtered MotherDuck-based dataset."""

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


MotherDuckDataset.filtered_dataset_cls = FilteredMotherDuckDataset


@validate_call
def connect_motherduck(
    db_name: str,
    auth_token: str | None = None,
) -> None:
    """Open a connection to MotherDuck.

    Args:
        db_name: MotherDuck database name.
        auth_token: Authentication token for MotherDuck. If omitted, this must be
            specified using the MOTHERDUCK_AUTH_TOKEN env var.

    Returns:
        None
    """
    params = _get_motherduck_connection_params(db_name=db_name, auth_token=auth_token)
    open_db_connection(dbms_name='MotherDuck', params=params)
    return None


def _get_motherduck_connection_params(
    db_name: str,
    auth_token: str | None = None,
) -> ConnectionParams:
    params = resolve_connection_params(
        args={'db_name': db_name, 'auth_token': auth_token},
        env_vars={'auth_token': 'MOTHERDUCK_AUTH_TOKEN'},
        required_args={'auth_token'},
    )
    # db_name is required as an argument by connect_motherduck; keep this payload
    # explicit and stable for server dispatch.
    return {
        'db_name': params['db_name'], 'auth_token': params['auth_token'],
    }


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