import asyncio
import logging
import math
import os
import stat
from functools import partial
from typing import Tuple, Optional, List, Union

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,
)
from clarabot_nano.common.protocol.request import NanoRequestContext
from clarabot_nano.common.protocol.response import NanoBaseResponse
from clarabot_nano.common.request.server_general import NanoFilesystemGetRequest
from clarabot_nano.common.solr.doc import query_iter_docs, SolrRequestFailed, decompose_docs_by_parent_accessor
from clarabot_nano.indexer.utility import create_accessor_path, 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
from clarabot_nano.server.general.handler.filesystem.utility import (
    read_file_range,
    list_directory,
    fingerprint_file_range,
    ls_entry_size,
)
from clarabot_nano.utility.path.validate import ValidationError

logger = logging.getLogger(__name__)


class NanoFilesystemGetResponse(NanoBaseResponse):
    FILE_SIZE = "fs"  # initial response if target is file
    LISTING = "ls"  # if target is directory
    FILE_CHUNK = "fc"  # part response if target is file


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

    HANDLER_KEY = NanoFilesystemGetRequest.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(NanoFilesystemGetRequest.PATH), str):
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Path is missing", self.context)
        self.no_thumbnails = bool(self.payload.get(NanoFilesystemGetRequest.NO_THUMBNAILS))
        self.path = None
        self.stat_or_ls: Union[os.stat_result, List[List[dict]]] = None

    async def _create_request_key(self) -> Tuple[bool, Optional[list], int]:
        try:
            paths, s = self.chroot.paths_stat(self.payload[NanoFilesystemGetRequest.PATH])
        except (OSError, ValidationError):
            raise MaybeConflictNanoHandlerError(
                NANO_REQUEST_ERROR_HANDLER_TARGET_NOT_FOUND, "Path is invalid", self.context
            )
        if self.path is None:  # pre-execution
            if stat.S_ISDIR(s.st_mode):
                list_args = [list_directory, partial(self.chroot.iter, self.payload[NanoFilesystemGetRequest.PATH])]
                # NOTE: If Solr has not started yet, or is experiencing issues, the listing will work without it.
                if not self.no_thumbnails and self.extracted_config.solr_query_url:
                    # Load thumbnail buckets for the directory and use it for generating the list.
                    parent_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 "",
                    )
                    image_buckets = {}
                    try:
                        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_bucket",
                            250,
                            exc_cls=SolrRequestFailed,
                            filter_queries=("file:1", create_docs_query_children(parent_accessor_path)),
                        ):
                            for name, doc in decompose_docs_by_parent_accessor(parent_accessor_path, ds):
                                if img_bucket := doc.get("img_bucket"):  # this is actually implied by the query
                                    image_buckets[name] = img_bucket
                    except SolrRequestFailed:
                        logger.warning("Solr query failed for fs-get: %s", paths[-1][0], exc_info=True)
                    else:
                        list_args.append(False)
                        list_args.append(image_buckets)
                # Directory listing must be performed in advance to calculate part count. Multipart directory listing
                # is seldom needed so it is okay if it has some overhead in that case.
                try:
                    ls = await asyncio.get_event_loop().run_in_executor(self.helper.pool, *list_args)
                except (OSError, ValidationError) as e:  # pragma: no cover; cannot simulate this
                    # Any error for actually reading the directory here is bad-request, except when it is a conflict.
                    # Permission error, deleted, renamed, moved... they cannot be recovered. If there is some
                    # temporary error type maybe that should be handled specifically.
                    raise MaybeConflictNanoHandlerError(
                        NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Directory listing failed", self.context
                    ) from e
                # Slicing a directory listing is not trivial. No OS supports pagination and arbitrary starting point
                # for iteration. The only property which we can build on is that the the items are in some arbitrary,
                # but deterministic order. It is mostly up to the type of filesystem the storage uses, which means
                # the properties of the ordering may be different even by folder.
                # In case there is some filesystem which does not return the file listing in a deterministic order,
                # manual sorting would be required, which is costly so we opt to not do it until we are sure we need
                # to.
                # Handle overlap merging by the entries' names.
                self.stat_or_ls = self._approximate_part_slicing(ls, ls_entry_size, 16)
            else:
                self.stat_or_ls = s
            self.path = paths[-1][0]
        path_parent, _, path_name = self.path.as_posix().rpartition("/")
        if isinstance(self.stat_or_ls, os.stat_result):
            parts = math.ceil(s.st_size / self.context[NanoRequestContext.PART_SIZE])
            # file retrieval
            return (
                bool(parts),  # empty file is a simple transfer-response without more parts
                [self.HANDLER_KEY, s.st_ino, s.st_size, s.st_mtime_ns, path_name, path_parent],
                1 + parts,
            )
        # directory listing
        # request_key is used to ensure relative consistency (no caching)
        return (
            False,
            [self.HANDLER_KEY, s.st_ino, path_name, path_parent],
            len(self.stat_or_ls),
        )

    async def _execute_initial(self) -> dict:
        d = await AbstractFileSystemHandler._execute_initial(self)
        if isinstance(self.stat_or_ls, os.stat_result):
            d[NanoFilesystemGetResponse.FILE_SIZE] = self.stat_or_ls.st_size
        else:
            d[NanoFilesystemGetResponse.LISTING] = self.stat_or_ls[0]
        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 isinstance(self.stat_or_ls, os.stat_result):
            try:
                part_start = (part - 1) * self.context[NanoRequestContext.PART_SIZE]
                file_chunk = await asyncio.get_event_loop().run_in_executor(
                    self.helper.pool, read_file_range, self.path, part_start, self.context[NanoRequestContext.PART_SIZE]
                )
                d[NanoFilesystemGetResponse.FILE_CHUNK] = file_chunk
                if s_hmac_start_part is not None:
                    range_fingerprint, bytes_count = await asyncio.get_event_loop().run_in_executor(
                        self.helper.pool,
                        fingerprint_file_range,
                        self.path,
                        s_hmac_start_part * self.context[NanoRequestContext.PART_SIZE],
                        self.context[NanoRequestContext.PART_SIZE],
                        part - s_hmac_start_part,
                        salt := os.urandom(16),
                    )
                    self._add_s_hmac(d, bytes_count, range_fingerprint, salt)
            except (OSError, ValidationError) as e:  # pragma: no cover; cannot simulate this
                # Any error for actually reading the file here is bad-request, except when it is a conflict.
                # Permission error, deleted, renamed, moved... they cannot be recovered. If there is some
                # temporary error type maybe that should be handled specifically.
                raise MaybeConflictNanoHandlerError(
                    NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "File chunk reading failed", self.context
                ) from e
        else:
            # NOTE: we cannot s_hmac the listing because it is volatile... no real benefit for p2p encryption anyways
            d[NanoFilesystemGetResponse.LISTING] = self.stat_or_ls[part]
        return d

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