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

import pybase64
from clarabot_nano_common.common.error import log_escape_str, maybe_rad32
from clarabot_nano_common.common.handler.error import NanoHandlerError
from clarabot_nano_common.common.protocol.adhoc import AdhocTokenFields
from clarabot_nano_common.common.protocol.errors import (
    NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST,
    NANO_REQUEST_ERROR_HANDLER_TEMPORARY_FAILURE,
)
from clarabot_nano_common.common.protocol.request import NanoRequestContext
from clarabot_nano_common.common.protocol.response import NanoBaseResponse
from clarabot_nano_common.common.request.server_general import NanoFilesystemThumbnailsRequest
from clarabot_nano_common.common.solr.doc import (
    query_iter_docs,
    SolrRequestFailed,
    decompose_docs_by_parent_accessor,
    query_docs_range,
)
from clarabot_nano_common.indexer.utility import (
    create_docs_query_children,
    create_docs_query_exact,
    convert_accessor_to_parent_accessor,
)
from clarabot_nano_common.server.base.handler.base import ServerHandleHelper
from clarabot_nano_common.server.general.config import GeneralServerConfig
from clarabot_nano_common.server.general.handler.base import (
    IndexCachedHandlerMixin,
    IndexCachedHandlerValueArgs,
    AbstractGeneralFileSystemHandler,
)
from clarabot_nano_common.server.general.handler.filesystem.utility import fingerprint_parts_data
from clarabot_nano_common.utility.serialize import msgpack_dumps, msgpack_loads, ulong2bytes, data_length_of_base64


class NanoFilesystemThumbnailsResponse(NanoBaseResponse):
    THUMBNAILS = "ts"


