from __future__ import annotations

import asyncio
import logging
import os
import shutil
import stat
import time
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager, nullcontext
from pathlib import Path
from typing import Optional, Callable, Tuple, Dict

import safe_radix32
from clarabot_crypto.context.nano import NANO_CRYPTO_CONTEXT, NanoConfigDumpArgs
from clarabot_crypto.context.nano_adhoc import NANO_ADHOC_CRYPTO_CONTEXT, NanoAdhocConfigDumpArgs
from clarabot_crypto.context.nano_own import NANO_OWN_CRYPTO_CONTEXT, NanoOwnConfigDumpArgs
from clarabot_crypto.keyring.account_base import AbstractSelfAccountKeyring, AbstractPeerAccountKeyring
from clarabot_crypto.keyring.adhoc import AdhocKeyring
from clarabot_crypto.keyring.room_base import AbstractRoomKeyring
from clarabot_nano_common.common.handler.filesystem.get import merge_fs_get_listing_responses
from clarabot_nano_common.common.protocol.errors import (
    NANO_REQUEST_ERROR_HANDLER_RESET_REQUEST,
    NANO_REQUEST_ERROR_PART_RANGE,
    NANO_REQUEST_ERROR_HANDLER_TEMPORARY_FAILURE,
    NANO_REQUEST_ERROR_HANDLER_TARGET_NOT_FOUND,
    NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST,
    NANO_REQUEST_ERROR,
    NANO_REQUEST_ERROR_UNKNOWN_TYPE,
    NANO_REQUEST_ERROR_TOO_MANY_REQUESTS,
)
from clarabot_nano_common.common.protocol.nano_request import NanoResponseError, NanoResponseStreamProcessor
from clarabot_nano_common.common.protocol.request import NanoRequestContext
from clarabot_nano_common.utility.path.chroot import ChrootPath
from clarabot_nano_common.utility.path.error import attempt_share_locked_operation
from clarabot_nano_common.utility.path.validate import ValidationError, validate_name, sanitize_name
from clarabot_nano_common.utility.serialize import msgpack_dumps

from clarabot_nano.common.filesystem.tx_queue import FilesystemTxQueue
from clarabot_nano.common.handler.filesystem.fstx import AUTO_PATH_FILE
from clarabot_nano.common.handler.filesystem.get import write_buffered_fs_get_file
from clarabot_nano.inter.tmpfiledb import InterTemporaryFileDatabase
from clarabot_nano.utility.file import DRIVE_PUT_FILE

logger = logging.getLogger(__name__)


class AbstractInterNanoRequester:
    def __init__(
        self,
        self_kr: AbstractSelfAccountKeyring,
        remote_kr: AbstractPeerAccountKeyring,
        remote_room_id: int,
        # not necessary when remote's owner is self as well
        room_kr: Optional[AbstractRoomKeyring] = None,
        adhoc_token: Optional[str] = None,
        adhoc_kr: Optional[AdhocKeyring] = None,
    ):
        self._loop = asyncio.get_running_loop()
        self.self_kr = self_kr
        self.remote_kr = remote_kr
        self.remote_room_id = safe_radix32.encode_fw(remote_room_id)
        self.room_kr = room_kr
        self.adhoc_token = adhoc_token
        self.adhoc_kr = adhoc_kr
        if room_kr and room_kr.id.decode("utf-8") != self.remote_room_id:
            raise Exception("Inconsistent remote room id in room keyring")
        if adhoc_kr and adhoc_kr.id.decode("utf-8") != self.remote_room_id:
            raise Exception("Inconsistent remote room id in adhoc keyring")
        if (self.adhoc_token is None) is not (self.adhoc_kr is None):
            raise Exception("Both adhoc token and secret must be provided or not")

    def encrypt_request(self, payload: dict) -> bytes:
        payload["room_id"] = self.remote_room_id
        request_payload = msgpack_dumps(payload)
        if self.self_kr.id == self.remote_kr.id:
            req = NANO_OWN_CRYPTO_CONTEXT.dump(NanoOwnConfigDumpArgs(request_payload, self.self_kr, time.time()))
        else:
            if self.adhoc_kr:
                req = NANO_ADHOC_CRYPTO_CONTEXT.dump(
                    NanoAdhocConfigDumpArgs(request_payload, self.self_kr, self.remote_kr, self.adhoc_kr, time.time())
                )
            elif self.room_kr:
                req = NANO_CRYPTO_CONTEXT.dump(
                    NanoConfigDumpArgs(request_payload, self.self_kr, self.remote_kr, self.room_kr, time.time())
                )
            else:
                raise Exception("Room-like keyring missing for inter-nano request")
        return req

    async def perform(self, pool: ThreadPoolExecutor, stream_processor: NanoResponseStreamProcessor):
        """
        It is the responsibility of the implementation to utilize threadpool for the invocation of
        stream_processor.feed() and also keep it serial in execution.
        """
        raise NotImplementedError()


