import hmac
import os.path
import re
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Optional, Tuple
from urllib.parse import unquote, urlparse, parse_qsl, quote, urlencode

import aiohttp
import pybase64
from aiohttp import web
from clarabot_crypto.base.cipher.aes import AESCipherSIVDeterministic
from clarabot_nano_common.common.config.constants import LocalSecrets
from clarabot_nano_common.common.handler.error import NanoHandlerError
from clarabot_nano_common.common.protocol.errors import (
    NANO_REQUEST_ERROR_DENIED,
    NANO_REQUEST_ERROR_HANDLER_TARGET_NOT_FOUND,
    NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST,
)
from clarabot_nano_common.common.protocol.request import NanoRequestContext
from clarabot_nano_common.common.protocol.server import NanoServerResponse
from clarabot_nano_common.server.base.handler.base import ServerHandleHelper, AbstractNanoServerHandler
from clarabot_nano_common.server.base.handler.filesystem.base import AbstractFileSystemHandler
from clarabot_nano_common.server.base.handler.filesystem.mixin import AbstractNanoDriveHandlerEvalMixin
from clarabot_nano_common.utility.serialize import msgpack_dumps, msgpack_loads
from lxml import etree
from ttldict2 import TTLDict

from clarabot_nano.common.gateway.message import InternalGatewayRequest, InternalGatewayResponse
from clarabot_nano.common.request.server_document import DocumentWsSessionRequestBase
from clarabot_nano.common.zmq.constants import INTERNAL_REQUEST_MESSAGE
from clarabot_nano.common.zmq.requester import AioClientRequester
from clarabot_nano.server.document.config import DocumentServerConfig
from clarabot_nano.server.document.proxy import DocumentWsManager


