"""Functions for interacting with a tempora server.

- login(host, username, password): login to a tempora server
- logout(): logout from the server
- list_datasets(): list datasets currently loaded on the server
"""
from __future__ import annotations

import contextvars
import logging
import os
from contextlib import contextmanager
from getpass import getpass
from importlib import resources
from pathlib import Path
from typing import cast
from typing import Iterator
from typing import Literal
from typing import NoReturn
from typing import TypedDict

import pyarrow.flight as flight

from .exceptions import parse_arrow_exception
from .exceptions import TemporaClientError
from .exceptions import TemporaServerError


log = logging.getLogger(__name__)


class SessionParams(TypedDict):
    client: flight.FlightClient | None
    auth_token: tuple[str, str] | None


class ConnectionArgs(TypedDict):
    host: str
    username: str
    password: str | None


SESSION: contextvars.ContextVar[SessionParams] = contextvars.ContextVar(
    'session', default={'client': None, 'auth_token': None},
)


def login(
    host: str | None = None, username: str | None = None, password: str | None = None,
    use_tls: bool = True, verify_server: bool = True,
) -> None:
    """Login to the tempora server instance.

    Args:
        host: Host URL of the server in "hostname:port" format. If omitted, the value of
            the TEMPORA_SERVER_HOST environment variable is used. If the port is not
            specified in the host URL it defaults to 7070.
        username: Username to authenticate with the server. If omitted, the value of
            the TEMPORA_SERVER_USERNAME environment variable is used.
        password: Password used to authenticate with the server. If omitted, the
            value of the `TEMPORA_SERVER_PASSWORD` environment variable is used.
            If that is also unset, the user is prompted to enter the password.
        use_tls: Set this to False if the server is running without TLS enabled.
        verify_server: Set this to False to disable server verification when using TLS.
    """
    con_args = _process_connection_args(
        host=host, username=username, password=password,
    )
    _authenticate_with_server(
        con_args=con_args, use_tls=use_tls, verify_server=verify_server,
    )
    return None


def _process_connection_args(
    host: str | None = None, username: str | None = None, password: str | None = None,
) -> ConnectionArgs:
    # Process the arguments: try to read the environment variable if passed arg is None
    args_dict = {'host': host, 'username': username, 'password': password}
    con_args: ConnectionArgs = {'host': '', 'username': '', 'password': None}

    for arg_name, arg_value in args_dict.items():
        arg_name = cast(Literal['host', 'username', 'password'], arg_name)
        if arg_value is not None:
            con_args[arg_name] = arg_value
        else:
            env_var = f'TEMPORA_SERVER_{arg_name.upper()}'
            env_value = os.environ.get(env_var)
            if env_value:
                con_args[arg_name] = env_value
            elif arg_name != 'password':
                msg = (
                    f'No value passed for "{arg_name}" argument and '
                    f'environment variable {env_var} is not set...'
                )
                log.error(msg)
                raise TemporaClientError(msg)

    return con_args


def _authenticate_with_server(
    con_args: ConnectionArgs, use_tls: bool = True, verify_server: bool = True,
) -> contextvars.Token | NoReturn:
    # Append the default port to the hostname if one was not specified
    host = con_args['host']
    if ':' not in host:
        host += ':7070'

    tls_root_certs: bytes | None = None
    if use_tls and verify_server:
        # Get the TLS root certificate
        p = os.getenv('TEMPORA_TLS_ROOT_CERT_PATH')
        if p:
            tls_root_certs = Path(p).read_bytes()
        else:
            # Fallback to the bundled trial certificate
            resource = resources.files('tempora.resources.certs')
            tls_root_certs = resource.joinpath('tempora_trial_ca.pem').read_bytes()

    scheme = 'grpc+tls' if use_tls else 'grpc'
    client = flight.connect(
        location=f'{scheme}://{host}',
        tls_root_certs=tls_root_certs,
        disable_server_verification=not verify_server,
    )

    # Before trying to authenticate check that we can communicate with the server, see:
    # https://github.com/apache/arrow/blob/main/python/examples/flight/client.py#L178
    try:
        client.wait_for_available()
    except flight.FlightUnauthenticatedError:
        # The server is up and running at the given host:port address
        log.info(f'Successfully connected to server at {host}')
    except flight.FlightUnavailableError as e:
        flight_msg_prefix = 'Flight returned unavailable error, with message:'
        msg = e.args[0].replace(flight_msg_prefix, '').strip()
        msg = f'Unable to connect to server at {host}. {msg}'
        log.error(msg)
        client.close()
        raise TemporaServerError(msg) from None

    # Authenticate
    try:
        if not con_args['password']:
            con_args['password'] = getpass(
                prompt=f'Enter password for user {con_args["username"]}: ',
            )
        assert con_args['password'] is not None
        auth_token = client.authenticate_basic_token(
            con_args['username'], con_args['password'],
        )
        reset_token = SESSION.set({'client': client, 'auth_token': auth_token})
        msg = f'Successfully authenticated with server at {host}'
        print(msg)
        log.info(msg)
    except flight.FlightUnauthenticatedError as e:
        client.close()
        msg = f'Unable to authenticate with server at {host}. {e.args[0]}'
        log.error(msg)
        raise TemporaServerError(msg) from None

    return reset_token


