import asyncio
import functools

from typing import Any, List, Callable, Tuple, Optional

from kazoo.client import KazooClient, ReadLock, WriteLock
from kazoo.protocol.states import ZnodeStat

SAAS_ZK_HOSTS = """
    saas-zookeeper1.search.yandex.net:14880,
    saas-zookeeper2.search.yandex.net:14880,
    saas-zookeeper3.search.yandex.net:14880,
    saas-zookeeper4.search.yandex.net:14880,
    saas-zookeeper5.search.yandex.net:14880
"""


class ZKClient:
    _TIMEOUT = 10

    def __init__(self, hosts: str = SAAS_ZK_HOSTS, *args, **kwargs) -> None:
        self._kazoo_client: KazooClient = KazooClient(hosts, *args, **kwargs)

    async def _wait_for_execution(self, func: Callable) -> Any:
        loop = asyncio.get_event_loop()

        try:
            data = await loop.run_in_executor(None, functools.partial(func, timeout=self._TIMEOUT))
        except TimeoutError:
            raise OperationTimeout('Unable to complete the operation within the given timeout=%d', self._TIMEOUT)

        return data

    async def _start(self) -> None:
        result = self._kazoo_client.start_async()
        await self._wait_for_execution(result.wait)

    async def __aenter__(self) -> 'ZKClient':
        await self._start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        self._kazoo_client.stop()

    async def set(self, path: str, value: bytes) -> ZnodeStat:
        result = self._kazoo_client.set_async(path, value)
        return await self._wait_for_execution(result.get)

    async def get(self, path: str) -> Tuple[bytes, ZnodeStat]:
        result = self._kazoo_client.get_async(path)
        return await self._wait_for_execution(result.get)

    async def list(self, path: str) -> List[ZnodeStat]:
        result = self._kazoo_client.get_children_async(path)
        children = await self._wait_for_execution(result.get)
        return children

    async def ensure_path(self, *args, **kwargs) -> Optional[ZnodeStat]:
        result = self._kazoo_client.ensure_path_async(*args, **kwargs)
        return await self._wait_for_execution(result.get)

    def create_write_lock(self, *args, **kwargs) -> WriteLock:
        return self._kazoo_client.WriteLock(*args, **kwargs)

    def create_read_lock(self, *args, **kwargs) -> ReadLock:
        return self._kazoo_client.ReadLock(*args, **kwargs)


class OperationTimeout(TimeoutError):
    ...
