## Copyright (c) 2021, 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.
##
import asyncio
import omni.client

from typing import Callable


class CheckpointHelper:

    server_cache = {}

    @staticmethod
    def extract_server_from_url(url: str) -> str:
        server_url = ""
        if url:
            client_url = omni.client.break_url(url)
            server_url = omni.client.make_url(scheme=client_url.scheme, host=client_url.host, port=client_url.port)
        return server_url

    @staticmethod
    async def is_checkpoint_enabled_async(url: str) -> bool:
        server_url = CheckpointHelper.extract_server_from_url(url)
        if not server_url:
            enabled = False
        elif server_url in CheckpointHelper.server_cache:
            enabled = CheckpointHelper.server_cache[server_url]
        else:
            result, server_info = await omni.client.get_server_info_async(server_url)
            supports_checkpoint = result and server_info and server_info.checkpoints_enabled
            CheckpointHelper.server_cache[server_url] = supports_checkpoint
            enabled = supports_checkpoint
        return enabled

    @staticmethod
    def is_checkpoint_enabled_with_callback(url: str, callback: Callable):
        async def is_checkpoint_enabled_async(url: str, callback: Callable):
            enabled = await CheckpointHelper.is_checkpoint_enabled_async(url)
            if callback:
                callback(CheckpointHelper.extract_server_from_url(url), enabled)
        asyncio.ensure_future(is_checkpoint_enabled_async(url, callback))