import asyncio
import base64
import os
import stat
import tempfile
from concurrent.futures.thread import ThreadPoolExecutor
from contextlib import asynccontextmanager
from datetime import datetime
from io import StringIO
from pathlib import Path
from typing import Sequence, Type, List, AsyncContextManager, Tuple

import aiohttp
import safe_radix32
from aiohttp import web, StreamReader
from clarabot_nano_common.common.config.constants import CommonConfigConstants
from clarabot_nano_common.common.handler.error import NanoHandlerError
from clarabot_nano_common.common.protocol.request import NanoRequestContext
from clarabot_nano_common.server.base.handler.base import AbstractNanoServerHandler
from clarabot_nano_common.server.base.handler.filesystem.mixin import AbstractNanoDriveHandlerEvalMixin
from clarabot_nano_common.utility.path.chroot import ChrootPath
from clarabot_nano_common.utility.path.validate import validate_name, ValidationError
from clarabot_nano_common.utility.timeout import IncreasingTimeout
from tarka.utility.file.reserve import ReserveFile

from clarabot_nano.common.zmq.inproc_client import AioInprocClient
from clarabot_nano.common.zmq.requester import AioClientRequester
from clarabot_nano.resource import get_resource_path
from clarabot_nano.server.base.process import ServerProcess
from clarabot_nano.server.document.config import DocumentServerConfig
from clarabot_nano.server.document.handler.base import DocumentServerHandleHelper
from clarabot_nano.server.document.handler.document.ws_close import DocumentWsCloseHandler
from clarabot_nano.server.document.handler.document.ws_open import DocumentWsOpenHandler
from clarabot_nano.server.document.handler.document.ws_prepare import DocumentWsPrepareHandler
from clarabot_nano.server.document.handler.document.ws_read import DocumentWsReadHandler
from clarabot_nano.server.document.handler.document.ws_write import DocumentWsWriteHandler


class WopiAccess(AbstractNanoDriveHandlerEvalMixin):
    def __init__(
        self,
        context: dict,
        extracted_config: DocumentServerConfig,
        file_path_id: str,
        can_write: bool,
        require_write: bool,
    ):
        self.context = context
        self.chroot = self.resolve_chroot(extracted_config)
        self.file_path_id = file_path_id
        resolved_path = Path(file_path_id)
        try:
            self.drive_path = resolved_path.relative_to(self.chroot.root).as_posix()
        except ValueError:
            raise web.HTTPNotFound()
        self.path, self.s = self.chroot.new_path(
            self.drive_path, name_preprocess_fn=extracted_config.normalize_remote_path
        )
        if self.s and (not stat.S_ISREG(self.s.st_mode) or self.path != resolved_path):
            raise web.HTTPNotFound()  # drive change made the file-id invalid or ambiguous
        can_write_now = self.can_write_drive()
        if can_write != can_write_now:
            raise web.HTTPNotFound()  # require session to be consistent with write level permission
        # enforce write permission level
        if require_write:
            if not can_write_now:
                raise web.HTTPNotFound()
        elif not self.s:  # the file must exist on the WOPI api when access is not for writing
            raise web.HTTPNotFound()

    def build_error_message(self, buf: StringIO) -> None:
        pass  # N/A

    @property
    def owner_id(self) -> int:
        return self.context[NanoRequestContext.OWNER_ACCOUNT_ID]

    @property
    def user_id(self) -> int:
        return self.context[NanoRequestContext.ACCOUNT_ID]

    @property
    def room_id(self) -> int:
        return self.context[NanoRequestContext.ROOM][NanoRequestContext.ROOM_ID]

    def reserve_relative_file(self, suggested_name: str) -> Tuple[str, Path, str]:
        new_name = ReserveFile().reserve(self.path.parent, suggested_name)
        new_path = self.path.parent / new_name
        new_drive_path = new_path.relative_to(self.chroot.root).as_posix()
        return new_name, new_path, new_drive_path


