#  Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
#  NVIDIA CORPORATION and its licensors retain all intellectual property
#  and proprietary rights in and to this software, related documentation
#  and any modifications thereto.  Any use, reproduction, disclosure or
#  distribution of this software and related documentation without an express
#  license agreement from NVIDIA CORPORATION is strictly prohibited.

from idl.connection.transport.ws import WebSocketClient
from omni.discovery import DiscoverySearch
from ._generated.client import *
import asyncio


class TaggingClientHelper:
    default_fb: str = "localhost:3020"

    def __init__(self, nucleus_server: str,
                 tagging_fallback: str = default_fb,
                 subscription_fallback: str = default_fb,
                 suggestions_fallback: str = default_fb,
                 discovery_timeout: float = 10):
        self.server = nucleus_server
        self.tagging_fallback = tagging_fallback
        self.subscription_fallback = subscription_fallback
        self.suggestions_fallback = suggestions_fallback
        self.transports = {}
        self.interfaces = {}
        self.timeout = discovery_timeout

    async def __aenter__(self):
        return self
       
    async def __aexit__(self, *args, **kwargs):
        await self.close()

    async def close(self):
        try:
            for i in self.interfaces.values():
                await i.transport.close()
            for t in self.transports.values():
                await t.close()
        except BaseException as e:
            print(e)

    async def get_tagging_service(self) -> TaggingService:
        return await self._get_service(TaggingService, self.tagging_fallback)

    async def get_tag_subscription_service(self) -> TagSubscription:
        return await self._get_service(TagSubscription, self.subscription_fallback)

    async def get_tag_suggestions_service(self) -> TagSuggestions:
        return await self._get_service(TagSuggestions, self.suggestions_fallback)

    async def _get_transport(self, fallback):
        if fallback in self.transports:
            return self.transports[fallback]
        else:
            print(f'discovery failed, trying fallback connection to {fallback}')
            transport = WebSocketClient(uri=f'ws://{fallback}')
            await transport.prepare()
            self.transports[fallback] = transport
            return transport

    async def _find_service(self, interface):
        async with DiscoverySearch(self.server) as search:
            return await search.find(interface, meta={"deployment": "external"})

    async def _get_service(self, interface, fallback):
        if interface in self.interfaces:
            return self.interfaces[interface]
        else:
            instance = await self._create_interface(interface, fallback)
            self.interfaces[interface] = instance
            return instance

    async def _create_interface(self, interface, fallback):
        try:
            return await asyncio.wait_for(self._find_service(interface), timeout=self.timeout)
        except Exception:
            transport = await self._get_transport(fallback)
            return interface(transport) if transport else None
