# MIT licensed
# Copyright (c) 2020 lilydjwg <lilydjwg@gmail.com>, et al.

from __future__ import annotations

import asyncio
from asyncio import Queue
import contextlib
from typing import (
  Dict, Optional, List, AsyncGenerator, NamedTuple, Union,
  Any, Tuple, Callable, TypeVar, Coroutine, Generic,
  TYPE_CHECKING,
)
from pathlib import Path

import toml
import structlog

from .httpclient import session # type: ignore

logger = structlog.get_logger(logger_name=__name__)

Entry = Dict[str, Any]
Entries = Dict[str, Entry]
VersData = Dict[str, str]
VersionResult = Union[None, str, List[str], Exception]

class KeyManager:
  def __init__(
    self, file: Optional[Path],
  ) -> None:
    if file is not None:
      with file.open() as f:
        keys = toml.load(f)['keys']
    else:
      keys = {}
    self.keys = keys

  def get_key(self, name: str) -> Optional[str]:
    return self.keys.get(name)

class RawResult(NamedTuple):
  name: str
  version: VersionResult
  conf: Entry

class Result(NamedTuple):
  name: str
  version: str
  conf: Entry

class BaseWorker:
  def __init__(
    self,
    token_q: Queue[bool],
    result_q: Queue[RawResult],
    tasks: List[Tuple[str, Entry]],
    tries: int,
    keymanager: KeyManager,
  ) -> None:
    self.token_q = token_q
    self.result_q = result_q
    self.tries = tries
    self.keymanager = keymanager
    self.tasks = tasks

  @contextlib.asynccontextmanager
  async def acquire_token(self) -> AsyncGenerator[None, None]:
    token = await self.token_q.get()
    logger.debug('got token')
    try:
      yield
    finally:
      await self.token_q.put(token)
      logger.debug('return token')

T = TypeVar('T')
S = TypeVar('S')

class AsyncCache(Generic[T, S]):
  cache: Dict[T, Union[S, asyncio.Task]]
  lock: asyncio.Lock

  def __init__(self) -> None:
    self.cache = {}
    self.lock = asyncio.Lock()

  async def _get_json(self, key: Tuple[str, str]) -> Any:
    url = key[1]
    async with session.get(url) as res:
      return await res.json(content_type=None)

  async def get_json(self, url: str) -> Any:
    return await self.get(
      ('_jsonurl', url), self._get_json) # type: ignore

  async def get(
    self,
    key: T,
    func: Callable[[T], Coroutine[None, None, S]],
  ) -> S:
    async with self.lock:
      cached = self.cache.get(key)
      if cached is None:
        coro = func(key)
        fu = asyncio.create_task(coro)
        self.cache[key] = fu

    if asyncio.isfuture(cached): # pending
      return await cached # type: ignore
    elif cached is not None: # cached
      return cached # type: ignore
    else: # not cached
      r = await fu
      self.cache[key] = r
      return r

if TYPE_CHECKING:
  from typing_extensions import Protocol
  class GetVersionFunc(Protocol):
    async def __call__(
      self,
      name: str, conf: Entry,
      *,
      cache: AsyncCache,
      keymanager: KeyManager,
    ) -> VersionResult:
      ...
else:
  GetVersionFunc = Any

class FunctionWorker(BaseWorker):
  func: GetVersionFunc
  cache: AsyncCache

  def initialize(self, func: GetVersionFunc) -> None:
    self.func = func
    self.cache = AsyncCache()

  async def run(self) -> None:
    futures = [
      self.run_one(name, entry)
      for name, entry in self.tasks
    ]
    for fu in asyncio.as_completed(futures):
      await fu

  async def run_one(
    self, name: str, entry: Entry,
  ) -> None:
    assert self.func is not None

    try:
      async with self.acquire_token():
        version = await self.func(
          name, entry,
          cache = self.cache,
          keymanager = self.keymanager,
        )
      await self.result_q.put(RawResult(name, version, entry))
    except Exception as e:
      await self.result_q.put(RawResult(name, e, entry))

class GetVersionError(Exception):
  def __init__(self, msg: str, **kwargs: Any) -> None:
    self.msg = msg
    self.kwargs = kwargs

@@ -270,17 +270,8 @@
 ble,
- TypeVar,
  Cor
@@ -278,23 +278,24 @@
 outine, 
-Generic
+Hashable
 ,%0A  TYPE
@@ -1648,42 +1648,8 @@
 ')%0A%0A
-T = TypeVar('T')%0AS = TypeVar('S')%0A
 %0Acla
@@ -1665,23 +1665,8 @@
 ache
-(Generic%5BT, S%5D)
 :%0A  
@@ -1681,33 +1681,21 @@
 ict%5B
-T, Union%5BS, asyncio.Task%5D
+Hashable, Any
 %5D%0A  
@@ -2124,17 +2124,24 @@
    key: 
-T
+Hashable
 ,%0A    fu
@@ -2154,17 +2154,24 @@
 llable%5B%5B
-T
+Hashable
 %5D, Corou
@@ -2179,21 +2179,21 @@
 ine%5B
-None, None, S
+Any, Any, Any
 %5D%5D,%0A
@@ -2203,9 +2203,11 @@
  -%3E 
-S
+Any
 :%0A  
@@ -2521,39 +2521,24 @@
 eturn cached
- # type: ignore
 %0A    else: #

update typing