from __future__ import annotations

import asyncio
from concurrent.futures.thread import ThreadPoolExecutor
from functools import partial
from pathlib import Path
from typing import Optional, List, Tuple, Callable, Any

import aiohttp

from clarabot_nano.common.solr.doc import query_iter_docs, SolrRequestFailed
from clarabot_nano.common.zmq.requester import AioClientRequester
from clarabot_nano.indexer.utility import (
    create_docs_query_exact,
    create_accessor_path,
    convert_accessor_to_parent_accessor,
)
from clarabot_nano.server.base.config import ServerConfig
from clarabot_nano.server.base.handler.base import ServerHandleHelper
from clarabot_nano.server.general.cachedb import IndexCacheDatabase
from clarabot_nano.server.general.solr.reflect import SolrSchemaQueryReflector
from clarabot_nano.utility.serialize import msgpack_dumps, ulong2bytes


class IndexCachedHandlerMixin(object):
    HANDLER_KEY: str  # fix type hinting
    helper: GeneralServerHandleHelper  # fix type hinting
    extracted_config: ServerConfig  # fix type hinting

    def index_cached_value(self, paths: List[Tuple[Path, str]]) -> asyncio.Future[Tuple[Any, str, Optional[Any]]]:
        """
        Initiate the retrieval or the generation of a cached value for the resolved paths.
        The value shall be cacheable for the resolved posix path, independent of the accessor-path for example.
        The argument shall be the paths list structure from the chroot utility.
        The caller is responsible for handling SolrRequestFailed.
        """
        cache_key = self.index_cache_key(paths)
        return self.helper.get_cached_value_future(cache_key, self._index_cached_value, paths, cache_key)

    def index_cache_key(self, paths: List[Tuple[Path, str]]) -> bytes:
        return msgpack_dumps([self.HANDLER_KEY, paths[-1][0].as_posix()])

    async def _index_cached_value(
        self, paths: List[Tuple[Path, str]], cache_key: bytes
    ) -> Tuple[Any, str, Optional[Any]]:
        """
        Indexed value caching implementation.
        """
        accessor_path = create_accessor_path(
            paths[0][0].as_posix(),
            paths[-1][0].relative_to(paths[0][0]).as_posix() if len(paths) > 1 else "",
        )
        parent_validity = await self.get_parent_validity(accessor_path)
        cached = await self.helper.cache_db.get(cache_key)
        if cached and cached[1] == parent_validity:
            cached = self._load_cached(cached[0])
            if cached is not None:
                return cached, accessor_path, None
        # Value is either not cached, expired or obsolete.
        cached, source = await self._make_cached(paths, accessor_path)
        # Only store new value if it is still correct. Implicit retry may be implemented here.
        if await self.get_parent_validity(accessor_path) == parent_validity:
            await self.helper.cache_db.set(cache_key, msgpack_dumps(cached), parent_validity)
        return cached, accessor_path, source

    async def get_parent_validity(self, accessor_path: str) -> bytes:
        """
        Get the current validity value for the indexed parent.
        """
        validity = 0
        async for docs in query_iter_docs(
            self.helper.session,
            self.extracted_config.solr_query_url,
            self.extracted_config.solr_auth_header,
            create_docs_query_exact(*convert_accessor_to_parent_accessor(accessor_path)),
            "path asc",
            "accessors,mtime",
            50,
            exc_cls=SolrRequestFailed,
            filter_queries=("file:0",),
        ):
            for doc in docs:
                if accessor_path in doc["accessors"]:
                    if doc.get("mtime"):
                        validity = doc["mtime"]
                    break
        return ulong2bytes(validity)

    def _load_cached(self, cached_value: bytes) -> Any:
        """
        Load and validate the msgpack-ed value. This method is customizable to provide a way to make old values
        obsolete during an upgrade.
        """
        raise NotImplementedError()

    async def _make_cached(self, paths: List[Tuple[Path, str]], accessor_path: str) -> Tuple[Any, Optional[Any]]:
        """
        Build the value to be cached. Optionally other data (used for building the value) can be returned and will
        be transferred to the caller code.
        """
        raise NotImplementedError()


class GeneralServerHandleHelper(ServerHandleHelper, SolrSchemaQueryReflector):
    """
    Solr schema reflection, core-requester and indexed-value cache database.
    """

    def __init__(
        self,
        pool: ThreadPoolExecutor,
        session: aiohttp.ClientSession,
        cache_db: IndexCacheDatabase,
        core_requester: AioClientRequester,
    ):
        ServerHandleHelper.__init__(self, pool)
        SolrSchemaQueryReflector.__init__(self)
        self.session = session
        self.cache_db = cache_db
        self.core_requester = core_requester
        self._cached_value_tasks = dict()

    async def close(self):
        ts = list(self._cached_value_tasks.values())
        if ts:
            for t in ts:
                t.cancel()
            await asyncio.wait(ts, return_when=asyncio.ALL_COMPLETED)

    def _cached_value_task_done(self, key: bytes, f):
        self._cached_value_tasks.pop(key, None)

    def get_cached_value_future(self, key: bytes, coro_fn: Callable, *args) -> asyncio.Future:
        if not (t := self._cached_value_tasks.get(key)):
            self._cached_value_tasks[key] = (t := asyncio.create_task(coro_fn(*args)))
            t.add_done_callback(partial(self._cached_value_task_done, key))
        return asyncio.shield(t)


# class AbstractGeneralNanoHandler(AbstractNanoHandler):
#     # fix type hinting
#     helper: GeneralServerHandleHelper
