import subprocess
from functools import partial, wraps
from pathlib import Path
from typing import (
    Any,
    Awaitable,
    Callable,
    Coroutine,
    Iterable,
    Optional,
    TypeVar,
    Union,
    overload,
)

import requests
import tqdm  # type: ignore
import trio

T = TypeVar("T")
U = TypeVar("U")
V = TypeVar("V")


Some = Optional[Union[T, list[T]]]


def unsome_(x: Some[T]) -> Optional[list[T]]:
    if x is None:
        return None
    if isinstance(x, list):
        return x
    return [x]


def unsome(x: Some[T]) -> list[T]:
    return unsome_(x) or []


def intercalate(delimiter: T, items: Iterable[Iterable[T]]) -> list[T]:
    it = iter(items)
    try:
        result = list(next(it))
    except StopIteration:
        return []
    for x in it:
        result.append(delimiter)
        result.extend(list(x))
    return result


TERMINAL_LOCK = trio.Lock()


RunArg = Union[str, Path]


async def run(*args: RunArg) -> subprocess.CompletedProcess:
    return await trio.run_process(args, check=True)


async def run_ok(*args: RunArg) -> bool:
    try:
        result = await trio.run_process(
            args, check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )
        return result.returncode == 0
    except FileNotFoundError:
        return False


async def run_output(*args: RunArg, silent: bool = False) -> str:
    result = await trio.run_process(
        args, capture_stdout=True, capture_stderr=silent, check=True
    )
    return result.stdout.decode().strip()


def flatten(items: Iterable[Iterable[T]]) -> list[T]:
    return [item for sublist in items for item in sublist]


async def parallel_map_tqdm(
    action: Callable[[T], Awaitable[U]], items: list[T]
) -> list[U]:
    with tqdm.tqdm(total=len(items)) as progress:

        async def action_and_update(item: T) -> U:
            result = await action(item)
            async with TERMINAL_LOCK:
                progress.update(1)
            return result

        return await gather(*(partial(action_and_update, item) for item in items))


async def gather(*tasks: Callable[[], Awaitable[T]]) -> list[T]:
    async def collect(
        index: int, task: Callable[[], Awaitable[T]], results: dict[int, T]
    ):
        results[index] = await task()

    results: dict[int, T] = {}
    async with trio.open_nursery() as nursery:
        for index, task in enumerate(tasks):
            nursery.start_soon(collect, index, task, results)
    return [results[i] for i in range(len(tasks))]


def url_version(url: str) -> str:
    head_response = requests.head(url, allow_redirects=True)
    return head_response.headers["etag"]


def raise_(exception: BaseException) -> Any:
    raise exception


AsyncRet = Coroutine[Any, Any, T]


@overload
def async_cached(fn: Callable[[], AsyncRet[T]]) -> Callable[[], AsyncRet[T]]:
    ...


@overload
def async_cached(fn: Callable[[U], AsyncRet[T]]) -> Callable[[U], AsyncRet[T]]:
    ...


@overload
def async_cached(fn: Callable[[U, V], AsyncRet[T]]) -> Callable[[U, V], AsyncRet[T]]:
    ...


def async_cached(fn: Callable[..., AsyncRet[T]]) -> Callable[..., AsyncRet[T]]:
    return _async_cached_lock(None, fn)


@overload
def async_cached_lock(fn: Callable[[], AsyncRet[T]]) -> Callable[[], AsyncRet[T]]:
    ...


@overload
def async_cached_lock(fn: Callable[[U], AsyncRet[T]]) -> Callable[[U], AsyncRet[T]]:
    ...


@overload
def async_cached_lock(
    fn: Callable[[U, V], AsyncRet[T]]
) -> Callable[[U, V], AsyncRet[T]]:
    ...


def async_cached_lock(fn: Callable[..., AsyncRet[T]]) -> Callable[..., AsyncRet[T]]:
    return _async_cached_lock(trio.Lock(), fn)


class NoLock:
    async def __aenter__(self) -> None:
        pass

    async def __aexit__(self, *args: Any) -> None:
        pass


NO_LOCK = NoLock()


def _async_cached_lock(
    lock: Optional[trio.Lock], fn: Callable[..., AsyncRet[T]]
) -> Callable[..., AsyncRet[T]]:
    cache: dict[Any, T] = {}

    @wraps(fn)
    async def cached_fn(*args: Any) -> T:
        if args not in cache:
            async with lock or NO_LOCK:
                # Check again, someone holding the lock might have updated the cache
                if args not in cache:
                    cache[args] = await fn(*args)
        return cache[args]

    return cached_fn
