import asyncio
import hashlib
import json
import logging
import os
import shutil
import sys
import time
import weakref
from contextlib import contextmanager
from pathlib import Path
from typing import Optional, Set

from aiohttp import web, WSMsgType, WSMessage, WSCloseCode
from aiohttp.web_response import json_response
from clarabot_nano_common.utility.serialize import msgpack_loads, msgpack_dumps
from ttldict2 import TTLDict

from clarabot_nano.common.auth.two_factor import AbstractTwoFactorErrorHandler, TwoFactorError
from clarabot_nano.common.config.client import AbstractAioConfiguredWorkerProcess
from clarabot_nano.common.core.message import (
    MSG_TYPE_CONTROL_REQ_STATS,
    MSG_TYPE_CONTROL_INDEXER_WORK,
    MSG_TYPE_CONTROL_SERVICE_ERROR,
    MSG_TYPE_GATEWAY_LOGOUT,
    MSG_TYPE_GATEWAY_LOGIN,
    MSG_TYPE_REQUEST_CONFIG_RESET,
    MSG_TYPE_REQUEST_STATE_UPDATE,
    MSG_TYPE_REQUEST_CONFIG_UPDATE,
)
from clarabot_nano.common.spawn.constants import ShortNames
from clarabot_nano.common.zmq.requester import AioClientRequester
from clarabot_nano.control_api.config import ControlApiConfig

logger = logging.getLogger(__name__)


class TfaCodeRequired(Exception):
    pass


class ControlApiTwoFactorHandler(AbstractTwoFactorErrorHandler):
    def _query_user_2fa_code(self) -> Optional[str]:
        raise TfaCodeRequired()


