import asyncio
from typing import Callable, TypeVar

from wandb.sdk import wandb_setup
from wandb.sdk.lib import asyncio_compat, printer

_T = TypeVar("_T")


def run_async_with_spinner(
    spinner_printer: printer.Printer,
    text: str,
    func: Callable[[], _T],
) -> _T:
    """Run a slow function while displaying a loading icon.

    Args:
        spinner_printer: The printer to use to display text.
        text: The text to display next to the spinner while the function runs.
        func: The function to run.

    Returns:
        The result of func.
    """

    async def loop_run_with_spinner() -> _T:
        async def update_spinner() -> None:
            tick = 0
            with spinner_printer.dynamic_text() as text_area:
                if not text_area:
                    spinner_printer.display(text)
                    return

                while True:
                    spinner = spinner_printer.loading_symbol(tick)
                    text_area.set_text(f"{spinner} {text}")
                    tick += 1
                    await asyncio.sleep(0.1)

        async with asyncio_compat.cancel_on_exit(update_spinner()):
            return await asyncio.get_running_loop().run_in_executor(None, func)

    asyncer = wandb_setup.singleton().asyncer
    return asyncer.run(loop_run_with_spinner)