import contextlib
import logging
import os
import socket
import time
from typing import Any

with contextlib.suppress(ImportError):  # Module not available when building docs
    from kola.kola import KolaConnector, KolaIOError

logger = logging.getLogger("kola")


# Q use IPC protocol version 6, which is compatible with kdb+
class Q(object):
    def __init__(
        self,
        host: str,
        port: int,
        user: str = "",
        passwd: str = "",
        enable_tls: bool = False,
        retries: int = 0,
        timeout: int = 0,
    ):
        if not user:
            try:
                user = os.getlogin()
            except Exception:
                user = "unknown"
        if (not host) or host == socket.gethostname():
            host = "127.0.0.1"
        self.host = host
        self.port = port
        self.user = user
        self.retries = retries
        self.q = KolaConnector(host, port, user, passwd, enable_tls, timeout, 6)

    def connect(self):
        self.q.connect()

    def disconnect(self):
        self.q.shutdown()

    def sync(self, expr: str, *args) -> Any:
        if self.retries <= 0:
            return self.q.sync(expr, *args)
        else:
            n = 0
            # exponential backoff
            while n < self.retries:
                try:
                    return self.q.sync(expr, *args)
                except KolaIOError as e:
                    logging.info(
                        "Failed to sync - '%s', retrying in %s seconds", e, 2**n
                    )
                    time.sleep(2**n)
                    n += 1
                    if n == self.retries:
                        raise (e)

    def asyn(self, expr: str, *args):
        if self.retries <= 0:
            return self.q.asyn(expr, *args)
        else:
            n = 0
            # exponential backoff
            while n < self.retries:
                try:
                    return self.q.asyn(expr, *args)
                except KolaIOError as e:
                    logging.info(
                        "Failed to async - '%s', retrying in %s seconds", e, 2**n
                    )
                    time.sleep(2**n)
                    n += 1
                    if n == self.retries:
                        raise (e)

    def receive(self) -> Any:
        if self.retries <= 0:
            return self.q.receive()
        else:
            n = 0
            # exponential backoff
            while n < self.retries:
                try:
                    return self.q.receive()
                except KolaIOError as e:
                    logging.info(
                        "Failed to receive - '%s', retrying in %s seconds", e, 2**n
                    )
                    self.connect()
                    time.sleep(2**n)
                    n += 1
                    if n == self.retries:
                        raise (e)