class DocumentServerHandleHelper(ServerHandleHelper):
    WSD_VERSION_HASH_RE = re.compile(r"^https?://[^/]+/browser/([^/]+)/cool\.html\?$")
    WSD_WS_URL_RE = re.compile(r"^wss?://[^/?#]+/cool/([^/?#]+)/ws(?:\?(.+)$)?")
    WOPI_URL_PATH_RE = re.compile(r"^/wopi/files/(.+)$")

    def __init__(
        self,
        pool: ThreadPoolExecutor,
        requester: AioClientRequester,
        extracted_config: DocumentServerConfig,
        session: aiohttp.ClientSession,
    ):
        ServerHandleHelper.__init__(self, pool)
        self.requester = requester
        self.extracted_config = extracted_config
        self.session = session
        self.document_manager = DocumentWsManager(session, extracted_config)
        self._requesting_template_sources = TTLDict(30.0)
        self._empty_templates = {
            ".odt": "empty/writer.ott",
            ".odp": "empty/impress.otp",
            ".ods": "empty/calc.ots",
        }

    async def get_wsd_version_hash(self) -> str:
        async with self.session.get(
            f"http://127.0.0.1:{self.extracted_config.cool_server_port}/hosting/discovery"
        ) as resp:
            resp.raise_for_status()
            discovery_text = await resp.text()
        discovery = etree.fromstring(discovery_text)
        for action in discovery.xpath("/wopi-discovery/net-zone/app[@name='writer']/action"):
            url_src = action.get("urlsrc")
            if url_src:
                m = self.WSD_VERSION_HASH_RE.match(url_src)
                if m:
                    return m.group(1)
        raise ValueError("urlsrc not found in discovery xml")

    def create_wopi_url_args(
        self, resolved_path: Path, account_id: int, room_id: int, can_write: bool
    ) -> Tuple[str, str]:
        """
        The WOPI client (wsd) needs two arguments for work, a "file-id" and an "access-token".

        === FILE-ID ===
        The "file-id" is a value that the clients (user, wopi, wsd) use to refer to a specific resource.

        This value does not need to be secured against manipulation. Operations that use the file-id reference,
        do so by validating strict, live access control of the nano for the requesting account by the "access-token".

        ENCRYPT THE VALUE
        The value is designed to appear in URLs as a path or query component. This is unfortunate, because it makes
        leaking file names and paths easy or unavoidable. To resolve this privacy issue we will obfuscate or encrypt
        the value so only the relevant Nano in question can understand it.

        REQUIRED PROPERTIES FOR DEFINED BEHAVIOUR
        The WSD service opens a new editor websocket and connects to the document broker by the file-id if it already
        exists. This means that this value governs how collaboration can be accessed. Since we do support nested drive
        layouts, the file-id must be unique for each file, but must be the same in each drive. This effectively means
        that the full local path of the file must be used with a deterministic transformation. For consistent behaviour
        the path needs to be fully resolved and case corrected as well.

        NESTED-DRIVE USER ACCOUNT DISCLOSURE
        The collaborative editor could reveal different users from different drives to each other. Before, even if
        they have had access to the same file they would not be able to find out about each other. To keep separation
        we could make the file-id value drive or room bound, but that would fragment the collaborative editor groups,
        they would prevent each other from working on the same file by the locking mechanism. Such a state would be
        complicated to the users to understand, and it can't actually be resolved, the owner may only kick out one of
        the groups from editing so the other can lock the file to work on it.
        As a mitigation the account-ids could be obfuscated in the editor, but that would be inconvenient as well for
        the user to understand what is happening.
        TODO: We need to evaluate how security or privacy is affected. This privacy issue is only relevant if the owner
              actually wants to secretly keep different groups separate and hidden from each other. Which is not fully
              possible, because they will see changes by other people, just not know what their identity is. Later
              there might even be file-tx history logs available that would also reveal such information.

        === ACCESS-TOKEN ===
        The "access-token" is a value that the user client acquires and passes to the wsd client to get access to a
        storage traditionally.

        Each access-token has to be bound to a specific file-id, using any of the two values with a different
        combination shall not be valid for security reasons. Token expiration would also be required normally,
        but in our case such is not not necessary, because all token use already presumes an authenticated client.
        Since the users can't change their access-token arbitrarily (managed and overseen by the document-server),
        the "token trading" issue is naturally prevented.

        In our case the access-token is superfluous on the user-client side, because the integrated storage and
        authorization system of Nano talks directly to the wsd. TODO: invalid constant or omitted from user client?

        As the access token will always be created on the Nano from the current authorized request-context and
        verified by the same Nano instance. As long as the wsd process behaves as expected (does not manipulate
        authorization tokens either by design or bug) the token would not even need to be signed. There is a websocket
        command to change the access-token of the session, this command must be verified by the Nano to ensure
        consistent session authorization. The real-word use-case would be to refresh the access-token when it expires
        most likely, but it should be N/A for us. Regardless, the access-token passed to the wsd will be HMAC signed
        and verified on the WOPI api.

        The access-token shall contain the account-id and relevant room-id to be checked for consistency.
        """
        file_id = AESCipherSIVDeterministic.encrypt(
            resolved_path.as_posix().encode("utf-8"), self.extracted_config.secrets[LocalSecrets.WOPI_FILE_ID_KEY]
        )
        access_token = msgpack_dumps([account_id, room_id, can_write])
        signature = hmac.digest(
            self.extracted_config.secrets[LocalSecrets.WOPI_ACCESS_TOKEN_KEY],
            b".2.".join((file_id, access_token)),
            "sha1",  # secure because of HMAC construct
        )
        return pybase64.urlsafe_b64encode(file_id).decode("ascii").rstrip("="), ".".join(
            (
                pybase64.urlsafe_b64encode(access_token).decode("ascii").rstrip("="),
                pybase64.urlsafe_b64encode(signature).decode("ascii").rstrip("="),
            )
        )

    def verify_wopi_url_args(self, file_id_arg: str, access_token_arg: str) -> Tuple[str, int, int, bool]:
        """
        This just verifies the signature and decrypts the file_id, the using code shall check the authorization of
        the room and the existence of the path within the drive as necessary!
        """
        try:
            file_id = pybase64.urlsafe_b64decode(file_id_arg.encode("ascii") + b"==")
            access_token_b64, _, signature_b64 = access_token_arg.partition(".")
            access_token = pybase64.urlsafe_b64decode(access_token_b64.encode("ascii") + b"==")
            signature = pybase64.urlsafe_b64decode(signature_b64.encode("ascii") + b"==")
        except Exception:  # formatting or padding errors
            raise web.HTTPUnauthorized()
        signature_ref = hmac.digest(
            self.extracted_config.secrets[LocalSecrets.WOPI_ACCESS_TOKEN_KEY],
            b".2.".join((file_id, access_token)),
            "sha1",  # secure because of HMAC construct
        )
        if not hmac.compare_digest(signature_ref, signature):
            raise web.HTTPUnauthorized()
        try:
            file_path = AESCipherSIVDeterministic.decrypt(
                file_id, self.extracted_config.secrets[LocalSecrets.WOPI_FILE_ID_KEY]
            ).decode("utf-8")
            account_id, room_id, can_write = msgpack_loads(access_token)
        except Exception:
            raise web.HTTPUnauthorized()
        return file_path, account_id, room_id, can_write

    def parse_wsd_ws_url(self, wsd_ws_url: str) -> Tuple[str, int, int, bool, str]:
        """
        Use to verify the file_id and access_token in a websocket initiator URL.
        The translation from the public constant is optional, an already translated local service wopi url is
        accepted as-is. This is relevant when "saveas" is used and the WOPI api itself returns a local URL.
        """
        m = self.WSD_WS_URL_RE.match(wsd_ws_url)
        if not m:
            raise Exception(f"Invalid WSD ws url: {wsd_ws_url}")
        wopi_url = m.group(1)
        ws_url_query = m.group(2)  # WOPISrc={wopi_url}&...
        if wopi_url.startswith("0x"):  # un-hexify
            wopi_url = bytes.fromhex(wopi_url[2:]).decode("utf-8")
        unqouted_wopi_url = unquote(wopi_url)  # https://<host>/wopi/files/{file_id}}?access_token={access_token}&...
        if not unqouted_wopi_url.startswith(
            self.extracted_config.WOPI_HOST_TRANSLATE_CONST
        ) and not unqouted_wopi_url.startswith(self.extracted_config.wopi_host):
            raise Exception("Invalid WOPI host")
        parsed_wopi_url = urlparse(unqouted_wopi_url)
        m = self.WOPI_URL_PATH_RE.match(parsed_wopi_url.path)
        if not m:
            raise Exception("Invalid WOPI url path in WSD ws url")
        file_id_arg = m.group(1)
        for key, value in parse_qsl(parsed_wopi_url.query):
            if key == "access_token":
                access_token_arg = value
                break
        else:
            raise Exception("Invalid WOPI url query in WSD ws url")
        if ws_url_query:
            parsed_ws_url_query = dict(parse_qsl(ws_url_query))
            if "WOPISrc" not in parsed_ws_url_query:
                raise Exception("Invalid WOPISrc in WSD ws url")
            if not unqouted_wopi_url.startswith(parsed_ws_url_query["WOPISrc"] + "?"):
                raise Exception("Inconsistent WOPISrc and WOPI url path in WSD ws url")
        else:
            parsed_ws_url_query = {
                "WOPISrc": f"{parsed_wopi_url.scheme}://{parsed_wopi_url.netloc}/{parsed_wopi_url.path}"
            }
        # translate constant host to local service if necessary
        if not unqouted_wopi_url.startswith(self.extracted_config.wopi_host):
            unqouted_wopi_url = self.extracted_config.translate_to_wopi_host(unqouted_wopi_url)
            parsed_ws_url_query["WOPISrc"] = self.extracted_config.translate_to_wopi_host(
                parsed_ws_url_query["WOPISrc"]
            )
        file_path, account_id, room_id, can_write = self.verify_wopi_url_args(file_id_arg, access_token_arg)
        ws_url = (
            f"ws://127.0.0.1:{self.extracted_config.cool_server_port}"
            f"/cool/{quote(unqouted_wopi_url, safe='')}/ws?{urlencode(parsed_ws_url_query)}"
        )
        if ws_url.endswith("&compat=%2Fws"):  # NOTE: more correct, but may be unnecessary
            ws_url = ws_url[:-5] + "/ws"
        return file_path, account_id, room_id, can_write, ws_url

    async def get_internal_auth_context(self, account_id: int, room_id: int, load_avatar_name: bool = False) -> dict:
        result = await self.get_internal_response(account_id, room_id, None, load_avatar_name)
        context = result.get(InternalGatewayResponse.CONTEXT)
        if e := result.get(NanoServerResponse.ERROR) or not isinstance(context, dict):
            if e == NANO_REQUEST_ERROR_DENIED:
                return {}
            raise Exception(f"Gateway internal-auth response error {e}")
        return context

    async def internal_get_lock(self, account_id: int, room_id: int, path: str) -> Optional[str]:
        result = await self.get_internal_response(account_id, room_id, {"type": "FS-L-G", "path": path})
        response = result.get(InternalGatewayResponse.RESPONSE)
        if e := result.get(NanoServerResponse.ERROR) or not isinstance(response, dict):
            if e == NANO_REQUEST_ERROR_HANDLER_TARGET_NOT_FOUND:
                raise web.HTTPNotFound()
            raise Exception(f"Gateway internal get-lock error {e}")
        token = response.get("t")
        if token is not None and not isinstance(token, str):
            raise Exception("Gateway internal get-lock invalid")
        return token

    async def internal_lock(
        self, account_id: int, room_id: int, path: str, op: str, token: str, old=None
    ) -> Optional[str]:
        req = {"type": "FS-L-X", "path": path, "op": op, "token": token}
        if op == "replace":
            req["old"] = old
        result = await self.get_internal_response(account_id, room_id, req)
        response = result.get(InternalGatewayResponse.RESPONSE)
        if e := result.get(NanoServerResponse.ERROR) or not isinstance(response, dict):
            if e == NANO_REQUEST_ERROR_HANDLER_TARGET_NOT_FOUND:
                raise web.HTTPNotFound()
            raise Exception(f"Gateway internal lock error {e}")
        conflict = response.get("c")
        if conflict is not None and not isinstance(conflict, str):
            raise Exception("Gateway internal lock invalid")
        return conflict

    async def get_internal_response(
        self, account_id: int, room_id: int, request: Optional[dict], load_avatar_name: bool = False
    ) -> dict:
        """
        Acquire the current auth of the account_id for the room_id.
        """
        gw_req = {InternalGatewayRequest.ACCOUNT_ID: account_id, InternalGatewayRequest.ROOM_ID: room_id}
        if request is not None:
            gw_req[InternalGatewayRequest.REQUEST] = request
        if load_avatar_name:
            gw_req[InternalGatewayRequest.LOAD_AVATAR] = True
        response_frames = await self.requester.request([msgpack_dumps(gw_req)], message_type=INTERNAL_REQUEST_MESSAGE)
        if len(response_frames) != 1:
            raise Exception("Gateway internal-auth response framing is invalid")
        try:
            response = msgpack_loads(response_frames[0])
        except Exception:
            raise Exception("Gateway internal-auth response format is invalid")
        if not isinstance(response, dict):
            raise Exception("Gateway internal-auth response type is invalid")
        return response

    def validate_template_source_path(self, file_path: str, template: str) -> Optional[str]:
        if template == "empty":
            return self._empty_templates.get(os.path.splitext(file_path)[1].lower())

    def store_template_source(self, file_path: str, account_id: int, room_id: int, template: str) -> None:
        self._requesting_template_sources[(file_path, account_id, room_id)] = template

    def get_template_source(self, file_path: str, account_id: int, room_id: int) -> Optional[str]:
        template_source = self._requesting_template_sources.get((file_path, account_id, room_id))
        if template_source:
            return template_source
        return self.validate_template_source_path(file_path, "empty")


