import asyncio
import logging
import os

from typing import Dict, List

from infra.yp_service_discovery.api import api_pb2
from infra.yp_service_discovery.python.resolver.resolver import Resolver

from saas.library.python.deploy_manager_api import DeployManagerApiClient
from saas.library.python.saas_ctype import SaasCtype


class SearchProxyDiscovery:
    _CACHE_UPDATE_SECONDS = 60*30

    def __init__(self, saas_ctypes: List[SaasCtype]):
        self._ctype_to_sp_endpoints: Dict[str, List[str]] = {}
        self._saas_ctypes: List[SaasCtype] = saas_ctypes

        self._dm: DeployManagerApiClient = DeployManagerApiClient()
        self._loop: asyncio.AbstractEventLoop = asyncio.get_event_loop()
        self._first_update: asyncio.Event = asyncio.Event()

        client_id = os.getenv('PORTO_HOST', 'saas-robot@')
        self._resolver: Resolver = Resolver(client_name=f'saas_canary: {client_id}', timeout=5)

    def get_sp_endpoints(self, ctype: str) -> List[str]:
        return self._ctype_to_sp_endpoints.get(ctype, [])

    @property
    def first_update(self) -> asyncio.Event:
        return self._first_update

    async def _update_endpoints_for_ctype(self, ctype: str):
        endpoint_sets = await self._loop.run_in_executor(None, lambda: self._dm.get_proxy_endpoint_sets(ctype))
        endpoints = []

        for endpoint_set in endpoint_sets:
            cluster, id_ = endpoint_set.split('@')

            request = api_pb2.TReqResolveEndpoints()
            request.cluster_name = cluster
            request.endpoint_set_id = id_

            result = await self._loop.run_in_executor(None, lambda: self._resolver.resolve_endpoints(request))
            curr_endpoints = list(map(lambda x: f'{x.fqdn}:{x.port}', result.endpoint_set.endpoints))
            endpoints += curr_endpoints

        logging.debug('Updated searchproxy endpoints for ctype %s, new value = %s', ctype, endpoints)
        self._ctype_to_sp_endpoints[ctype] = endpoints

    async def _update_endpoints(self):
        for saas_ctype in self._saas_ctypes:
            try:
                await self._update_endpoints_for_ctype(saas_ctype.name)
            except asyncio.CancelledError:
                raise
            except Exception:
                logging.warning('Unable to update endpoints for ctype %s, skipping...', saas_ctype.name)

        self._first_update.set()

    async def run(self):
        while True:
            try:
                await self._update_endpoints()
            except asyncio.CancelledError:
                raise
            except Exception:
                logging.exception(
                    'Unable to update %s data, waiting for the next iteration...',
                    self.__class__.__name__
                )

            await asyncio.sleep(self._CACHE_UPDATE_SECONDS)
