import hashlib
import logging
import os
import stat
from datetime import date
from io import StringIO
from pathlib import Path
from typing import Tuple, Optional, List, Sequence

from clarabot_nano_common.common.error import log_escape_str
from clarabot_nano_common.common.handler.error import NanoHandlerError
from clarabot_nano_common.common.protocol.errors import (
    NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST,
    NANO_REQUEST_ERROR_HANDLER_DENIED,
    NANO_REQUEST_ERROR_HANDLER_RESET_REQUEST,
    NANO_REQUEST_ERROR_HANDLER_TARGET_EXISTS,
    NANO_REQUEST_ERROR_SHARE_LOCKED_OPERATION,
    NANO_REQUEST_ERROR_INVALID_NEW_NAME,
    NANO_REQUEST_ERROR_LOCKED_PATH,
)
from clarabot_nano_common.common.protocol.request import NanoRequestContext
from clarabot_nano_common.common.protocol.response import NanoBaseResponse
from clarabot_nano_common.server.base.config import ServerConfig
from clarabot_nano_common.utility.file import path_stat_differs
from clarabot_nano_common.utility.path.chroot import NewNameValidationError
from clarabot_nano_common.utility.path.error import attempt_share_locked_operation, SharedLockError
from clarabot_nano_common.utility.path.validate import ValidationError, validate_name
from clarabot_nano_common.utility.serialize import msgpack_dumps, msgpack_loads, bigger_int_from_js_msgpack

from clarabot_nano.common.filesystem.operation.replace import ReplaceOperation
from clarabot_nano.common.filesystem.operation.write import (
    WriteTruncateOperation,
    SeekWriteOperation,
    AbstractWriteOperation,
)
from clarabot_nano.common.filesystem.tx_queue import TokenLockError
from clarabot_nano.common.handler.filesystem.fstx import AbstractEnsureDirectoryNanoHandlerMixin, AUTO_PATH_FILE
from clarabot_nano.common.request.server_fstx import NanoFilesystemPutRequest
from clarabot_nano.server.fstx.handler.base import (
    AbstractFilesystemTransactorNanoHandler,
    FilesystemTransactorServerHandleHelper,
)
from clarabot_nano.utility.file import DRIVE_PUT_FILE

logger = logging.getLogger(__name__)


class AutoNameReplaceOperation(ReplaceOperation):
    """
    Usually the paths to be locked are consistent and the operation types change, but at _move_to_final() with
    auto-name the original destination is substituted with a safely created alternate that needs no explicit locking.
    It's directory is locked by the original paths so it is safe to work with.
    NOTE: If more such operation will arise, maybe the exact path-matching requirement at submission could be loosened.
    """

    OPERATION_KEY = "AN-REP"

    def __init__(self, account_id: int, src: str, dst: str, orig_dst: str, mtime_ns: Optional[int] = None):
        ReplaceOperation.__init__(self, account_id, src, dst, mtime_ns)
        self.orig_dst = orig_dst

    def paths(self) -> Sequence[str]:
        return (self.src, self.orig_dst)


class NanoFilesystemPutResponse(NanoBaseResponse):
    # only added to initialization response
    NORMALIZED_NAME = "norm_name"  # may be different from the input name due to normalization
    # added to initialization and finalization response if auto_name is enabled
    NORMALIZED_PATH = "norm_path"  # returns the upload path or the final path selected by the Nano respectively


