"""Helper functions for DB connections."""
from __future__ import annotations

import json
import logging
import os
from typing import Any
from typing import Union

import pyarrow.flight as flight

from tempora.utils.exceptions import parse_arrow_exception
from tempora.utils.exceptions import TemporaClientError
from tempora.utils.exceptions import TemporaServerError
from tempora.utils.server import get_connection


log = logging.getLogger(__name__)


ConnectionParams = dict[str, Union[str, dict[str, str], None]]


def resolve_connection_params(
    *,
    args: dict[str, str | None],
    env_vars: dict[str, str],
    required_args: set[str],
    profile_params: dict[str, Any] | None = None,
) -> ConnectionParams:
    """Resolve connection params with precedence: args > profile > env."""
    resolved: ConnectionParams = {}
    profile_params = profile_params or {}

    for arg_name, arg_value in args.items():
        value = normalize_param_value(arg_value)

        if value is None:
            value = normalize_param_value(profile_params.get(arg_name))

        env_var = env_vars.get(arg_name)
        if value is None and env_var:
            value = normalize_param_value(os.getenv(env_var))

        if value is None and arg_name in required_args:
            if env_var:
                raise TemporaClientError(
                    f'No value passed for "{arg_name}" argument and environment '
                    f'variable {env_var} is not set.',
                )
            raise TemporaClientError(
                f'No value passed for required "{arg_name}" argument.',
            )

        resolved[arg_name] = value

    return resolved


def normalize_param_value(value: Any) -> str | None:
    if value is None:
        return None
    if isinstance(value, str):
        value = value.strip()
        return value if value else None
    return str(value)


def open_db_connection(dbms_name: str, params: ConnectionParams) -> None:
    connection_params = params.copy()
    connection_params['dbms_name'] = dbms_name
    json_params = json.dumps(connection_params)
    action = flight.Action('open_db_connection', json_params.encode())
    try:
        client, auth_token = get_connection()
        call_options = flight.FlightCallOptions(headers=[auth_token])
        # Server responds with an empty result iter upon success
        next(client.do_action(action, call_options))
    except flight.FlightError as e:
        msg = f'Problem connecting to {dbms_name}: {parse_arrow_exception(e)}'
        raise TemporaServerError(msg) from None
    except StopIteration:
        # Empty response from server - all good
        ...
    return None


def check_db_connection(dbms_name: str) -> bool:
    # Check the status of a DB connection
    action = flight.Action('check_db_connection', dbms_name.encode())
    try:
        client, auth_token = get_connection()
        call_options = flight.FlightCallOptions(headers=[auth_token])
        result = next(client.do_action(action, call_options))
        return bool(result.body.to_pybytes())
    except flight.FlightError as e:
        msg = f'Problem validating {dbms_name} connection: {parse_arrow_exception(e)}'
        log.error(msg)
        raise TemporaServerError(msg) from None
    except StopIteration:
        raise TemporaServerError('No response from server...') from None


def close_db_connection(dbms_name: str) -> None:
    # Close the DB connection on the server
    action = flight.Action('close_db_connection', dbms_name.encode())
    try:
        client, auth_token = get_connection()
        call_options = flight.FlightCallOptions(headers=[auth_token])
        # Server responds with an empty result iter upon success
        next(client.do_action(action, call_options))
    except flight.FlightError as e:
        msg = f'Problem closing {dbms_name} connection: {parse_arrow_exception(e)}'
        log.error(msg)
        raise TemporaServerError(msg) from None
    except StopIteration:
        # Empty response from server - all good
        ...
    return None