class AbstractDrivePathLocking:
    async def call_ex_lock(self, account_id: int, room_id: int, path: str, token: str, ensure_dir: bool = False):
        raise NotImplementedError()

    async def call_ex_refresh_lock(self, account_id: int, room_id: int, path: str, token: str):
        raise NotImplementedError()

    async def call_ex_unlock(self, account_id: int, room_id: int, path: str, token: str):
        raise NotImplementedError()

    async def resolve_chroot_by_room(self, room_id: int) -> Optional[ChrootPath]:
        raise NotImplementedError()


class AbstractInterNanoExecutor:
    """
    Base class to implement nano-2-nano communication based processes.
    """

    # NOTE: The main process is serial in logic, so all thread-invocations made by asyncio-await clauses are
    # also only performed serially.
    POOL_WORKERS = 2
    NAME_PREFIX: str

    REQUEST_RETRY_DELAYS = (0.5, 1.0, 1.5, 5.0, 15.0, 30.0, 60.0)  # NOTE: length is number of retries

    # Names that are temporary or are a technical detail of some feature should be skipped.
    IGNORE_TEMPORARY_NAME_RE = DRIVE_PUT_FILE.tag_re

    def __init__(
        self,
        inr: AbstractInterNanoRequester,
        dpl: AbstractDrivePathLocking,
        owner_account_id: int,
        part_size: int,
        drive_room_id: int,
        chroot: ChrootPath,
        temp_db: InterTemporaryFileDatabase,
    ):
        self._loop = asyncio.get_running_loop()
        # Locking implementation for drive operation coordination.
        self.dpl = dpl
        self._dpl_data: Dict[Tuple[str, str], float] = dict()
        self._dpl_lock = asyncio.Lock()
        # This bunch is for addressing and performing requests to a remote Nano.
        self.inr = inr
        self.context = {
            NanoRequestContext.OWNER_ACCOUNT_ID: owner_account_id,
            NanoRequestContext.ACCOUNT_ID: owner_account_id,  # Operations are executed by this Nano, thus its owner.
            NanoRequestContext.PART_SIZE: part_size,
        }
        self.drive_room_id = drive_room_id
        # This is for state persistence and support for resuming large file downloads.
        self.temp_db = temp_db
        # This is mostly ineffective, but good to have.
        self.chroot = chroot
        # State tracking to adapt API behaviour to remote Nano.
        self._use_old_fs_get_listing = False

    async def run(self):
        with ThreadPoolExecutor(thread_name_prefix=self.NAME_PREFIX, max_workers=self.POOL_WORKERS) as pool:
            tasks = [
                asyncio.create_task(self._keep_lock_fresh()),
                asyncio.create_task(self._run_custom(pool)),
            ]
            try:
                done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
                for t in done:
                    return t.result()  # raise if error occurred
            finally:
                for t in tasks:
                    if not t.done():
                        t.cancel()
                await asyncio.gather(*tasks, return_exceptions=True)

    async def _run_custom(self, pool: ThreadPoolExecutor):
        """
        Use ._fstx_lock(...) to lock specific path then work under it.
        """
        raise NotImplementedError()

    @asynccontextmanager
    async def _fstx_lock(self, drive_path: str, ensure_dir: bool = False):
        while True:
            lock_token = f"{self.NAME_PREFIX}LOCK-{int(time.time() * 1000000)}-{os.urandom(8).hex()}"
            async with self._dpl_lock:
                if (drive_path, lock_token) not in self._dpl_data:
                    break
        await self.dpl.call_ex_lock(
            self.context[NanoRequestContext.ACCOUNT_ID], self.drive_room_id, drive_path, lock_token, ensure_dir
        )
        async with self._dpl_lock:
            self._dpl_data[(drive_path, lock_token)] = time.perf_counter()
        try:
            yield
        finally:
            async with self._dpl_lock:
                del self._dpl_data[(drive_path, lock_token)]
            await self.dpl.call_ex_unlock(
                self.context[NanoRequestContext.ACCOUNT_ID], self.drive_room_id, drive_path, lock_token
            )

    async def _keep_lock_fresh(self):
        while True:
            await asyncio.sleep(FilesystemTxQueue.DEFAULT_TOKEN_LOCK_TTL / 8)
            async with self._dpl_lock:
                for lock_path__lock_token, lock_ts in self._dpl_data.items():
                    if (time.perf_counter() - lock_ts) > FilesystemTxQueue.DEFAULT_TOKEN_LOCK_TTL / 2:
                        await self.dpl.call_ex_refresh_lock(
                            self.context[NanoRequestContext.ACCOUNT_ID], self.drive_room_id, *lock_path__lock_token
                        )
                        self._dpl_data[lock_path__lock_token] = time.perf_counter()

    async def _sync_remote_item_recursively(
        self,
        pool: ThreadPoolExecutor,
        tmp_dir: Optional[Path],
        old_folder: Optional[Path],  # previous local mirror to optimize file copies with if any
        new_folder: Path,
        remote_path: str,
        raise_on_peek_fail: bool = False,  # not recursive behaviour
        granular_locking: bool = True,
    ):
        peek = await self._peek_remote_item(pool, remote_path, raise_on_fail=raise_on_peek_fail)
        if isinstance(peek, dict) and isinstance(peek.get("name"), str):
            await self._sync_remote_recursively_impl(
                pool, tmp_dir, old_folder, new_folder, remote_path, peek, granular_locking
            )

    async def _sync_remote_directory_recursively(
        self,
        pool: ThreadPoolExecutor,
        tmp_dir: Optional[Path],
        old_folder: Optional[Path],  # previous local mirror to optimize file copies with if any
        new_folder: Path,
        remote_path: str,
        raise_on_listing_fail: bool = False,  # not recursive behaviour
        granular_locking: bool = True,
    ):
        remote_listing = await self._list_remote_dir(pool, remote_path, raise_on_fail=raise_on_listing_fail)
        for remote_item in remote_listing:
            if not isinstance(remote_item, dict) or remote_item.get("unresolved"):
                continue
            remote_name = remote_item.get("name")
            if not isinstance(remote_name, str):
                continue
            remote_item_path = "/".join([remote_path, remote_name] if remote_path else [remote_name])
            await self._sync_remote_recursively_impl(
                pool, tmp_dir, old_folder, new_folder, remote_item_path, remote_item, granular_locking
            )

    async def _sync_remote_recursively_impl(
        self,
        pool: ThreadPoolExecutor,
        tmp_dir: Optional[Path],
        old_folder: Optional[Path],  # previous local mirror to optimize file copies with if any
        new_folder: Path,
        remote_item_path: str,
        remote_item: dict,
        granular_locking: bool = True,
    ):
        # Validate name of remote item before use.
        remote_name = remote_item.get("name")
        if not isinstance(remote_name, str):
            return
        name = self._sanitize_check_name(remote_name)
        if not name:
            return
        # NOTE: since the name is validated before, we'll use it with safety from now on.
        mtime = remote_item.get("mtime")
        if not isinstance(mtime, float):
            return
        size = remote_item.get("size")
        if size is not None and not isinstance(size, int):
            return
        new_path = new_folder / name
        # Check item already in new folder if any.
        async with (
            self._fstx_lock(new_path.relative_to(self.chroot.root).as_posix()) if granular_locking else nullcontext()
        ):
            try:
                new_stat = new_path.lstat()
            except OSError:
                new_stat = None
            else:  # Conflict check & handle item if it does not match the remote.
                if stat.S_ISDIR(new_stat.st_mode):
                    if size is not None:
                        await self._loop.run_in_executor(pool, self._fix_conflicting_item, new_path)
                        new_stat = None
                elif size is None or size != new_stat.st_size or mtime != new_stat.st_mtime:  # file
                    await self._loop.run_in_executor(pool, self._fix_conflicting_item, new_path)
                    new_stat = None
            old_path = (old_folder / name) if old_folder else None
            if new_stat is None:  # Create, copy or download item into new.
                if size is None:
                    new_path.mkdir()
                else:  # file
                    await self._sync_remote_file(pool, tmp_dir, old_path, new_path, remote_item_path, size, mtime)
        if size is None:  # recurse into directory
            await self._sync_remote_directory_recursively(
                pool, tmp_dir, old_path, new_path, remote_item_path, granular_locking=granular_locking
            )

    def _sanitize_check_name(self, name: str) -> Optional[str]:
        if not isinstance(name, str):
            return
        try:
            validate_name(name)
        except ValidationError:
            try:
                name = sanitize_name(name)
            except ValidationError:
                logger.warning("Can't sanitize name of remote entity %s", name, exc_info=True)
                return
        if self.IGNORE_TEMPORARY_NAME_RE.match(name):
            return  # do not copy over temporary files
        return name

    async def _sync_remote_file(
        self,
        pool: ThreadPoolExecutor,
        tmp_dir: Optional[Path],
        old_path: Optional[Path],  # previous local mirror to optimize file copy with if any
        new_path: Path,
        remote_item_path: str,
        size: int,
        mtime: float,
    ):
        if old_path:
            try:
                old_stat = old_path.lstat()
            except OSError:
                old_stat = None
            else:
                if not stat.S_ISREG(old_stat.st_mode) or size != old_stat.st_size or mtime != old_stat.st_mtime:
                    old_stat = None
        else:
            old_stat = None
        if old_stat is not None:  # reflink copy file from old to new
            if not (await self._copy_file_old_to_new(pool, old_path, new_path)):
                old_stat = None
        if old_stat is None:  # download if not copied/moved from old
            await self._save_remote_file(pool, tmp_dir, remote_item_path, new_path, mtime, size)

    def _fix_conflicting_item(self, path: Path):
        alt_path = path.parent / AUTO_PATH_FILE.reserve(path, try_as_is=False)
        if path.is_dir():
            alt_path.unlink()  # can't replace file with folder
        os.replace(path, alt_path)

    def _prepare_tmp_directory(self, path: Path, force: bool = False):
        if path.exists() and not path.is_dir():
            self._fix_conflicting_item(path)
        path.mkdir(exist_ok=True)
        old_limit = time.time() - 60 * 60
        for item in path.iterdir():
            if item.is_dir():
                shutil.rmtree(item)
            elif force or item.lstat().st_mtime < old_limit:
                item.unlink()

    async def _copy_file_old_to_new(self, pool: ThreadPoolExecutor, src: Path, trg: Path) -> bool:
        raise NotImplementedError()

    async def _peek_remote_item(
        self, pool: ThreadPoolExecutor, path: str, raise_on_fail: bool = False, backup_done: str = None
    ) -> Optional[dict]:
        request = {"type": "FS-PEEK", "path": path, "o": 0}
        if backup_done:
            request["backup_done"] = backup_done
        result = None

        def _resp_cb(resp_i: int, resp: dict, sp: NanoResponseStreamProcessor):  # NOTE: Called in threadpool.
            nonlocal result
            tmp = resp["response"].get("pk")
            if (
                isinstance(tmp, dict)
                and isinstance(tmp.get("name"), str)
                and isinstance(tmp.get("mtime"), float)
                and (tmp.get("size") is None or isinstance(tmp.get("size"), int))
            ):
                result = tmp

        stream_processor = await self._remote_nano_request(pool, request, _resp_cb)
        if stream_processor.error:
            if raise_on_fail:
                raise NanoResponseError(f"Peek item failed {stream_processor.error}")
        return result

    async def _list_remote_dir(self, pool: ThreadPoolExecutor, path: str, raise_on_fail: bool = False) -> list:
        while True:
            request = {"type": "FS-GET" if self._use_old_fs_get_listing else "FS-LS", "path": path}
            accum = []

            def _resp_cb(resp_i: int, resp: dict, sp: NanoResponseStreamProcessor):  # NOTE: Called in threadpool.
                nonlocal accum
                if self._use_old_fs_get_listing:
                    if resp.get("part") == 0:
                        del accum[:]
                    accum.append(resp["response"])
                elif resp.get("part") != 0:  # FS-LS initial response is for caching-transfer setup
                    if resp.get("part") == 1:
                        del accum[:]
                    accum.extend(resp["response"]["ls"])

            stream_processor = await self._remote_nano_request(pool, request, _resp_cb)
            if stream_processor.error:
                if stream_processor.error == NANO_REQUEST_ERROR_UNKNOWN_TYPE and not self._use_old_fs_get_listing:
                    # Switch over to FS-GET for drive version prior to 5.
                    self._use_old_fs_get_listing = True
                    continue
                if raise_on_fail:
                    raise NanoResponseError(f"Listing directory failed {stream_processor.error}")
                return []  # Skip single failures to progress on.
            if self._use_old_fs_get_listing:
                return merge_fs_get_listing_responses(accum)
            return accum

    async def _save_remote_file(
        self, pool: ThreadPoolExecutor, tmp_dir: Optional[Path], path: str, new_path: Path, mtime: float, size: int
    ):
        request = {"type": "FS-GET", "path": path}
        if not tmp_dir:
            tmp_dir = new_path.parent
        # get tmp-file tracking info or set it up
        tmp_row = await self.temp_db.get_tmp_file(new_path.as_posix())
        if not tmp_row:
            tmp_name = DRIVE_PUT_FILE.reserve(tmp_dir, name=new_path.name, try_as_is=False)
            meta = b""  # not yet set
            tmp_id = await self.temp_db.create_tmp_file(
                self.context[NanoRequestContext.ACCOUNT_ID], new_path.as_posix(), tmp_name, meta
            )
        else:
            tmp_id, account_id, tmp_name, meta = tmp_row
        buf = []
        part = wrote = 0
        save_path = tmp_dir / tmp_name  # tmp prefix will prevent it from being indexed
        try:
            sps = save_path.lstat()
        except OSError:
            pass  # does not exist most likely
        else:
            if stat.S_ISREG(sps.st_mode):
                wrote = sps.st_size
                part = wrote // self.context[NanoRequestContext.PART_SIZE]
            else:
                await self.temp_db.drop_tmp_file(tmp_id)
                raise Exception("Local temp file invalid for inter-nano download")
        if part <= 1:
            part = 0
            meta = b""
        last_key = meta

        def _resp_cb(resp_i: int, resp: dict, sp: NanoResponseStreamProcessor):  # NOTE: Called in threadpool.
            nonlocal buf, wrote, last_key
            if sp.cache_key != last_key or resp.get("part") == 1:
                # Any time the cache-key changes or we're receiving the first data-part we are restarting the transfer.
                # NOTE: sp.transfer_key could be set instead (for empty file transfer, this is fine)
                with open(save_path, "wb") as f:
                    f.truncate()
                wrote = 0
                del buf[:]
                last_key = sp.cache_key
            wrote += write_buffered_fs_get_file(save_path, resp["response"], buf, wrote, self.inr.remote_kr)

        try:
            stream_processor = await self._remote_nano_request(pool, request, _resp_cb, part, cache_key=meta or None)
        finally:
            # After running set the meta, so resuming is possible.
            if last_key and last_key != meta:
                await self.temp_db.set_tmp_file_meta(tmp_id, last_key)
        if not stream_processor.error:
            await attempt_share_locked_operation(
                3,
                3.0,
                self._loop.run_in_executor,
                pool,
                self._finalize_save_remote_file,
                save_path,
                new_path,
                size,
                mtime,
            )
            await self.temp_db.drop_tmp_file(tmp_id)
        elif stream_processor.error == NANO_REQUEST_ERROR_HANDLER_TARGET_NOT_FOUND:
            if await self._loop.run_in_executor(pool, self._abandon_save_remote_file, save_path):
                await self.temp_db.drop_tmp_file(tmp_id)

    def _finalize_save_remote_file(self, save_path: Path, new_path: Path, size: int, mtime: float):
        if os.lstat(save_path).st_size == size:
            os.utime(save_path, (time.time(), mtime))
            os.replace(save_path, new_path)
        else:
            raise NanoResponseError("File saving result size does not match")

    def _abandon_save_remote_file(self, save_path: Path):
        try:
            save_path.unlink(missing_ok=True)
            return True
        except Exception:
            pass  # fallback to tmp db cleanup

    async def _remote_nano_request(
        self,
        pool: ThreadPoolExecutor,
        payload: dict,
        response_handler_cb: Callable[[int, dict, NanoResponseStreamProcessor], None],
        part: int = 0,
        cache_key: Optional[bytes] = None,
        transfer_key: Optional[bytes] = None,
    ) -> NanoResponseStreamProcessor:
        """
        Simple request processing, with simple retry support and continue.
        We're implementing an error-forgiving logic here, as it may be better to give up on a specific file/directory
        now after a few retries than blocking all processing on some broken item. If strict error handling is desired,
        raise an appropriate exception if  the `error` attribute is set on the returned processor object.
        """
        okay = 0
        retry = 0
        bad = 0
        request = await self._loop.run_in_executor(pool, self.inr.encrypt_request, payload)

        def _response_handler_error_wrapper(i_, r_):  # NOTE: Called in threadpool.
            nonlocal okay, retry
            if stream_processor.error:
                raise NanoResponseError(f"Error response received {stream_processor.error}", r_)
            if stream_processor.server_error:
                raise NanoResponseError(f"Error response received from server {stream_processor.server_error}", r_)
            if okay >= 15 and retry > 0:
                retry -= 1
                okay = 0
            else:
                okay += 1
            response_handler_cb(i_, r_, stream_processor)

        stream_processor = NanoResponseStreamProcessor(
            request,
            self.inr.self_kr,
            self.inr.remote_kr,
            self.inr.room_kr,
            self.inr.adhoc_kr,
            False,
            part,
            cache_key,
            transfer_key,
            _response_handler_error_wrapper,
        )
        while True:
            reset = False
            try:
                await self.inr.perform(pool, stream_processor)
                stream_processor.finish()
                return stream_processor
            except NanoResponseError as e:
                if stream_processor.server_error:
                    # if stream_processor.server_error == SERVER_ERROR_TOO_MANY_REQUESTS:
                    raise e  # Server errors are related to network connection, just propagate and retry whole process.
                if stream_processor.error is None:
                    # NOTE: This should not really happen actually. This indicates some precondition in the response
                    # processor has failed. There is not much to do specifically, so we just retry a few times.
                    stream_processor.error = NANO_REQUEST_ERROR  # generic error, just for indication of some issue
                    logger.exception("Nano response error without reported-error code")
                elif stream_processor.error in (
                    NANO_REQUEST_ERROR_HANDLER_RESET_REQUEST,
                    NANO_REQUEST_ERROR_PART_RANGE,
                ):
                    reset = True
                elif stream_processor.error in (NANO_REQUEST_ERROR_HANDLER_TEMPORARY_FAILURE,):
                    pass  # Retry transient errors.
                elif stream_processor.error == NANO_REQUEST_ERROR_HANDLER_BAD_REQUEST:
                    # Bad request may be some transient state, but it is unclear so count it separately.
                    if bad >= 1:  # max one retry
                        break
                    bad += 1
                elif stream_processor.error == NANO_REQUEST_ERROR_HANDLER_TARGET_NOT_FOUND:
                    break  # Concrete errors that we know not to retry.
                elif stream_processor.error == NANO_REQUEST_ERROR_TOO_MANY_REQUESTS:
                    raise e  # Propagate and retry whole process later.
                # else: Unspecified errors are retried like transient.
            if retry >= len(self.REQUEST_RETRY_DELAYS):
                break
            # delay before retry
            await asyncio.sleep(self.REQUEST_RETRY_DELAYS[retry])
            if reset:
                stream_processor.reset()
            else:
                stream_processor.retry()
            retry += 1
            okay = 0
        return stream_processor