def logout() -> None:
    """Logout from the tempora server that the user has previously connected to."""
    action = flight.Action('logout', b'')
    client, auth_token = get_connection()
    call_options = flight.FlightCallOptions(headers=[auth_token])
    try:
        next(client.do_action(action, call_options))
    except flight.FlightError as e:
        client.close()
        msg = parse_arrow_exception(e)
        log.error(msg)
        raise TemporaServerError(msg) from None
    except StopIteration:
        # Empty response from server - all good
        client.close()
        SESSION.set({'client': None, 'auth_token': None})  # Clear the session ctx var
        return None


@contextmanager
def server_connection(
    host: str | None = None, username: str | None = None, password: str | None = None,
    use_tls: bool = True, verify_server: bool = True,
) -> Iterator[None]:
    """Context manager to manage the connection to the tempora server instance.

    Args:
        host: Host URL of the server in "hostname:port" format. If omitted, the value of
            the TEMPORA_SERVER_HOST environment variable is used. If the port is not
            specified in the host URL it defaults to 7070.
        username: Username to authenticate with the server. If omitted, the value of
            the TEMPORA_SERVER_USERNAME environment variable is used.
        password: Password used to authenticate with the server. If omitted, the
            value of the `TEMPORA_SERVER_PASSWORD` environment variable is used.
            If that is also unset, the user is prompted to enter the password.
        use_tls: Set this to False if the server is running without TLS enabled.
        verify_server: Set this to False to disable server verification when using TLS.
    """
    con_args = _process_connection_args(
        host=host, username=username, password=password,
    )
    reset_token = _authenticate_with_server(
        con_args=con_args, use_tls=use_tls, verify_server=verify_server,
    )
    try:
        yield
    finally:
        try:
            logout()
        finally:
            # Reset the SESSION var to its previous state even if logout() fails
            SESSION.reset(reset_token)


def get_connection() -> tuple[flight.FlightClient, tuple[str, str]] | NoReturn:
    """Get the server connection details.

    Returns:
        A tuple containing the FlightClient object and the authenication token for the
        server.
    """
    session = SESSION.get()
    if session['client'] is None:
        msg = 'No server connection. Please authenticate by calling: tempora.login()'
        raise TemporaClientError(msg)
    if session['auth_token'] is None:
        msg = 'Missing authentication token. Call tempora.login() to re-authenticate.'
        raise TemporaClientError(msg)
    return session['client'], session['auth_token']


def list_datasets() -> None:
    """List datasets currently loaded on the server."""
    client, auth_token = get_connection()
    try:
        datasets_found = False
        call_options = flight.FlightCallOptions(headers=[auth_token])
        for flight_info in client.list_flights(options=call_options):
            datasets_found = True
            dataset_desc = flight_info.descriptor.command.decode()
            ref_count = flight_info.total_records
            sfx = 's' if ref_count > 1 else ''
            print(f'Referenced by {ref_count} client dataset{sfx}:\n{dataset_desc}\n')
        if not datasets_found:
            print('No datasets on the server')
    except flight.FlightError as e:
        msg = f'Problem encountered in "list_datasets": {parse_arrow_exception(e)}'
        raise TemporaServerError(msg) from None
    except StopIteration:
        raise TemporaServerError('No response from server...') from None
    return None
