import asyncio
import hashlib
import os
import stat
from pathlib import Path
from typing import Tuple, Optional, List, Any

import pybase64

from clarabot_nano.common.handler.base import NanoHandlerError, MaybeConflictNanoHandlerError
from clarabot_nano.common.protocol.errors import (
    NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST,
    NANO_REQUEST_ERROR_HANDLER_TARGET_NOT_FOUND,
    NANO_REQUEST_ERROR_HANDLER_TEMPORARY_FAILURE,
)
from clarabot_nano.common.protocol.request import NanoRequestContext
from clarabot_nano.common.protocol.response import NanoBaseResponse
from clarabot_nano.common.request.server_general import NanoFilesystemThumbnailsRequest
from clarabot_nano.common.solr.doc import query_iter_docs, SolrRequestFailed, decompose_docs_by_parent_accessor
from clarabot_nano.indexer.utility import create_docs_query_children
from clarabot_nano.server.base.config import ServerConfig
from clarabot_nano.server.base.handler.base import ServerHandleHelper
from clarabot_nano.server.base.handler.filesystem.base import AbstractFileSystemHandler
from clarabot_nano.server.general.handler.base import GeneralServerHandleHelper, IndexCachedHandlerMixin
from clarabot_nano.server.general.handler.filesystem.utility import fingerprint_parts_data
from clarabot_nano.utility.algorithm import multi_islice
from clarabot_nano.utility.path.validate import ValidationError
from clarabot_nano.utility.serialize import msgpack_dumps, msgpack_loads


class NanoFilesystemThumbnailsResponse(NanoBaseResponse):
    THUMBNAILS = "ts"