def _modification_timestamp(s: os.stat_result) -> str:
    return datetime.fromtimestamp(s.st_mtime).isoformat()


class DocumentServerProcess(ServerProcess):
    _extracted_config: DocumentServerConfig = None  # to fix type hinting

    LOCK_RETRY_DELAY = IncreasingTimeout(0.1, 1.0, 0.05)

    def __init__(self, pool: ThreadPoolExecutor, server_prefix: str):
        ServerProcess.__init__(self, pool, server_prefix, log_with_pid=True)

    def _get_extracted_config_cls(self):
        return DocumentServerConfig

    @asynccontextmanager
    async def _create_handle_helper(
        self, extracted_config, gateway_client, handler
    ) -> AsyncContextManager[DocumentServerHandleHelper]:
        requester = AioClientRequester(gateway_client, handler, setup_control_notifying_error_handler=False)
        async with aiohttp.ClientSession() as session:
            helper = DocumentServerHandleHelper(self._pool, requester, extracted_config, session)
            try:
                yield helper
            finally:
                await helper.document_manager.close_document_sessions()

    def _init_get_handler_classes(self) -> Sequence[Type[AbstractNanoServerHandler]]:
        return [
            DocumentWsPrepareHandler,
            DocumentWsOpenHandler,
            DocumentWsReadHandler,
            DocumentWsWriteHandler,
            DocumentWsCloseHandler,
        ]

    def _create_work_tasks(
        self, gateway_client: AioInprocClient, helper: DocumentServerHandleHelper
    ) -> List[asyncio.Future]:
        return [
            asyncio.create_task(self._wopi_api(helper)),
        ]

    async def _wopi_api(self, helper: DocumentServerHandleHelper):
        # setup web-app api to serve local WOPI client
        app = web.Application()
        app["helper"] = helper
        app.add_routes(
            [
                web.get("/wopi/files/{file_id}", self._wopi_handle_check_file_info),
                web.get("/wopi/files/{file_id}/contents", self._wopi_handle_get_file),
                web.post("/wopi/files/{file_id}", self._wopi_handle_file),
                web.post("/wopi/files/{file_id}/contents", self._wopi_handle_put_file),
                # Public templates access point.
                web.get("/template/{path:.+}", self._handle_template_get),
            ]
        )
        runner = web.AppRunner(app)
        await runner.setup()
        site = web.TCPSite(runner, "127.0.0.1", self._extracted_config.cool_wopi_port)
        await site.start()
        try:
            while True:
                await asyncio.sleep(60)
        finally:
            await runner.cleanup()

    async def _handle_template_get(self, request: web.Request):
        template_root = ChrootPath(
            get_resource_path("template", "document"),
            allow_symlinks=False,
            contain_symlinks=True,
            allow_path_cycles=False,
            restrict_sensitive=False,
        )
        return web.FileResponse(template_root.path(request.match_info["path"]))

    async def _wopi_handle_check_file_info(self, request: web.Request):
        access = await self._auth_wopi_access(request, require_write=False, load_avatar_name=True)
        anonymous = access.context[NanoRequestContext.ACCOUNT_ID] == 0
        read_only = not access.can_write_drive()
        info = {
            # ==== Required properties ====
            "BaseFileName": access.path.name,
            "OwnerId": safe_radix32.encode_fw(access.owner_id),
            "Size": access.s.st_size,
            "UserId": safe_radix32.encode_fw(access.user_id),
            # "Version": "",  # NOTE: Not implemented, because not utilized by wsd.
            # ==== WOPI host capabilities ====
            "SupportsUpdate": True,
            "SupportsGetLock": True,
            "SupportsLocks": True,
            # "SupportsUserInfo": False,
            # "SupportsRename": False,
            # "SupportsDeleteFile": False,
            # "SupportsExtendedLockLength": False,
            # "SupportsFolders": False,
            # ==== User metadata properties ====
            "IsAnonymousUser": anonymous,
            "UserFriendlyName": access.context.get(NanoRequestContext.ACCOUNT_AVATAR_NAME) or "",
            # "UserInfo": "",  # get operation of SupportsUserInfo
            # ==== User permission properties =====
            "ReadOnly": read_only,
            # "RestrictedWebViewOnly": False,
            "UserCanAttend": True,  # see update broadcast
            "UserCanNotWriteRelative": read_only,
            "UserCanPresent": not read_only,  # broadcast updates
            # "UserCanRename": False,
            "UserCanWrite": not read_only,
            # ==== File URL properties ====
            # "CloseUrl": "",  # client app navigates here when ir closes or crashes
            # "DownloadUrl": "",  # for the file
            # "FileEmbedCommandUrl": "",  # embedding specific solution
            # "FileSharingUrl": "",  # sharing
            # "FileUrl": "",  # the host can tell the client to get the file from a different host (performance)
            # "FileVersionUrl": "",  # version history
            # "HostEditUrl": "",  # editing url
            # "HostEmbeddedViewUrl": "",  # embedded viewing url
            # "HostViewUrl": "",  # viewing url
            # "SignoutUrl": "",
            # ==== Other properties ====
            # "AllowErrorReportPrompt": False,
            # "AllowExternalMarketplace": False,
            # "ClientThrottlingProtection": "",
            # "CloseButtonClosesWindow": False,
            # "DisablePrint": False,
            # "DisableTranslation": False,
            "FileExtension": access.path.suffix,
            # "FileNameMaxLength": 250,  # excluding the extension... crazy world eh?
            "LastModifiedTime": _modification_timestamp(access.s),
            # "RequestedCallThrottling": "",
            # "SHA256": "",  # base64 hash of file for view caching
            # "TemporarilyNotWritable": False,
            # ==== wsd optional properties ====
            # "PostMessageOrigin": "",  # ???
            "HidePrintOption": True,  # DownloadAsPDF is equivalent, and we do not implement custom printing yet... ?
            # "HideSaveOption": False,
            # "HideExportOption": False,
            # "DisableExport": False,
            # "DisableCopy": False,
            "EnableOwnerTermination": True,
            # "IsUserLocked": False,
            # "IsUserRestricted": False,
            # ==== extended properties ====
            # "TemplateSource": "<url>",
            # "UserExtraInfo": {"avatar": "<url>"},
        }
        avatar_key = access.context.get(NanoRequestContext.ACCOUNT_AVATAR_KEY)
        if avatar_key:
            info["UserExtraInfo"] = {
                "avatar": CommonConfigConstants.api_url(f"r/avatar/{info['UserId']}/{avatar_key}"),
            }
        # When file is empty, try to create it from a template. Plain text files (.txt) are special, as they are
        # already valid with zero size, so we can't safely template them.
        if access.s.st_size == 0 and os.path.splitext(info["BaseFileName"])[-1].lower() != ".txt":
            if read_only:
                raise web.HTTPNotFound()  # read-only session cannot create a new document from template
            helper: DocumentServerHandleHelper = request.app["helper"]
            template = helper.get_template_source(access.file_path_id, access.user_id, access.room_id)
            if not template:
                raise web.HTTPNotFound()  # no template to create with
            info["TemplateSource"] = f"http://127.0.0.1:{helper.extracted_config.cool_wopi_port}/template/{template}"
        return web.json_response(info)

    async def _wopi_handle_get_file(self, request: web.Request):
        max_size = 2147483647
        if max_size_h := request.headers.get("X-WOPI-MaxExpectedSize"):
            try:
                max_size = max(int(max_size_h), 0)
            except Exception:
                raise web.HTTPBadRequest()
        access = await self._auth_wopi_access(request, require_write=False)
        if access.s.st_size > max_size:
            raise web.HTTPPreconditionFailed()
        return web.FileResponse(access.path)

    async def _wopi_handle_put_file(self, request: web.Request):
        access = await self._auth_wopi_access(request)
        override = request.headers.get("X-WOPI-Override")
        if override != "PUT":
            raise web.HTTPBadRequest()
        if not access.s:  # NOTE: The GetLock check would fail on a non-existent file with 404 later anyway.
            raise web.HTTPNotFound()
        # NOTE: If the external conflict detection X-COOL-WOPI-Timestamp header is not present, the client is
        # forcing save (maybe by user request). The header will always be present otherwise, because there is
        # no path on which the CheckFileInfo returns successfully without a timestamp value.
        req_ts = request.headers.get("X-COOL-WOPI-Timestamp")
        if req_ts:
            file_ts = _modification_timestamp(access.s)
            if req_ts != file_ts:
                return web.json_response({"COOLStatusCode": 1010}, status=web.HTTPConflict.status_code)
        helper: DocumentServerHandleHelper = request.app["helper"]
        locking_token = await helper.internal_get_lock(access.user_id, access.room_id, access.drive_path)
        # NOTE: We require locking for put to ensure consistency, but this is not WOPI standard.
        # NOTE: The WSD has deficient conflict handling in case it has not locked or lost the lock.
        if locking_token is None:
            return web.Response(status=web.HTTPConflict.status_code, headers={"X-WOPI-Lock": ""})
        lock = request.headers.get("X-WOPI-Lock")
        # lock policy by WOPI standard
        if (lock and lock != locking_token) or (
            not lock and (locking_token is not None or (access.s is not None and access.s.st_size != 0))
        ):
            return web.Response(status=web.HTTPConflict.status_code, headers={"X-WOPI-Lock": locking_token or ""})
        assert locking_token is not None
        s = await self._write_file_body(access.path, request.content)  # write new file
        return web.json_response({"LastModifiedTime": _modification_timestamp(s)})

    async def _wopi_handle_file(self, request: web.Request):
        access = await self._auth_wopi_access(request)
        override = request.headers.get("X-WOPI-Override")
        if override == "GET_LOCK":
            return await self._wopi_handle_file_get_lock(request, access)
        if override in ("LOCK", "UNLOCK", "REFRESH_LOCK"):
            return await self._wopi_handle_file_lock(request, access, override)
        if override == "PUT_RELATIVE":
            return await self._wopi_handle_file_put_relative(request, access)
        raise web.HTTPNotImplemented()

    async def _wopi_handle_file_get_lock(self, request: web.Request, access: WopiAccess):
        if not access.s:
            raise web.HTTPNotFound()
        helper: DocumentServerHandleHelper = request.app["helper"]
        locking_token = await helper.internal_get_lock(access.user_id, access.room_id, access.drive_path)
        return web.Response(
            status=web.HTTPConflict.status_code if locking_token == "" else web.HTTPOk.status_code,
            headers={"X-WOPI-Lock": locking_token or ""},
        )

    async def _wopi_handle_file_lock(self, request: web.Request, access: WopiAccess, override: str):
        lock = request.headers.get("X-WOPI-Lock")
        if not lock:
            raise web.HTTPBadRequest()
        if not access.s:
            raise web.HTTPNotFound()
        old_lock = None
        if override == "LOCK":
            old_lock = request.headers.get("X-WOPI-OldLock")
            if old_lock:
                op = "replace"
            else:
                op = "create"
        elif override == "REFRESH_LOCK":
            op = "refresh"
        elif override == "UNLOCK":
            op = "delete"
        else:
            raise NotImplementedError()
        helper: DocumentServerHandleHelper = request.app["helper"]
        conflicting_token = await helper.internal_lock(
            access.user_id, access.room_id, access.drive_path, op, lock, old_lock
        )
        if conflicting_token is not None:
            return web.Response(status=web.HTTPConflict.status_code, headers={"X-WOPI-Lock": conflicting_token})
        return web.Response()

    async def _wopi_handle_file_put_relative(self, request: web.Request, access: WopiAccess):
        if (
            request.headers.get("X-WOPI-RelativeTarget")
            or request.headers.get("X-WOPI-OverwriteRelativeTarget")
            or request.headers.get("X-WOPI-FileConversion")
        ):
            raise web.HTTPNotImplemented()  # NOTE: Not implementing them, they are not utilized.
        try:  # SuggestedTarget is required.
            suggested_target = request.headers.get("X-WOPI-SuggestedTarget").encode("ascii").decode("utf-7")
        except Exception:
            raise web.HTTPBadRequest()
        # NOTE: It is actually irrelevant if the source file is locked or even exists.
        try:
            validate_name(suggested_target, dos=False)
        except ValidationError:
            suggested_target = Path(access.drive_path).name
        else:
            if suggested_target.startswith("."):
                suggested_target = Path(access.drive_path).with_suffix(suggested_target).name
        new_name, new_path, new_drive_path = access.reserve_relative_file(suggested_target)
        # NOTE: We need to lock the path internally to work on it with safety after atomic creation.
        helper: DocumentServerHandleHelper = request.app["helper"]
        locking_token = f"prel-{base64.urlsafe_b64encode(os.urandom(9)).decode()}"
        for i in range(8):  # Optimally we should wait for the lock to be acquired, but we can retry for now.
            conflicting_token = await helper.internal_lock(
                access.user_id, access.room_id, new_drive_path, "create", locking_token
            )
            if conflicting_token is None:
                break
            elif suggested_target == new_name:
                # Try and reserve a new random name if the suggested name is locked. This can work around a broken
                # or lost file not being saved by PutFile, but may succeed here.
                new_name, new_path, new_drive_path = access.reserve_relative_file(suggested_target)
                continue
            await asyncio.sleep(self.LOCK_RETRY_DELAY.get(i))
        else:
            # NOTE: Returning 409 in case of "SuggestedTarget" is not standard, so we raise internal error instead.
            raise web.HTTPInternalServerError(reason="Locking failure")
        try:
            s = await self._write_file_body(new_path, request.content)  # write new file
        finally:  # do not leave path locked in case of error
            await helper.internal_lock(access.user_id, access.room_id, new_drive_path, "delete", locking_token)
        # File written, lock released. We just need to respond accordingly.
        new_file_id, new_access_token = helper.create_wopi_url_args(
            new_path, access.user_id, access.room_id, access.can_write_drive()
        )
        new_url = (
            f"http://127.0.0.1:{helper.extracted_config.cool_wopi_port}"
            f"/wopi/files/{new_file_id}?access_token={new_access_token}"
        )
        return web.json_response(
            {
                "Name": new_name,
                "Url": new_url,
                "LastModifiedTime": _modification_timestamp(s),
            }
        )

    async def _auth_wopi_access(
        self, request: web.Request, require_write: bool = True, load_avatar_name: bool = False
    ) -> WopiAccess:
        helper: DocumentServerHandleHelper = request.app["helper"]
        file_path, account_id, room_id, can_write = helper.verify_wopi_url_args(
            request.match_info["file_id"], request.query.get("access_token", "")
        )
        context = await helper.get_internal_auth_context(account_id, room_id, load_avatar_name=load_avatar_name)
        if not context:  # access is denied
            raise web.HTTPNotFound()
        try:
            access = WopiAccess(context, helper.extracted_config, file_path, can_write, require_write)
        except NanoHandlerError:
            raise web.HTTPNotFound()  # either the requesting path is invalid or the drive is inaccessible
        return access

    async def _write_file_body(self, path: Path, body: StreamReader) -> os.stat_result:
        fd, temp_path = tempfile.mkstemp(dir=path.parent, prefix=path.name, text=False)
        try:
            with os.fdopen(fd, "wb") as f:  # FIXME: aiofiles.open ?
                while True:
                    chunk = await body.readany()
                    f.write(chunk)
                    if not chunk:
                        break
            os.replace(temp_path, path)
        finally:
            try:
                os.unlink(temp_path)  # do not leave the temp file for whatever reason.
            except OSError:
                pass
        return os.lstat(path)