class AbstractDocumentNanoServerHandler(AbstractNanoServerHandler):
    helper: DocumentServerHandleHelper  # fix type hinting
    extracted_config: DocumentServerConfig  # fix type hinting


class AbstractDocumentFileSystemHandler(AbstractFileSystemHandler):
    helper: DocumentServerHandleHelper  # fix type hinting
    extracted_config: DocumentServerConfig  # fix type hinting


class AbstractDocumentSessionNanoServerHandler(AbstractDocumentNanoServerHandler, AbstractNanoDriveHandlerEvalMixin):
    ENSURE_WRITE_PERMISSION_CONSISTENCY = True

    def __init__(
        self, helper: DocumentServerHandleHelper, extracted_config: DocumentServerConfig, context: dict, payload: dict
    ):
        AbstractDocumentNanoServerHandler.__init__(self, helper, extracted_config, context, payload)
        self.session = self.payload.get(DocumentWsSessionRequestBase.SESSION)
        if not isinstance(self.session, str):
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Missing session", self)
        if not isinstance(room := self.context.get(NanoRequestContext.ROOM), dict):
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Missing room", self)
        self.document_session = self.helper.document_manager.get_document_session(
            self.session,
            self.context[NanoRequestContext.ACCOUNT_ID],
            room[NanoRequestContext.ROOM_ID],
            self.can_write_drive() if self.ENSURE_WRITE_PERMISSION_CONSISTENCY else None,
        )
        if not self.document_session:
            raise NanoHandlerError(NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST, "Invalid session", self)
        # Document-WS-Session security notes:
        #   - user prepares file-id + access-token for specific file in a drive, these values are secured by HMAC
        #   - when a session is opened by file-id + access-token, they are checked for consistency for being bound
        #     to a valid account, room and file in the drive of the room and write permission
        #   - the following read and write handlers ensure account, room and write consistency, but do not check for
        #     file existence in the drive, the WOPI client may handle if the file is lost or manipulated as appropriate
        #       - write permission grant and revocation will disrupt the document viewing editing for security
        #   - the close handler only checks account and room consistency to allow closure at write permission change
        #   - any access on the WOPI API (WSD side) are checked for account, room, write permission consistency,
        #     drive containment and file existence depending on if the request is read-only or not
        #   - on the WSD side, a new session's access-token is ensured to be the one that we've checked for open
        #       - since no other file-id may be valid for that token, the client/user cannot ask the WSD to load
        #         a different file than authorized
        #       - the only command that could change the access-token in a WSD session is blocked in the proxy
        #         communication handler ("resetaccesstoken")
        #   Thus any access to a session for any side is appropriately authorized and verified for consistency.