class ControlApiProcess(AbstractAioConfiguredWorkerProcess):
    STATE_KEYS = [
        "gateway_status",
        "indexer_status",
        "gateway_update",
        "vendor",
        "license",
        "device_id",
        "license_local_validity",
        "license_server_validity",
    ]
    REMOTE_KEYS = [
        "name",
    ]

    def __init__(self):
        AbstractAioConfiguredWorkerProcess.__init__(self, ShortNames.CONTROL_API, extracted_config_cls=ControlApiConfig)
        self._core_requester: AioClientRequester = None
        self._core_handler.add_handler(MSG_TYPE_CONTROL_REQ_STATS, self._handle_request_stats)
        self._core_handler.add_handler(MSG_TYPE_CONTROL_INDEXER_WORK, self._handle_indexer_work)
        self._core_handler.add_handler(MSG_TYPE_CONTROL_SERVICE_ERROR, self._handle_service_error)
        self._login_states = TTLDict(300.0)
        self._streams: Set[asyncio.Queue] = set()
        self._websockets = weakref.WeakSet()
        # init config state to keep in sync over ws
        config = self.get_config()
        initial_remote = config["remote"]
        initial_state = config["state"]
        self._config_state = {}
        for k in self.STATE_KEYS:
            self._config_state[k] = initial_state.get(k)
        for k in self.REMOTE_KEYS:
            self._config_state[k] = initial_remote.get(k)

    async def _handle_core_updated_config(self, message):
        await AbstractAioConfiguredWorkerProcess._handle_core_updated_config(self, message)
        config = self.get_config()
        changed = False
        for attr in self.REMOTE_KEYS:
            value = config["remote"].get(attr)
            if value != self._config_state.get(attr):
                self._config_state[attr] = value
                changed = True
        for attr in self.STATE_KEYS:
            value = config["state"].get(attr)
            if value != self._config_state.get(attr):
                self._config_state[attr] = value
                changed = True
        if changed:
            self._push_state()

    @contextmanager
    def _stream(self):
        s = asyncio.Queue()
        try:
            self._streams.add(s)
            yield s
        finally:
            self._streams.discard(s)

    def _push_message(self, message: dict):
        for s in self._streams:
            s.put_nowait(json.dumps(message))

    def _push_state(self):
        self._push_message({"type": "state", **self._config_state})

    def _handle_request_stats(self, message):
        req_per_s = msgpack_loads(message[1]).get("reqs")
        if not isinstance(req_per_s, float):
            req_per_s = 0.0
        self._push_message({"type": "req_stat", "req_per_s": req_per_s})

    def _handle_indexer_work(self, message):
        work = msgpack_loads(message[1]).get("work")
        if not isinstance(work, str):
            work = ""
        self._push_message({"type": "idx_work", "work": work})

    def _handle_service_error(self, message):
        error_message = msgpack_loads(message[1])["msg"]
        self._push_message({"type": "err_message", "msg": error_message})

    async def _main_wrapper(self):
        self._core_requester = self.create_core_requester()
        with self._core_requester:
            await AbstractAioConfiguredWorkerProcess._main_wrapper(self)

    async def _work(self):
        # reset the socket path
        socket_path: Path = self._extracted_config.control_api_path
        socket_path.parent.mkdir(parents=True, exist_ok=True)
        socket_path.unlink(missing_ok=True)
        # setup web-app api
        app = web.Application()
        app.add_routes(
            [
                web.get("/status", self.handle_ws_status),
                # web.get("/access-log", self.handle_get_access_log),
                # web.get("/process-log", self.handle_get_process_log),
                web.post("/do-login", self.handle_do_login),
                # web.post("/do-logout", self.handle_do_logout),
                web.post("/do-reset", self.handle_do_reset),
                web.post("/new-license", self.handle_do_new_license),
                web.post("/rename", self.handle_do_rename),
            ]
        )
        app.on_shutdown.append(self._on_shutdown)
        # run aiohttp app in current loop and exit with the normal process
        runner = web.AppRunner(app)
        await runner.setup()
        site = (web.NamedPipeSite if sys.platform.startswith("win") else web.UnixSite)(runner, str(socket_path))
        await site.start()
        if not sys.platform.startswith("win"):
            if self._extracted_config.control_api_group is not None:
                shutil.chown(socket_path, group=self._extracted_config.control_api_group)
            if self._extracted_config.control_api_mode is not None:
                os.chmod(socket_path, self._extracted_config.control_api_mode)
        try:
            while True:
                await asyncio.sleep(60)
        finally:
            await runner.cleanup()

    async def _on_shutdown(self, app):
        for ws in set(app["websockets"]):
            await ws.close(code=WSCloseCode.GOING_AWAY, message="Server shutdown")

    # convenience functions
    async def _handle_core_request(self, command, request):
        reply = await self._core_requester.request([command, request], timeout=10)  # timeout for debug
        if len(reply) == 1:
            result = msgpack_loads(reply[0])
        else:
            result = {"ok": None}
        return result

    async def _gateway_request(self, request_frames, timeout: float = 16.0):
        gcrn = self.get_config()["state"].get("gateway_command_result_number")
        self._core_requester.notify(request_frames)
        # mega-simple wait for request completion; this is sufficient for debugging...
        until = time.perf_counter() + timeout
        while gcrn == self.get_config()["state"].get("gateway_command_result_number"):
            if time.perf_counter() > until:
                return {"result": "gateway request timeout"}
            await asyncio.sleep(0.1)
        return {"result": self.get_config()["state"].get("gateway_command_result")}

    # nano api
    async def handle_do_login(self, request: web.Request):
        data = await request.json()
        state_identity = hashlib.sha256(
            f"GsQ2d0BfT2AMuTJVpLaT\n{data['email']}\n{data['password']}".encode("utf-8")
        ).hexdigest()
        state = self._login_states.get(state_identity)
        if not state:
            state = {"server_session_key": os.urandom(16)}
        self._login_states[state_identity] = state  # set/keep alive
        two_factor_handler = ControlApiTwoFactorHandler()
        try:
            i = 0
            while True:
                i += 1
                req = {"email": data["email"], "password": data["password"], **state}
                if "2faCode" in data:
                    req["2faCode"] = data["2faCode"]
                result = await self._gateway_request([MSG_TYPE_GATEWAY_LOGIN, msgpack_dumps(req)])
                if tfa := two_factor_handler.get_two_factor(result["result"]):
                    extra_args = two_factor_handler.handle_two_factor_response(*tfa)
                    assert "2faCode" not in extra_args
                    if "otpSalt" in extra_args:
                        state["otpSalt"] = extra_args["otpSalt"]
                    if i <= 8:  # prevent endless loop, just in case
                        continue
                break
        except TwoFactorError as e:
            return json_response({"result": e.message})
        except TfaCodeRequired:
            return json_response({"result": "2faCode-required"})
        return json_response(result)

    async def handle_do_reset(self, request: web.Request):
        result = await self._gateway_request([MSG_TYPE_GATEWAY_LOGOUT, b""])
        # NOTE: The gateway always destroys the local tokens, even if remote can't be notified, so explicit success
        # or failure responses shall continue the reset here.
        if isinstance(result, dict) and result.get("result") in ("logout successful", "logout failed"):
            result = await self._handle_core_request(MSG_TYPE_REQUEST_CONFIG_RESET, msgpack_dumps({}))
            if result.get("ok"):
                result["result"] = "ok"
        return json_response(result)

    async def handle_do_new_license(self, request: web.Request):
        data = await request.json()
        new_license = data.get("new_license")
        if not isinstance(new_license, str):
            return json_response({"error": "new_license must be str"})
        result = await self._handle_core_request(
            MSG_TYPE_REQUEST_STATE_UPDATE, msgpack_dumps({"_handler_key": "new_license", "license": new_license})
        )
        if result.get("ok"):
            result["result"] = "ok"
        return json_response(result)

    async def handle_do_rename(self, request: web.Request):
        data = await request.json()
        name = data.get("name")
        if not isinstance(name, str):
            return json_response({"error": "name must be str"})
        result = await self._handle_core_request(
            MSG_TYPE_REQUEST_CONFIG_UPDATE, msgpack_dumps({"_handler_key": "name", "name": name})
        )
        if result.get("ok"):
            result["result"] = "ok"
        return json_response(result)

    async def handle_ws_status(self, request: web.Request):
        ws = web.WebSocketResponse()
        await ws.prepare(request)

        async def _reader():
            async for msg in ws:
                msg: WSMessage
                if msg.type == WSMsgType.TEXT:
                    try:
                        await self._handle_ws_status_request(msg.json())
                    except Exception:
                        logger.error("ws-status error for %s", msg)
                        raise
                elif msg.type == WSMsgType.ERROR:
                    logger.error("ws-status connection error %s %s", msg, ws.exception())
                    break
                else:
                    logger.error("ws-status connection unexpected message %s", msg)
                    break

        async def _writer():
            with self._stream() as s:
                self._push_state()
                s: asyncio.Queue
                while True:
                    x = await s.get()
                    await ws.send_str(x)

        _r = asyncio.create_task(_reader())
        _w = asyncio.create_task(_writer())
        self._websockets.add(ws)
        try:
            try:
                _, pending = await asyncio.wait([_r, _w], return_when=asyncio.FIRST_COMPLETED)
            finally:
                _r.cancel()
                _w.cancel()
                await asyncio.gather(_r, _w, return_exceptions=True)
                await ws.close()
        finally:
            self._websockets.discard(ws)
        return ws

    async def _handle_ws_status_request(self, data):
        if not isinstance(data, dict):
            raise Exception("ws-req not dict")
