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

import configparser
import dataclasses as dc
import os
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
from tempora.utils.exceptions import TemporaClientError


@dc.dataclass(frozen=True)
class _DatabricksDatasetArgs(_DBDatasetArgs):
    catalog: str | None = None


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

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

    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 FilteredDatabricksDataset(FilterMixin, DatabricksDataset):
    """Class representing a filtered Databricks-based dataset."""

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


DatabricksDataset.filtered_dataset_cls = FilteredDatabricksDataset


@validate_call
def connect_databricks(
    host: str | None = None,
    auth_token: str | None = None,
    http_path: str | None = None,
    warehouse: str | None = None,
    port: str | None = None,
    profile: str | None = None,
    connection_name: str | None = None,
) -> None:
    """Open a connection to Databricks.

    Args:
        host: Databricks workspace host. If omitted, this must be specified using
            the DATABRICKS_HOST env var or a named Databricks profile.
        auth_token: Databricks authentication token (typically a PAT). If omitted,
            this must be specified using the DATABRICKS_AUTH_TOKEN env var or a named
            Databricks profile.
        http_path: SQL warehouse or endpoint HTTP path. Optional if warehouse is
            provided. If omitted, uses DATABRICKS_HTTP_PATH or profile value.
        warehouse: Warehouse identifier or name. Optional if http_path is provided.
            If omitted, uses DATABRICKS_WAREHOUSE or profile value.
        port: Workspace port. If omitted, uses DATABRICKS_PORT, profile value, or
            defaults to '443'.
        profile: Name of the profile in .databrickscfg to load defaults from.
        connection_name: Alias for profile.

    Returns:
        None
    """
    params = _get_databricks_connection_params(
        host=host,
        auth_token=auth_token,
        http_path=http_path,
        warehouse=warehouse,
        port=port,
        profile=profile,
        connection_name=connection_name,
    )
    open_db_connection(dbms_name='Databricks', params=params)
    return None


def _get_databricks_connection_params(
    host: str | None = None,
    auth_token: str | None = None,
    http_path: str | None = None,
    warehouse: str | None = None,
    port: str | None = None,
    profile: str | None = None,
    connection_name: str | None = None,
) -> ConnectionParams:
    profile_params = _load_databricks_profile(profile_name=profile or connection_name)
    params = resolve_connection_params(
        args={
            'host': host,
            'auth_token': auth_token,
            'http_path': http_path,
            'warehouse': warehouse,
            'port': port,
        },
        env_vars={
            'host': 'DATABRICKS_HOST',
            'auth_token': 'DATABRICKS_AUTH_TOKEN',
            'http_path': 'DATABRICKS_HTTP_PATH',
            'warehouse': 'DATABRICKS_WAREHOUSE',
            'port': 'DATABRICKS_PORT',
        },
        required_args={'host', 'auth_token'},
        profile_params=profile_params,
    )

    if not params.get('http_path') and not params.get('warehouse'):
        raise TemporaClientError(
            'Databricks requires at least one of "http_path" or "warehouse".',
        )
    if not params.get('port'):
        params['port'] = '443'

    return params


def _load_databricks_profile(profile_name: str | None) -> dict[str, str]:
    if not profile_name:
        return {}

    cfg_path = os.path.expanduser(
        os.getenv('DATABRICKS_CONFIG_FILE', '~/.databrickscfg'),
    )

    if not os.path.exists(cfg_path):
        raise TemporaClientError(f'Missing Databricks config file: {cfg_path}')

    parser = configparser.ConfigParser()
    parser.read(cfg_path)

    if profile_name == 'DEFAULT':
        profile_values = dict(parser.defaults())
    elif parser.has_section(profile_name):
        profile_values = dict(parser[profile_name])
    else:
        raise TemporaClientError(
            f'Missing Databricks profile "{profile_name}" in config file: {cfg_path}',
        )

    params: dict[str, str] = {}
    if 'host' in profile_values:
        params['host'] = profile_values['host']
    if 'auth_token' in profile_values:
        params['auth_token'] = profile_values['auth_token']
    elif 'token' in profile_values:
        params['auth_token'] = profile_values['token']
    if 'http_path' in profile_values:
        params['http_path'] = profile_values['http_path']
    if 'warehouse' in profile_values:
        params['warehouse'] = profile_values['warehouse']
    if 'port' in profile_values:
        params['port'] = profile_values['port']

    return params


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