class ThumbnailsFileSystemHandler(AbstractFileSystemHandler, IndexCachedHandlerMixin):
    REJECTS_ANONYMOUS_IDENTITY = False  # Further authorization is by room membership and group flags.

    HANDLER_KEY = NanoFilesystemThumbnailsRequest.TYPE_KEY

    helper: GeneralServerHandleHelper  # fix type hinting

    def __init__(self, helper: ServerHandleHelper, extracted_config: ServerConfig, context: dict, payload: dict):
        AbstractFileSystemHandler.__init__(self, helper, extracted_config, context, payload)
        if not isinstance(self.payload.get(NanoFilesystemThumbnailsRequest.PATH), str):
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Path is missing", self.context)
        self.path = None
        self.part_slices: List[int] = None
        self.data_hash: bytes = None
        self.approx_size: int = None
        self.accessor_path: str = None
        self.thumbnails: List[Tuple[str, str, int]] = None

    async def _create_request_key(self) -> Tuple[bool, Optional[list], int]:
        try:
            paths, s = self.chroot.paths_stat(self.payload[NanoFilesystemThumbnailsRequest.PATH])
        except (OSError, ValidationError):
            raise MaybeConflictNanoHandlerError(
                NANO_REQUEST_ERROR_HANDLER_TARGET_NOT_FOUND, "Path is invalid", self.context
            )
        if not stat.S_ISDIR(s.st_mode):
            raise MaybeConflictNanoHandlerError(
                NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Path is not a directory", self.context
            )
        if self.path is None:  # pre-execution
            try:
                cached_value, accessor_path, thumbnails = await self.index_cached_value(paths)
            except SolrRequestFailed:
                raise NanoHandlerError(
                    NANO_REQUEST_ERROR_HANDLER_TEMPORARY_FAILURE,
                    "Preparing thumbnails failed by Solr query",
                    self.context,
                )
            # cached_value: Tuple[int, int, bytes, int, List[int]]
            self.accessor_path = accessor_path
            self.data_hash = cached_value[2]
            self.approx_size = cached_value[3]
            self.part_slices = cached_value[4]
            self.path = paths[-1][0]
        path_parent, _, path_name = self.path.as_posix().rpartition("/")
        return (
            cached := bool(self.approx_size),  # empty list is a simple transfer-response without more parts
            [self.HANDLER_KEY, s.st_ino, path_name, path_parent, self.approx_size, self.data_hash],
            cached + 1 + len(self.part_slices),
        )

    async def _execute_initial(self) -> dict:
        d = await AbstractFileSystemHandler._execute_initial(self)
        if not self.approx_size:
            d[NanoFilesystemThumbnailsResponse.THUMBNAILS] = None
        return d

    async def _execute_part(self, part: int, s_hmac_start_part: Optional[int]) -> dict:
        d = await AbstractFileSystemHandler._execute_part(self, part, s_hmac_start_part)
        if self.thumbnails is None:
            try:
                self.thumbnails = await self._get_thumbnails(self.accessor_path)
            except SolrRequestFailed:
                raise NanoHandlerError(
                    NANO_REQUEST_ERROR_HANDLER_TEMPORARY_FAILURE,
                    "Gathering thumbnails failed by Solr query",
                    self.context,
                )
        # NOTE: The returned mtime will match up with the st_mtime that is returned by LS, so cache punch-through
        # can be initiated by the client.
        parts_data = [
            msgpack_dumps(
                [(name, pybase64.b64decode(img_data), mtime / 1000000000) for name, img_data, mtime in thumbnails_slice]
            )
            for thumbnails_slice in multi_islice(
                self.thumbnails, self.part_slices, part - 1 if s_hmac_start_part is None else s_hmac_start_part, part
            )
        ]
        d[NanoFilesystemThumbnailsResponse.THUMBNAILS] = parts_data[-1]
        if s_hmac_start_part is not None:
            range_fingerprint, bytes_count = await asyncio.get_event_loop().run_in_executor(
                self.helper.pool,
                fingerprint_parts_data,
                parts_data,
                salt := os.urandom(16),
            )
            self._add_s_hmac(d, bytes_count, range_fingerprint, salt)
        return d

    def _log_message(self) -> str:
        if part := self.context.get(NanoRequestContext.PART):
            return f"Get {self.path} thumbnails part {part}"
        return f"Get {self.path} thumbnails"

    # index-cached value impl
    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.
        """
        try:
            value = msgpack_loads(cached_value)
        except Exception:
            return None
        if (
            not isinstance(value, (list, tuple))
            or len(value) != 5
            or value[0] != 1
            or value[1] != self.context[NanoRequestContext.PART_SIZE]
        ):
            return None
        return value

    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.
        """
        thumbnails = await self._get_thumbnails(accessor_path)
        part_slices = []
        all_size = 0
        last_slice_size = 0
        thumbnails_hash = hashlib.sha256()
        part_size = self.context[NanoRequestContext.PART_SIZE]
        for i, (name, img_data, mtime) in enumerate(thumbnails):
            if last_slice_size > part_size:
                part_slices.append(i)
                all_size += last_slice_size
                last_slice_size = 0
            # NOTE: Including the whole image data in the hash should not be necessary. The hash should change
            # anyways because of the inclusion of the mtime if the thumbnail is updated in the index.
            thumbnails_hash.update(f"{(img_size := len(img_data))}:{mtime}:{name}".encode("utf-8"))
            last_slice_size += 13 + len(name) + int(img_size * 0.75)
        return (1, part_size, thumbnails_hash.digest(), all_size + last_slice_size, part_slices), thumbnails

    async def _get_thumbnails(self, accessor_path: str) -> List[Tuple[str, str, int]]:
        thumbnails = []
        # NOTE: The img_bucket field is not used here yet. It only serves as a lightweight indication to
        # clients when listing a directory and caching the images.
        async for ds in query_iter_docs(
            self.helper.session,
            self.extracted_config.solr_query_url,
            self.extracted_config.solr_auth_header,
            "img_bucket:[1 TO *]",
            "path asc",
            "accessors,img_data,mtime",
            250,
            exc_cls=SolrRequestFailed,
            filter_queries=("file:1", create_docs_query_children(accessor_path)),
        ):
            for name, doc in decompose_docs_by_parent_accessor(accessor_path, ds):
                if img_data := doc.get("img_data"):  # this is actually implied by the query
                    thumbnails.append((name, img_data, mtime if (mtime := doc.get("mtime")) else 0))
        return thumbnails
