import asyncio
import json
import os
import signal
import uuid
from datetime import datetime
from typing import Dict, Iterable, List, Optional, cast

from fastapi import WebSocket, WebSocketDisconnect  # type: ignore
from starlette.websockets import WebSocketState

from ..kernel_driver.connect import (
    cfg_t,
    connect_channel,
    launch_kernel,
    read_connection_file,
)
from ..kernel_driver.connect import (
    write_connection_file as _write_connection_file,  # type: ignore
)
from ..kernel_driver.message import create_message, receive_message, send_message
from .message import (  # type: ignore
    deserialize_msg_from_ws_v1,
    from_binary,
    get_msg_from_parts,
    get_parent_header,
    get_zmq_parts,
    send_raw_message,
    serialize_msg_to_ws_v1,
    to_binary,
)

kernels: dict = {}


class AcceptedWebSocket:
    _websocket: WebSocket
    _accepted_subprotocol: Optional[str]

    def __init__(self, websocket, accepted_subprotocol):
        self._websocket = websocket
        self._accepted_subprotocol = accepted_subprotocol

    @property
    def websocket(self):
        return self._websocket

    @property
    def accepted_subprotocol(self):
        return self._accepted_subprotocol


class KernelServer:
    def __init__(
        self,
        kernelspec_path: str = "",
        kernel_cwd: str = "",
        connection_cfg: Optional[cfg_t] = None,
        connection_file: str = "",
        write_connection_file: bool = True,
        capture_kernel_output: bool = True,
    ) -> None:
        self.capture_kernel_output = capture_kernel_output
        self.kernelspec_path = kernelspec_path
        self.kernel_cwd = kernel_cwd
        self.connection_cfg = connection_cfg
        self.connection_file = connection_file
        self.write_connection_file = write_connection_file
        self.channel_tasks: List[asyncio.Task] = []
        self.sessions: Dict[str, AcceptedWebSocket] = {}
        # blocked messages and allowed messages are mutually exclusive
        self.blocked_messages: List[str] = []
        self.allowed_messages: Optional[List[str]] = None  # when None, all messages are allowed
        # when [], no message is allowed
        self.setup_connection_file()

    def setup_connection_file(self):
        if self.write_connection_file:
            self.connection_file_path, self.connection_cfg = _write_connection_file(
                self.connection_file
            )
        elif self.connection_file:
            self.connection_file_path = self.connection_file
            self.connection_cfg = read_connection_file(self.connection_file)
        else:
            if self.connection_cfg is None:
                raise RuntimeError("No connection_cfg")
        self.key = cast(str, self.connection_cfg["key"])

    def block_messages(self, message_types: Iterable[str] = []):
        # if using blocked messages, discard allowed messages
        self.allowed_messages = None
        if isinstance(message_types, str):
            message_types = [message_types]
        self.blocked_messages = list(message_types)

    def allow_messages(self, message_types: Optional[Iterable[str]] = None):
        # if using allowed messages, discard blocked messages
        self.blocked_messages = []
        if message_types is None:
            self.allowed_messages = None
            return
        if isinstance(message_types, str):
            message_types = [message_types]
        self.allowed_messages = list(message_types)

    @property
    def connections(self) -> int:
        return len(self.sessions)

    async def start(self) -> None:
        if not self.kernelspec_path:
            raise RuntimeError("Could not find a kernel, maybe you forgot to install one?")
        self.last_activity = {
            "date": datetime.utcnow().isoformat() + "Z",
            "execution_state": "starting",
        }
        self.kernel_process = await launch_kernel(
            self.kernelspec_path,
            self.connection_file_path,
            self.kernel_cwd,
            self.capture_kernel_output,
        )
        assert self.connection_cfg is not None
        identity = uuid.uuid4().hex.encode("ascii")
        self.shell_channel = connect_channel("shell", self.connection_cfg, identity=identity)
        self.stdin_channel = connect_channel("stdin", self.connection_cfg, identity=identity)
        self.control_channel = connect_channel("control", self.connection_cfg, identity=identity)
        self.iopub_channel = connect_channel("iopub", self.connection_cfg)
        await self._wait_for_ready()
        self.channel_tasks += [
            asyncio.create_task(self.listen("shell")),
            asyncio.create_task(self.listen("stdin")),
            asyncio.create_task(self.listen("control")),
            asyncio.create_task(self.listen("iopub")),
        ]

    async def stop(self) -> None:
        # FIXME: stop kernel in a better way
        try:
            self.kernel_process.send_signal(signal.SIGINT)
            self.kernel_process.kill()
            await self.kernel_process.wait()
        except Exception:
            pass
        try:
            os.remove(self.connection_file_path)
        except Exception:
            pass
        for task in self.channel_tasks:
            task.cancel()
        self.channel_tasks = []

    async def restart(self) -> None:
        await self.stop()
        self.setup_connection_file()
        await self.start()

    async def serve(
        self,
        websocket: AcceptedWebSocket,
        session_id: str,
        permissions: Optional[Dict[str, List[str]]],
    ):
        self.sessions[session_id] = websocket
        self.can_execute = permissions is None or "execute" in permissions.get("kernels", [])
        await self.listen_web(websocket)
        # the session could have been removed through the REST API, so check if it still exists
        if session_id in self.sessions:
            del self.sessions[session_id]

    async def listen_web(self, websocket: AcceptedWebSocket):
        try:
            await self.send_to_zmq(websocket)
        except WebSocketDisconnect:
            pass

    async def listen(self, channel_name: str):
        if channel_name == "shell":
            channel = self.shell_channel
        elif channel_name == "control":
            channel = self.control_channel
        elif channel_name == "iopub":
            channel = self.iopub_channel
        elif channel_name == "stdin":
            channel = self.stdin_channel

        while True:
            parts = await get_zmq_parts(channel)
            parent_header = get_parent_header(parts)
            if channel == self.iopub_channel:
                # broadcast to all web clients
                for websocket in self.sessions.values():
                    await self.send_to_ws(websocket, parts, parent_header, channel_name)
            else:
                session = parent_header["session"]
                if session in self.sessions:
                    websocket = self.sessions[session]
                    await self.send_to_ws(websocket, parts, parent_header, channel_name)

    async def _wait_for_ready(self):
        while True:
            msg = create_message("kernel_info_request")
            await send_message(msg, self.shell_channel, self.key)
            msg = await receive_message(self.shell_channel, timeout=0.2)
            if msg is not None and msg["msg_type"] == "kernel_info_reply":
                msg = await receive_message(self.iopub_channel, timeout=0.2)
                if msg is None:
                    # IOPub not connected, start over
                    pass
                else:
                    break

    async def send_to_zmq(self, websocket):
        if not websocket.accepted_subprotocol:
            while True:
                msg = await receive_json_or_bytes(websocket.websocket)
                if not self.can_execute:
                    continue
                msg_type = msg["header"]["msg_type"]
                if (msg_type in self.blocked_messages) or (
                    self.allowed_messages is not None and msg_type not in self.allowed_messages
                ):
                    continue
                channel = msg.pop("channel")
                if channel == "shell":
                    await send_message(msg, self.shell_channel, self.key)
                elif channel == "control":
                    await send_message(msg, self.control_channel, self.key)
                elif channel == "stdin":
                    await send_message(msg, self.stdin_channel, self.key)
        elif websocket.accepted_subprotocol == "v1.kernel.websocket.jupyter.org":
            while True:
                msg = await websocket.websocket.receive_bytes()
                if not self.can_execute:
                    continue
                channel, parts = deserialize_msg_from_ws_v1(msg)
                # NOTE: we parse the header for message filtering
                # it is not as bad as parsing the content
                header = json.loads(parts[0])
                msg_type = header["msg_type"]
                if (msg_type in self.blocked_messages) or (
                    self.allowed_messages is not None and msg_type not in self.allowed_messages
                ):
                    continue
                if channel == "shell":
                    await send_raw_message(parts, self.shell_channel, self.key)
                elif channel == "control":
                    await send_raw_message(parts, self.control_channel, self.key)
                elif channel == "stdin":
                    await send_raw_message(parts, self.stdin_channel, self.key)

    async def send_to_ws(self, websocket, parts, parent_header, channel_name):
        if not websocket.accepted_subprotocol:
            # default, "legacy" protocol
            msg = get_msg_from_parts(parts, parent_header=parent_header)
            msg["channel"] = channel_name
            await send_json_or_bytes(websocket.websocket, msg)
            if channel_name == "iopub":
                if "content" in msg and "execution_state" in msg["content"]:
                    self.last_activity = {
                        "date": msg["header"]["date"],
                        "execution_state": msg["content"]["execution_state"],
                    }
        elif websocket.accepted_subprotocol == "v1.kernel.websocket.jupyter.org":
            bin_msg = serialize_msg_to_ws_v1(parts, channel_name)
            try:
                await websocket.websocket.send_bytes(bin_msg)
            except Exception:
                pass
            # FIXME: update last_activity
            # but we don't want to parse the content!
            # or should we request it from the control channel?


async def receive_json_or_bytes(websocket):
    assert websocket.application_state == WebSocketState.CONNECTED
    message = await websocket.receive()
    websocket._raise_on_disconnect(message)
    if "text" in message:
        return json.loads(message["text"])
    msg = from_binary(message["bytes"])
    return msg


async def send_json_or_bytes(websocket, msg):
    bmsg = to_binary(msg)
    if bmsg is None:
        await websocket.send_json(msg)
    else:
        await websocket.send_bytes(bmsg)