class PutSystemHandler(AbstractFilesystemTransactorNanoHandler, AbstractEnsureDirectoryNanoHandlerMixin):
    """
    The optimal upload of a file is not straight forward to implement due to limitations of the clients and security
    requirements.
    The individual requests and data parts are safely encrypted and authenticated. Their integrity during transmission
    is thus guaranteed. The only point at which the file can get corrupted is when the client is reading from its
    storage or when it is written to the storage of the server. The server implementation should if possible check
    the file integrity without fully reading it back from the storage after the upload is finished. For this we assume
    all parts will be successfully written to disk or raise an error to tell the client the issue.
    The only additional integrity the server needs to check to ensure adequate safety is to ensure the order and
    number of parts the file consists of according to the client.
    The upload process will detect integrity errors if the source file's size of modification time changes, so the
    source file is somewhat checked to be consistent.
    """

    EXECUTION_STATE_CAN_CONFLICT = False  # cacheable and request_key are constant, no conflict is possible

    HANDLER_KEY = NanoFilesystemPutRequest.TYPE_KEY

    def __init__(
        self,
        helper: FilesystemTransactorServerHandleHelper,
        extracted_config: ServerConfig,
        context: dict,
        payload: dict,
    ):
        AbstractFilesystemTransactorNanoHandler.__init__(self, helper, extracted_config, context, payload)
        self.fstx_queue = helper.fstx_queue
        self.auto_name = self.payload.get(NanoFilesystemPutRequest.AUTO_NAME)
        if self.auto_name is None:
            pass  # Not using the automatic upload-path behaviour.
        elif isinstance(self.auto_name, str):
            self.auto_name = self.extracted_config.normalize_remote_path(self.auto_name)
            try:  # Pre-check this name to be a valid file-name.
                validate_name(self.auto_name, dos=False)
            except ValidationError as e:
                raise NanoHandlerError(NANO_REQUEST_ERROR_INVALID_NEW_NAME, "Invalid auto_name", self) from e
        else:
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid auto_name", self)
        self.offset = bigger_int_from_js_msgpack(self.payload.get(NanoFilesystemPutRequest.OFFSET))
        if not isinstance(self.offset, int) or self.offset < 0:
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid offset", self)
        self.path = self.path_s = None  # will be created with chroot in execute
        if not isinstance(self.payload.get(NanoFilesystemPutRequest.PATH), str):
            if self.auto_name and self.offset == 0 and self.payload.get(NanoFilesystemPutRequest.PATH) is None:
                # Select upload location automatically.
                # NOTE: Using the PATH as None here. This makes the auto-name start request invalid for prior Nanos that
                # do not support this functionality.
                pass
            else:
                raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Missing path", self)
        self.size = bigger_int_from_js_msgpack(self.payload.get(NanoFilesystemPutRequest.SIZE))
        if not isinstance(self.size, int) or self.size < -1:  # allow -1 for indeterminate
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid size", self)
        self.mtime = bigger_int_from_js_msgpack(self.payload.get(NanoFilesystemPutRequest.MTIME))
        if not isinstance(self.mtime, int) or self.mtime < -1:  # allow -1 for indeterminate
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid mtime", self)
        self.mtfrac = bigger_int_from_js_msgpack(self.payload.get(NanoFilesystemPutRequest.MTFRAC))
        if self.mtfrac is None:
            pass
        elif not isinstance(self.mtfrac, int) or self.mtfrac < 1:
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid mtfrac", self)
        self.data = self.payload.get(NanoFilesystemPutRequest.DATA)
        if not isinstance(self.data, bytes):
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid data", self)
        self.smac = self.payload.get(NanoFilesystemPutRequest.SMAC)
        if self.smac is not None and (not isinstance(self.smac, bytes) or len(self.smac) != 16):
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid data", self)
        if self.auto_name:  # NOTE: Overwrite is n/a with auto-name.
            self.overwrite = None
        else:
            self.overwrite = bigger_int_from_js_msgpack(self.payload.get(NanoFilesystemPutRequest.OVERWRITE))
            if self.overwrite is not None and not isinstance(self.overwrite, int):
                raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid overwrite", self)
        if not self.can_write_drive() and not self.can_upload_drive():  # anonymous is denied by mutation-id requirement
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_DENIED, "No permission", self)
        self.final_path = None  # only used for logging at the smac part

    async def _create_request_key(self, secure: bool = False) -> Tuple[bool, Optional[list], int]:
        return self._secure_request_key(False, None, 1) if secure else False, None, 1

    async def _execute_initial(self) -> dict:
        d = await AbstractFilesystemTransactorNanoHandler._execute_initial(self)
        try:
            drive_path = self.payload.get(NanoFilesystemPutRequest.PATH)
            if drive_path is None:  # For auto-name start.
                assert self.auto_name and self.offset == 0
                auto_parent = await self.ensure_get_drive_directory(f"Nano/{date.today().strftime('%Y-%m')}")
                # To avoid upload conflicts of the same user-upload fencing itself off, the generated target path
                # shall be relatively unique.
                async with self.helper.fstx_queue.wait_ready(paths=((auto_parent / self.auto_name).as_posix(),)):
                    # Queue-locked to avoid parent race condition operations like move or delete.
                    self.path = auto_parent / AUTO_PATH_FILE.reserve(auto_parent, self.auto_name, try_as_is=False)
                # NOTE: The path was created on disk to ensure it does not exist, but it should not exist before
                # the upload finishes so we remove it. Additionally, if it were to exist at finish because of any
                # other conflicts, a new random path might be chosen for auto-name upload instead anyways.
                self.path.unlink()  # NOTE: We can just leave .path_s as None after this.
                drive_path = self.path.relative_to(self.chroot.root).as_posix()
                d[NanoFilesystemPutResponse.NORMALIZED_PATH] = drive_path
            else:
                self.path, self.path_s = self.chroot.new_path(
                    drive_path, name_preprocess_fn=self.extracted_config.normalize_remote_path
                )
            upload_path = self.path.as_posix()
            # Tracking key is created to easily match up to a previously interrupted upload. This makes it near
            # automatic to continue or if needed restart an upload. The only bad scenario is that if the same
            # account wants to upload two different files to the same directory with the same name at the same
            # time. They will infinitely reset each other until one of them gives up and lets the other finish.
            tracking_key = (self.context[NanoRequestContext.ACCOUNT_ID], upload_path)
            tmp_row = await self.helper.temp_file_db.get_tmp_file(*tracking_key)
            if not tmp_row or self.offset == 0:  # starting or restarting
                if self.offset != 0:
                    raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_RESET_REQUEST, "Upload session not found", self)
                if not self.auto_name:
                    # If not using auto-name the target may already exist and prevent the upload. Checking it here
                    # at (re)start can prevent an empty temp-file from being created temporarily if the user does
                    # choose a different name.
                    self._validate_target_path()
                if not tmp_row:
                    async with self.helper.fstx_queue.wait_ready(paths=(upload_path,)):  # to avoid parent races
                        tmp_name = DRIVE_PUT_FILE.reserve(self.path)
                else:
                    _, tmp_name, _ = tmp_row
                tmp_id = await self.helper.temp_file_db.create_tmp_file(
                    *tracking_key,
                    tmp_name,
                    msgpack_dumps([self.size, self.mtime, self.mtfrac]),
                )
                d[NanoFilesystemPutResponse.NORMALIZED_NAME] = self.path.name
            else:
                tmp_id, tmp_name, tmp_meta = tmp_row
                if msgpack_loads(tmp_meta) != [self.size, self.mtime, self.mtfrac]:
                    raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_RESET_REQUEST, "Upload session mismatch", self)
            temp_path = (self.path.parent / tmp_name).as_posix()
            # Lock on the work paths.
            async with self.helper.fstx_queue.wait_ready(
                paths=AbstractWriteOperation(0, temp_path, upload_path, b"").paths()
            ) as submitter:
                # Final check before execution. This can fail if contesting operations have changed the filesystem.
                if path_stat_differs(
                    self.path,
                    self.path_s,
                    *self.chroot.new_path(drive_path, name_preprocess_fn=self.extracted_config.normalize_remote_path),
                ):
                    raise ValidationError("Depending paths are contested")
                if not self.auto_name:
                    # When auto-name is used, the target file is never overwritten. Skipping this check is safe as
                    # the _move_to_final() will call it and handle the exceptions to avoid the replacement.
                    # NOTE: The write operations below could also omit the lock on the final path as an optimization
                    # on this code-path.
                    self._validate_target_path()
                # Insert an incorrect mac-part for the current part before writing. This will prevent parallel
                # writes of the same file to seem complete, while one of the writes have failed and may have
                # broken a previous part. In other words: we invalidate the mac-store before each write and fix
                # it after the write is completed successfully by replacing the incorrect entry.
                await self.helper.temp_file_db.add_part_mac(tmp_id, self.offset, -1, b"")
                await submitter(
                    WriteTruncateOperation(
                        self.context[NanoRequestContext.ACCOUNT_ID], temp_path, upload_path, self.data
                    )  # starting or restarting
                    if not tmp_row or self.offset == 0
                    else SeekWriteOperation(
                        self.context[NanoRequestContext.ACCOUNT_ID],
                        temp_path,
                        upload_path,
                        self.data,
                        self.offset,
                    )
                )  # TODO: timeout
                await self.helper.temp_file_db.add_part_mac(
                    tmp_id,
                    self.offset,
                    len(self.data),
                    hashlib.sha256(self.data).digest()[:16],
                )
            if self.smac:
                # finalize upload by checking part consistency and MAC
                parts: List[int, int, bytes] = await self.helper.temp_file_db.get_all_part_mac(tmp_id)
                parts.sort(key=lambda part: part[0])
                smac = hashlib.sha256()
                prev_offset = 0
                for part_offset, part_size, part_mac in parts:
                    if prev_offset != part_offset:
                        raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_RESET_REQUEST, "Upload missing parts", self)
                    prev_offset += part_size
                    smac.update(part_mac)
                if smac.digest()[:16] != self.smac:
                    raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_RESET_REQUEST, "Upload failed to validate", self)
                try:
                    self.final_path = await attempt_share_locked_operation(
                        3, 3.0, self._move_to_final, drive_path, temp_path, upload_path
                    )
                except SharedLockError as e:
                    raise NanoHandlerError(
                        NANO_REQUEST_ERROR_SHARE_LOCKED_OPERATION, "Move to final path failed", self
                    ) from e
                await self.helper.temp_file_db.drop_tmp_file(tmp_id)
                if self.auto_name:
                    d[NanoFilesystemPutResponse.NORMALIZED_PATH] = self.final_path.relative_to(
                        self.chroot.root
                    ).as_posix()
        except NewNameValidationError as e:
            raise NanoHandlerError(NANO_REQUEST_ERROR_INVALID_NEW_NAME, "Invalid name", self) from e
        except ValidationError as e:
            raise NanoHandlerError(
                NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid path", self, maybe_conflict=True
            ) from e
        except OSError as e:
            raise NanoHandlerError(
                NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Unresolved path", self, maybe_conflict=True
            ) from e
        except TokenLockError as e:
            raise NanoHandlerError(NANO_REQUEST_ERROR_LOCKED_PATH, "Locked path", self) from e
        return d

    def _validate_target_path(self):
        # validate target path and permission to writing it
        try:
            s = self.path.lstat()
            if not stat.S_ISREG(s.st_mode):
                # can't be overwritten if it is not a file
                raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_TARGET_EXISTS, "Already exists", self)
        except OSError:  # the target path does not exist or inaccessible (can't discern)
            # NOTE: Permissions for CREATE are checked in __init__() so this is redundant currently.
            if not self.can_write_drive() and not self.can_upload_drive():  # CREATE
                raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_DENIED, "No permission", self)
        else:  # overwriting target path
            if not self.can_write_drive():  # MODIFY
                raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_DENIED, "No permission", self)
            # TODO: require timestamp of target file for more explicit protection?
            if not self.overwrite:
                raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_TARGET_EXISTS, "Already exists", self)

    async def _move_to_final(self, drive_path: str, temp_path: str, upload_path: str) -> Path:
        mtime_ns = None
        if self.mtfrac is not None and self.mtime is not None and self.mtime >= 0:
            mtfrac = str(self.mtfrac)[1:10]
            if len(mtfrac) < 9:
                mtfrac += "0" * (9 - len(mtfrac))
            mtime_ns = int(f"{self.mtime // 1000}{mtfrac}")
        # move tmp to final path (TOCTOU ofc)
        async with self.helper.fstx_queue.wait_ready(paths=(temp_path, upload_path)) as submitter:
            # Final check before execution. This can fail if contesting operations have changed
            # the filesystem.
            if path_stat_differs(
                self.path,
                self.path_s,
                *self.chroot.new_path(drive_path, name_preprocess_fn=self.extracted_config.normalize_remote_path),
            ):
                raise ValidationError("Depending paths are contested")
            try:  # Permission and race-condition checking.
                self._validate_target_path()
            except NanoHandlerError:
                if not self.auto_name:
                    raise  # User must handle conflict and select overwrite if they are allowed.
                # In case of auto-name, any conflict shall be avoided, so try and select any random path
                # in the target directory as final instead.
                final_path = self.path.parent / AUTO_PATH_FILE.reserve(
                    self.path.parent, self.auto_name, try_as_is=False
                )
                try:
                    op = AutoNameReplaceOperation(
                        self.context[NanoRequestContext.ACCOUNT_ID],
                        temp_path,
                        final_path.as_posix(),
                        orig_dst=upload_path,
                        mtime_ns=mtime_ns,
                    )
                    await submitter(op)  # TODO: timeout
                except Exception:
                    # We cautiously clean up the proposed final path in case of an error. Replace could be blocked by
                    # some security system temporarily, for which an implicit retry might succeed so leaving an extra
                    # empty file at the current final path is undesired.
                    try:
                        if os.lstat(temp_path).st_size != 0 and final_path.lstat().st_size == 0:
                            # Unlink the "unused" final path if it is still empty, while the temp file is not empty.
                            final_path.unlink()
                    except Exception:
                        logger.warning("Cleanup of proposed auto-name final path failed: %s", final_path, exc_info=True)
                    raise
            else:  # Move to upload-path is safe.
                final_path = self.path
                op = ReplaceOperation(
                    self.context[NanoRequestContext.ACCOUNT_ID], temp_path, upload_path, mtime_ns=mtime_ns
                )
                await submitter(op)  # TODO: timeout
            self._set_remote_file_security(final_path)
            return final_path

    async def _execute_part(self, part: int, s_hmac_start_part: Optional[int]) -> dict:
        raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Part request is invalid", self)

    def _log_message(self) -> str:
        if self.smac:
            if self.final_path is not None and self.path != self.final_path:
                return f"Write {log_escape_str(self.path)} as {log_escape_str(self.final_path)}"
            return f"Write {log_escape_str(self.path)}"
        return f"Write part {log_escape_str(self.path)} {self.offset}-{self.offset+len(self.data)}"

    def build_error_message(self, buf: StringIO) -> None:
        buf.write("Write ")
        auto_name = self.payload.get(NanoFilesystemPutRequest.AUTO_NAME)
        offset = bigger_int_from_js_msgpack(self.payload.get(NanoFilesystemPutRequest.OFFSET))
        # roughly mirror what is relevant to execution by the auto-name behaviour
        if auto_name:
            buf.write("auto-name ")
            buf.write(log_escape_str(auto_name))
            buf.write(" as ")
            if offset:
                buf.write(log_escape_str(self.payload.get(NanoFilesystemPutRequest.PATH)))
        else:
            buf.write(log_escape_str(self.payload.get(NanoFilesystemPutRequest.PATH)))
        if self.payload.get(NanoFilesystemPutRequest.SMAC):
            pass
        else:
            buf.write(" at ")
            buf.write(str(offset))
        buf.write(" in ")
        buf.write(log_escape_str(self.get_drive_root_for_error_message()))
