import os
import threading
import time
import warnings

import duckdb

from pyulysses.connector.pyarrow_client import (
    connect_to_dremio_flight_server_endpoint,
)
from pyulysses.env_var.ulysses_env_var import get_string_env_var
from pyulysses.logger.ulysses_logger import UlyssesLogger, set_logging_level

# Suppress Arrow buffer alignment warnings
warnings.filterwarnings('ignore', category=UserWarning, module='pyarrow')
os.environ['ARROW_DEFAULT_MEMORY_POOL'] = 'system'

LOGGER = set_logging_level(
    level_string=(
        get_string_env_var(UlyssesLogger.LOGGING_LEVEL, is_mandatory=False)
        or UlyssesLogger.DEFAULT_LOGGING_LEVEL
    ),
    module_name=__name__,
)


class Client:
    """
    The Client class represents a client for the Dremio server.

    Attributes:
    username (str): The username for authentication.
    token (str): The token for authentication.
    host (str): The hostname of the Dremio server.
    port (int): The port number of the Dremio server.
    """

    def __init__(self, username=None, token=None, host=None, port=None):
        """
        The constructor for Client class.

        Parameters:
        username (str): The username for authentication. If not provided, reads from DREMIO_USERNAME env var.
        token (str): The token for authentication. If not provided, reads from DREMIO_ACCESS_TOKEN env var.
        host (str): The hostname of the Dremio server. If not provided, reads from DREMIO_HOST env var.
        port (int): The port number of the Dremio server. If not provided, reads from DREMIO_PORT env var.
        """
        self.username = username or get_string_env_var(
            'DREMIO_USERNAME', is_mandatory=True
        )
        self.token = token or get_string_env_var(
            'DREMIO_ACCESS_TOKEN', is_mandatory=True
        )
        self.host = host or get_string_env_var(
            'DREMIO_HOST', default='dremio.example.com', is_mandatory=False
        )
        self.port = port or int(
            get_string_env_var(
                'DREMIO_PORT', default='9047', is_mandatory=False
            )
        )

    def query(
        self, query, retries=1, delay=2, query_timeout=60, arrow_format=False
    ):
        """
        This function will query the datahub but taking into consideration retries and
        some delay, measured in seconds.

        Parameters:
        query (str): The query to be executed.
        retries (int): Number of retries. Defaults to 1.
        delay (int): Delay in seconds between the retries of the query. Defaults to 2.
        query_timeout (int): Timeout for each query attempt in seconds. Defaults to 60.
        arrow_format(boolean): Use of pyarrow default output format (True) or convert output to DuckDB (false)

        Returns:
        duckdb.DuckDBPyRelation or pyarrow.Table: The result of the query.
        """

        if retries > 10 or retries < 0:
            LOGGER.error(
                'Number of retries must not be lower than 0 and must be up to 10.'
            )
            return None
        retries = retries + 1

        for attempt in range(retries):
            try:
                LOGGER.debug(f'[Retry {attempt + 1}] Querying Dremio: {query}')
                return self.query_with_timeout(
                    query, query_timeout, arrow_format
                )
            except Exception as e:
                LOGGER.error(f'Querying Dremio failed. Reason: {e}')

            if attempt < retries - 1:
                LOGGER.info(f'Retrying in {delay} seconds...')
                time.sleep(delay)

        LOGGER.error('All retry attempts failed.')
        return None


class TimeoutException(Exception):
    """Custom exception for query timeout."""

    pass