class ThumbnailsIndexCachedHandlerValueArgs(IndexCachedHandlerValueArgs):
    __slots__ = ("bucket",)

    def __init__(self, handler_key: str, paths: List[Tuple[Path, str]], bucket: int):
        self.bucket = bucket
        IndexCachedHandlerValueArgs.__init__(self, handler_key, paths)

    def _create_cache_key(self, handler_key: str) -> Sequence[Any]:
        return [handler_key, self.bucket, self.paths[-1][0].as_posix()]


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

    HANDLER_KEY = NanoFilesystemThumbnailsRequest.TYPE_KEY

    def __init__(self, helper: ServerHandleHelper, extracted_config: GeneralServerConfig, context: dict, payload: dict):
        AbstractGeneralFileSystemHandler.__init__(self, helper, extracted_config, context, payload)
        self.adhoc_account_id, self.adhoc_path = self._init_adhoc_account_path()
        self.path = None
        if not isinstance(self.payload.get(NanoFilesystemThumbnailsRequest.PATH), str):
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Missing path", self)
        self.bucket = self.payload.get(NanoFilesystemThumbnailsRequest.BUCKET, 1)
        if not isinstance(self.bucket, int) or self.bucket < 1:
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid bucket", self)
        self.part_slices: List[int] = None
        self.data_hash: bytes = None
        self.approx_size: int = None
        self.cache_args: ThumbnailsIndexCachedHandlerValueArgs = None

    async def _create_request_key(self, secure: bool = False) -> Tuple[bool, Optional[list], int]:
        paths, s = self.resolve_chroot_paths_stat(
            self.chroot, self.payload[NanoFilesystemThumbnailsRequest.PATH], sub_root=self.adhoc_path
        )
        if not stat.S_ISDIR(s.st_mode):
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Not a directory", self, maybe_conflict=True)
        if self.path is None:  # pre-execution
            self.cache_args = ThumbnailsIndexCachedHandlerValueArgs(self.HANDLER_KEY, paths, self.bucket)
            try:
                # cached_value: Tuple[int, int, bytes, int, List[int]]
                _, _, self.data_hash, self.approx_size, self.part_slices = await self.index_cached_value(
                    self.cache_args
                )
            except SolrRequestFailed:
                raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_TEMPORARY_FAILURE, "Prepare query failed", self)
            self.path = paths[-1][0]
        path_parent, _, path_name = paths[-1][0].as_posix().rpartition("/")
        req = (
            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 + len(self.part_slices),
        )
        return self._secure_request_key(*req) if secure else req

    async def _execute_initial(self) -> dict:
        d = await AbstractGeneralFileSystemHandler._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 AbstractGeneralFileSystemHandler._execute_part(self, part, s_hmac_start_part)
        cache_args = self.cache_args
        parts_data = []
        try:
            for p in range(part - 1 if s_hmac_start_part is None else s_hmac_start_part, part):
                start = 0 if p == 0 else self.part_slices[p - 1]
                rows = self.part_slices[p] - start
                parts_data.append(
                    await self._get_thumbnails_part(cache_args.accessor_path, cache_args.bucket, start, rows)
                )
        except SolrRequestFailed:
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_TEMPORARY_FAILURE, "Gather query failed", self)
        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:
        acc = f" as adhoc by {maybe_rad32(self.adhoc_account_id)}" if self.adhoc_path else ""
        if part := self.context.get(NanoRequestContext.PART):
            return f"Thumbnails {log_escape_str(self.path)} bucket {self.bucket}{acc} part {part}"
        return f"Thumbnails {log_escape_str(self.path)} bucket {self.bucket}{acc}"

    def build_error_message(self, buf: StringIO) -> None:
        buf.write("Thumbnails ")
        buf.write(log_escape_str(self.payload.get(NanoFilesystemThumbnailsRequest.PATH)))
        buf.write(" bucket ")
        buf.write(log_escape_str(self.payload.get(NanoFilesystemThumbnailsRequest.BUCKET, 1)))
        if adhoc := self.context.get(NanoRequestContext.ADHOC_DATA):
            buf.write(" as adhoc ")
            buf.write(log_escape_str(adhoc.get("path")))
            buf.write(" by ")
            buf.write(maybe_rad32(adhoc.get(AdhocTokenFields.DATA_ACCOUNT_ID)))
        if part := self.context.get(NanoRequestContext.PART):
            buf.write(" part ")
            buf.write(str(part))
        buf.write(" in ")
        buf.write(log_escape_str(self.get_drive_root_for_error_message()))

    # index-cached value impl
    async def _get_parent_validity(self, cache_args: ThumbnailsIndexCachedHandlerValueArgs) -> bytes:
        validity = 0
        bucket_attr = f"imgb_{cache_args.bucket}"
        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(cache_args.accessor_path)),
            "path asc",
            f"accessors,{bucket_attr}",
            50,
            exc_cls=SolrRequestFailed,
            filter_queries=("file:0",),
        ):
            for doc in docs:
                if cache_args.accessor_path in doc["accessors"]:
                    if doc.get(bucket_attr):
                        validity = doc[bucket_attr]
                    break
        return ulong2bytes(validity)

    def _load_cached(self, cached_value: bytes) -> Any:
        try:
            value = msgpack_loads(cached_value)
        except Exception:
            return None
        if (
            not isinstance(value, (list, tuple))
            or len(value) != 5
            or value[0] != 2
            or value[1] != self.context[NanoRequestContext.PART_SIZE]
        ):
            return None
        return value

    async def _make_cached(self, cache_args: ThumbnailsIndexCachedHandlerValueArgs) -> Any:
        thumbnails = await self._get_thumbnails(cache_args.accessor_path, cache_args.bucket)
        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 := data_length_of_base64(img_data))}:{mtime}:{name}".encode("utf-8"))
            last_slice_size += 13 + len(name) + img_size
        if not part_slices or part_slices[-1] != len(thumbnails):
            part_slices.append(len(thumbnails))
        return 2, part_size, thumbnails_hash.digest(), all_size + last_slice_size, part_slices

    async def _get_thumbnails(self, accessor_path: str, bucket: int) -> List[Tuple[str, str, int]]:
        thumbnails = []
        async for ds in query_iter_docs(
            self.helper.session,
            self.extracted_config.solr_query_url,
            self.extracted_config.solr_auth_header,
            f"img_bucket:{bucket}",
            "path asc",  # deterministic order is necessary
            "accessors,img_data,mtime",
            250,  # ~1MB of raw image data
            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

    async def _get_thumbnails_part(self, accessor_path: str, bucket: int, start: int, rows: int) -> bytes:
        """
        Used for serving part requests by the cached range.
        """
        thumbnails = []
        for name, doc in decompose_docs_by_parent_accessor(
            accessor_path,
            await query_docs_range(
                self.helper.session,
                self.extracted_config.solr_query_url,
                self.extracted_config.solr_auth_header,
                f"img_bucket:{bucket}",
                "path asc",  # deterministic order is necessary
                "accessors,img_data,mtime",
                start,
                rows,
                exc_cls=SolrRequestFailed,
                filter_queries=("file:1", create_docs_query_children(accessor_path)),
            ),
        ):
            if img_data := doc.get("img_data"):  # this is actually implied by the query
                thumbnails.append(
                    (
                        name,
                        pybase64.b64decode(img_data),
                        (mtime if (mtime := doc.get("mtime")) else 0) / 1000000000,
                    )
                )
        # 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.
        return msgpack_dumps(thumbnails)
