# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.

from __future__ import annotations

import asyncio
import json
import uuid
from logging import Logger
from typing import Any
from uuid import uuid4

from jupyter_server.auth import authorized
from jupyter_server.base.handlers import APIHandler, JupyterHandler
from jupyter_server.utils import ensure_async
from jupyter_ydoc import ydocs as YDOCS
from pycrdt import Doc, UndoManager
from pycrdt.store import BaseYStore
from pycrdt.websocket import YRoom
from tornado import web
from tornado.websocket import WebSocketHandler

from .loaders import FileLoaderMapping
from .rooms import DocumentRoom, TransientRoom
from .utils import (
    JUPYTER_COLLABORATION_AWARENESS_EVENTS_URI,
    JUPYTER_COLLABORATION_EVENTS_URI,
    JUPYTER_COLLABORATION_FORK_EVENTS_URI,
    LogLevel,
    decode_file_path,
    encode_file_path,
    room_id_from_encoded_path,
)
from .websocketserver import JupyterWebsocketServer, RoomNotFound

YFILE = YDOCS["file"]


SERVER_SESSION = str(uuid.uuid4())
FORK_DOCUMENTS = {}
FORK_ROOMS: dict[str, dict[str, str]] = {}


class YDocWebSocketHandler(WebSocketHandler, JupyterHandler):
    """`YDocWebSocketHandler` uses the singleton pattern for ``WebsocketServer``,
    which is a subclass of ypy-websocket's ``WebsocketServer``.

    In ``WebsocketServer``, we expect to use a WebSocket object as follows:

    - receive messages until the connection is closed with
       ``for message in websocket: pass``.
    - send a message with `await websocket.send(message)`.

    Tornado's WebSocket API is different, so ``YDocWebSocketHandler`` needs to be adapted:

    - ``YDocWebSocketHandler`` is an async iterator, that will yield received messages.
       Messages received in Tornado's `on_message(message)` are put in an async
       ``_message_queue``, from which we get them asynchronously.
    - The ``send(message)`` method is async and calls Tornado's ``write_message(message)``.
    - Although it's currently not used in ypy-websocket, ``recv()`` is an async method for
       receiving a message.
    """

    _message_queue: asyncio.Queue[Any]
    _background_tasks: set[asyncio.Task]
    _room_locks: dict[str, asyncio.Lock] = {}

    def _room_lock(self, room_id: str) -> asyncio.Lock:
        if room_id not in self._room_locks:
            self._room_locks[room_id] = asyncio.Lock()
        return self._room_locks[room_id]

    def create_task(self, aw):
        task = asyncio.create_task(aw)
        self._background_tasks.add(task)
        task.add_done_callback(self._background_tasks.discard)

    async def prepare(self):
        await ensure_async(super().prepare())
        if not self._websocket_server.started.is_set():
            self.create_task(self._websocket_server.start())
            await self._websocket_server.started.wait()

        # Get room
        self._room_id: str = room_id_from_encoded_path(self.request.path)

        async with self._room_lock(self._room_id):
            if self._websocket_server.room_exists(self._room_id):
                self.room: YRoom = await self._websocket_server.get_room(self._room_id)
            else:
                # Logging exceptions, instead of raising them here to ensure
                # that the y-rooms stay alive even after an exception is seen.
                def exception_logger(exception: Exception, log: Logger) -> bool:
                    """A function that catches any exceptions raised in the websocket
                    server and logs them.

                    The protects the y-room's task group from cancelling
                    anytime an exception is raised.
                    """
                    log.error(
                        f"Document Room Exception, (room_id={self._room_id or 'unknown'}): ",
                        exc_info=exception,
                    )
                    return True

                if self._room_id.count(":") >= 2:
                    # DocumentRoom
                    file_format, file_type, file_id = decode_file_path(self._room_id)
                    if file_id in self._file_loaders:
                        self._emit(
                            LogLevel.WARNING,
                            None,
                            "There is another collaborative session accessing the same file.\nThe synchronization between rooms is not supported and you might lose some of your changes.",
                        )

                    file = self._file_loaders[file_id]
                    updates_file_path = f".{file_type}:{file_id}.y"
                    ystore = self._ystore_class(
                        path=updates_file_path,
                        log=self.log,  # type:ignore[call-arg]
                    )
                    self.room = DocumentRoom(
                        self._room_id,
                        file_format,
                        file_type,
                        file,
                        self.event_logger,
                        ystore,
                        self.log,
                        exception_handler=exception_logger,
                        save_delay=self._document_save_delay,
                    )

                else:
                    # TransientRoom
                    # it is a transient document (e.g. awareness)
                    self.room = TransientRoom(
                        self._room_id,
                        log=self.log,
                        exception_handler=exception_logger,
                    )

                    if self._room_id == "JupyterLab:globalAwareness":
                        # Listen for the changes in GlobalAwareness to update users
                        self.room.awareness.observe(self._on_global_awareness_event)

            try:
                await self._websocket_server.start_room(self.room)
            except Exception as e:
                self.log.error("Room %s failed to start on websocket server", self._room_id)
                # Clean room
                await self.room.stop()
                self.log.info("Room %s deleted", self._room_id)
                self._emit(LogLevel.INFO, "clean", "Room deleted.")

                # Clean the file loader in file loader mapping if there are not any rooms using it
                _, _, file_id = decode_file_path(self._room_id)
                file = self._file_loaders[file_id]
                if file.number_of_subscriptions == 0 or (
                    file.number_of_subscriptions == 1 and self._room_id in file._subscriptions
                ):
                    self.log.info("Deleting file %s", file.path)
                    await self._file_loaders.remove(file_id)
                    self._emit(LogLevel.INFO, "clean", "file loader removed.")
                raise e
            self._websocket_server.add_room(self._room_id, self.room)

    def initialize(
        self,
        ywebsocket_server: JupyterWebsocketServer,
        file_loaders: FileLoaderMapping,
        ystore_class: type[BaseYStore],
        document_cleanup_delay: float | None = 60.0,
        document_save_delay: float | None = 1.0,
    ) -> None:
        self._background_tasks = set()
        # File ID manager cannot be passed as argument as the extension may load after this one
        self._file_id_manager = self.settings["file_id_manager"]
        self._file_loaders = file_loaders
        self._ystore_class = ystore_class
        self._cleanup_delay = document_cleanup_delay
        self._document_save_delay = document_save_delay
        self._websocket_server = ywebsocket_server
        self._message_queue = asyncio.Queue()
        self._room_id = ""
        self.room = None  # type:ignore

    @property
    def path(self):
        """
        Returns the room id. It needs to be called 'path' for compatibility with
        the WebsocketServer (websocket.path).
        """
        return self._room_id

    @property
    def max_message_size(self):
        """
        Override max_message size to 1GB
        """
        return 1024 * 1024 * 1024

    def __aiter__(self):
        # needed to be compatible with WebsocketServer (async for message in websocket)
        return self

    async def __anext__(self):
        # needed to be compatible with WebsocketServer (async for message in websocket)
        message = await self._message_queue.get()
        if not message:
            raise StopAsyncIteration()
        return message

    async def get(self, *args, **kwargs):
        """
        Overrides default behavior to check whether the client is authenticated or not.
        """
        if self.current_user is None:
            self.log.warning("Couldn't authenticate WebSocket connection")
            raise web.HTTPError(403)
        return await super().get(*args, **kwargs)

    async def open(self, room_id: str) -> None:  # type:ignore[override]
        """
        On connection open.
        """
        self.create_task(self._websocket_server.serve(self))

        if isinstance(self.room, DocumentRoom):
            # Close the connection if the document session expired
            session_id = self.get_query_argument("sessionId", "")
            if SERVER_SESSION != session_id:
                self.close(
                    1003,
                    f"Document session {session_id} expired. You need to reload this browser tab.",
                )

            # cancel the deletion of the room if it was scheduled
            if self.room.cleaner is not None:
                self.room.cleaner.cancel()

            try:
                # Initialize the room
                async with self._room_lock(self._room_id):
                    await self.room.initialize()
                self._emit_awareness_event(self.current_user.username, "join")
            except Exception as e:
                _, _, file_id = decode_file_path(self._room_id)
                file = self._file_loaders[file_id]

                # Close websocket and propagate error.
                if isinstance(e, web.HTTPError):
                    self.log.error(f"File {file.path} not found.\n{e!r}", exc_info=e)
                    self.close(1004, f"File {file.path} not found.")
                else:
                    self.log.error(f"Error initializing: {file.path}\n{e!r}", exc_info=e)
                    self.close(
                        1003,
                        f"Error initializing: {file.path}. You need to close the document.",
                    )

                # Clean up the room and delete the file loader
                if len(self.room.clients) == 0 or self.room.clients == {self}:
                    self._message_queue.put_nowait(b"")
                    self._cleanup_delay = 0
                    await self._clean_room()

            self._emit(LogLevel.INFO, "initialize", "New client connected.")
        else:
            if self._room_id != "JupyterLab:globalAwareness":
                self._emit_awareness_event(self.current_user.username, "join")

    async def send(self, message):
        """
        Send a message to the client.
        """
        # needed to be compatible with WebsocketServer (websocket.send)
        try:
            self.write_message(message, binary=True)
        except Exception as e:
            self.log.error("Failed to write message", exc_info=e)

    async def recv(self):
        """
        Receive a message from the client.
        """
        message = await self._message_queue.get()
        return message

    async def on_message(self, message):
        """
        On message receive.
        """
        self._message_queue.put_nowait(message)
        self._websocket_server.ypatch_nb += 1

    def on_close(self) -> None:
        """
        On connection close.
        """
        # stop serving this client
        self._message_queue.put_nowait(b"")
        if isinstance(self.room, DocumentRoom) and self.room.clients == {self}:
            # no client in this room after we disconnect
            # keep the document for a while in case someone reconnects
            self.log.info("Cleaning room: %s", self._room_id)
            self.room.cleaner = asyncio.create_task(self._clean_room())
        if self._room_id != "JupyterLab:globalAwareness":
            self._emit_awareness_event(self.current_user.username, "leave")

    def _emit(self, level: LogLevel, action: str | None = None, msg: str | None = None) -> None:
        _, _, file_id = decode_file_path(self._room_id)
        path = self._file_id_manager.get_path(file_id)

        data = {"level": level.value, "room": self._room_id, "path": path}
        if action:
            data["action"] = action
        if msg:
            data["msg"] = msg

        self.event_logger.emit(schema_id=JUPYTER_COLLABORATION_EVENTS_URI, data=data)

    def _emit_awareness_event(self, username: str, action: str, msg: str | None = None) -> None:
        data = {"roomid": self._room_id, "username": username, "action": action}
        if msg:
            data["msg"] = msg

        self.event_logger.emit(schema_id=JUPYTER_COLLABORATION_AWARENESS_EVENTS_URI, data=data)

    async def _clean_room(self) -> None:
        """
        Async task for cleaning up the resources.

        When all the clients of a room leave, we setup a task to clean up the resources
        after a certain amount of time. We need to wait a few seconds to clean up the room
        because sometimes websockets unintentionally disconnect.

        During the clean up, we need to delete the room to free resources since the room
        contains a copy of the document. In addition, we remove the file if there is no rooms
        subscribed to it.
        """
        assert isinstance(self.room, DocumentRoom)

        if self._cleanup_delay is None:
            return

        await asyncio.sleep(self._cleanup_delay)

        async with self._room_lock(self._room_id):
            # Remove the room from the websocket server
            self.log.info("Deleting Y document from memory: %s", self._room_id)
            await self._websocket_server.delete_room(room=self.room)

            # Clean room
            del self.room
            self.log.info("Room %s deleted", self._room_id)
            self._emit(LogLevel.INFO, "clean", "Room deleted.")

            # Clean the file loader if there are not rooms using it
            _, _, file_id = decode_file_path(self._room_id)
            file = self._file_loaders[file_id]
            if file.number_of_subscriptions == 0:
                self.log.info("Deleting file %s", file.path)
                await self._file_loaders.remove(file_id)
                self._emit(LogLevel.INFO, "clean", "Loader deleted.")
            del self._room_locks[self._room_id]

    def _on_global_awareness_event(self, topic: str, changes: tuple[dict[str, Any], Any]) -> None:
        """
        Update the users when the global awareness changes.

            Parameters:
                topic (str): `"update"` or `"change"` (`"change"` is triggered only if the states are modified).
                changes (tuple[dict[str, Any], Any]): The changes and the origin of the changes.
        """
        if topic != "change":
            return
        added_users = changes[0]["added"]
        removed_users = changes[0]["removed"]
        for user in added_users:
            u = self.room.awareness.states[user]
            if "user" in u:
                name = u["user"]["name"]
                self._websocket_server.connected_users[user] = name
                self.log.debug("Y user joined: %s", name)
        for user in removed_users:
            if user in self._websocket_server.connected_users:
                name = self._websocket_server.connected_users.pop(user)
                self.log.debug("Y user left: %s", name)

    def check_origin(self, origin):
        """
        Check origin
        """
        return True


class DocSessionHandler(APIHandler):
    """
    Jupyter Server's handler to retrieve the document's session.
    """

    auth_resource = "contents"

    @web.authenticated
    @authorized
    async def put(self, path):
        """
        Creates a new session for a given document or returns an existing one.
        """
        body = json.loads(self.request.body)
        format = body["format"]
        content_type = body["type"]
        file_id_manager = self.settings["file_id_manager"]

        idx = file_id_manager.get_id(path)
        if idx is not None:
            # index already exists
            self.log.info("Request for Y document '%s' with room ID: %s", path, idx)
            data = json.dumps(
                {
                    "format": format,
                    "type": content_type,
                    "fileId": idx,
                    "sessionId": SERVER_SESSION,
                }
            )
            self.set_status(200)
            return self.finish(data)

        # try indexing
        idx = file_id_manager.index(path)
        if idx is None:
            # file does not exists
            raise web.HTTPError(404, f"File {path!r} does not exist")

        # index successfully created
        self.log.info("Request for Y document '%s' with room ID: %s", path, idx)
        data = json.dumps(
            {
                "format": format,
                "type": content_type,
                "fileId": idx,
                "sessionId": SERVER_SESSION,
            }
        )
        self.set_status(201)
        return self.finish(data)


class TimelineHandler(APIHandler):
    def initialize(
        self, ystore_class: type[BaseYStore], ywebsocket_server: JupyterWebsocketServer
    ) -> None:
        self.ystore_class = ystore_class
        self.ywebsocket_server = ywebsocket_server

    async def get(self, path: str) -> None:
        idx = uuid4().hex
        file_id_manager = self.settings["file_id_manager"]
        file_id = file_id_manager.get_id(path)

        format = str(self.request.query_arguments.get("format")[0].decode("utf-8"))
        content_type = str(self.request.query_arguments.get("type")[0].decode("utf-8"))
        encoded_path = encode_file_path(format, content_type, file_id)
        try:
            room_id = room_id_from_encoded_path(encoded_path)
            room: YRoom = await self.ywebsocket_server.get_room(room_id)
            fork_ydoc: Doc = Doc()

            ydoc_factory = YDOCS.get(content_type)
            if ydoc_factory is None:
                self.set_status(404)
                self.finish(
                    {
                        "code": 404,
                        "error": f"No document factory found for content type: {content_type}",
                    }
                )
                return

            FORK_DOCUMENTS[idx] = ydoc_factory(fork_ydoc)
            undo_manager: UndoManager = FORK_DOCUMENTS[idx].undo_manager

            ystore = room.ystore
            assert ystore
            updates_and_timestamps = [(item[0], item[-1]) async for item in ystore.read()]

            result_timestamps = []

            for update, timestamp in updates_and_timestamps:
                undo_stack_len = len(undo_manager.undo_stack)
                fork_ydoc.apply_update(update)
                if len(undo_manager.undo_stack) > undo_stack_len:
                    result_timestamps.append(timestamp)

            fork_room = YRoom(ydoc=fork_ydoc)
            self.ywebsocket_server.add_room(idx, fork_room)

            data = {
                "roomId": room_id,
                "timestamps": result_timestamps,
                "forkRoom": idx,
                "sessionId": SERVER_SESSION,
            }
            self.set_status(200)
            self.finish(json.dumps(data))

        except RoomNotFound:
            self.set_status(404)
            self.finish({"code": 404, "error": "Room not found"})


class UndoRedoHandler(APIHandler):
    def initialize(self, ywebsocket_server: JupyterWebsocketServer) -> None:
        self._websocket_server = ywebsocket_server

    async def put(self, room_id):
        try:
            action = str(self.request.query_arguments.get("action")[0].decode("utf-8"))
            steps = int(self.request.query_arguments.get("steps")[0].decode("utf-8"))
            fork_room_id = str(self.request.query_arguments.get("forkRoom")[0].decode("utf-8"))

            fork_document = FORK_DOCUMENTS.get(fork_room_id)
            if not fork_document:
                self.set_status(404)
                self.log.warning(f"Fork document not found for room ID {fork_room_id}")
                return self.finish({"code": 404, "error": "Fork document not found"})

            undo_manager = fork_document.undo_manager

            if action == "undo":
                if undo_manager.can_undo():
                    await self._perform_undo_or_redo(undo_manager, "undo", steps)
                    self.set_status(200)
                    self.log.info(f"Undo operation performed in room {room_id} for {steps} steps")
                    return self.finish({"status": "undone"})
                else:
                    self.log.info(f"No more undo operations available in room {room_id}")
                    return self.finish({"error": "No more undo operations available"})
            elif action == "redo":
                if undo_manager.can_redo():
                    await self._perform_undo_or_redo(undo_manager, "redo", steps)
                    self.set_status(200)
                    self.log.info(f"Redo operation performed in room {room_id} for {steps} steps")
                    return self.finish({"status": "redone"})
                else:
                    self.log.info(f"No more redo operations available in room {room_id}")
                    return self.finish({"error": "No more redo operations available"})
            elif action == "restore":
                try:
                    # Cleanup undo manager after restoration
                    await self._cleanup_undo_manager(fork_room_id)

                    self.set_status(200)
                    self.log.info(f"Document in room {room_id} restored successfully")
                    return self.finish({"code": 200, "status": "Document restored successfully"})
                except Exception as e:
                    self.log.error(f"Error during document restore in room {room_id}: {str(e)}")
                    self.set_status(500)
                    return self.finish(
                        {"code": 500, "error": "Internal server error", "message": str(e)}
                    )

        except Exception as e:
            self.log.error(f"Error during undo/redo/restore operation in room {room_id}: {str(e)}")

    async def _perform_undo_or_redo(
        self, undo_manager: UndoManager, action: str, steps: int
    ) -> None:
        for _ in range(steps):
            if action == "undo" and len(undo_manager.undo_stack) > 1:
                undo_manager.undo()

            elif action == "redo" and undo_manager.can_redo():
                undo_manager.redo()
            else:
                break

    async def _cleanup_undo_manager(self, room_id: str) -> None:
        if room_id in FORK_DOCUMENTS:
            del FORK_DOCUMENTS[room_id]
            self.log.info(f"Fork Document for {room_id} has been removed.")


class DocForkHandler(APIHandler):
    """
    Jupyter Server handler to:
    - create a fork of a root document (optionally synchronizing with the root document),
    - delete a fork of a root document (optionally merging back in the root document).
    - get fork IDs of a root document.
    """

    auth_resource = "contents"

    def initialize(
        self,
        ywebsocket_server: JupyterWebsocketServer,
    ) -> None:
        self._websocket_server = ywebsocket_server

    @web.authenticated
    @authorized
    async def get(self, root_roomid):
        """
        Returns a dictionary of fork room ID to fork room information for the given root room ID.
        """
        self.write(
            {
                fork_roomid: fork_info
                for fork_roomid, fork_info in FORK_ROOMS.items()
                if fork_info["root_roomid"] == root_roomid
            }
        )

    @web.authenticated
    @authorized
    async def put(self, root_roomid):
        """
        Creates a fork of a root document and returns its ID.
        Optionally keeps the fork in sync with the root.
        """
        fork_roomid = uuid4().hex
        try:
            root_room = await self._websocket_server.get_room(root_roomid)
        except RoomNotFound:
            self.set_status(404)
            return self.finish({"code": 404, "error": "Root room not found"})

        update = root_room.ydoc.get_update()
        fork_ydoc: Doc = Doc()
        fork_ydoc.apply_update(update)
        model = self.get_json_body()
        synchronize = model.get("synchronize", False)
        if synchronize:
            root_room.ydoc.observe(lambda event: fork_ydoc.apply_update(event.update))
        FORK_ROOMS[fork_roomid] = fork_info = {
            "root_roomid": root_roomid,
            "synchronize": synchronize,
            "title": model.get("title", ""),
            "description": model.get("description", ""),
        }
        fork_room = YRoom(ydoc=fork_ydoc)
        self._websocket_server.rooms[fork_roomid] = fork_room
        await self._websocket_server.start_room(fork_room)
        self._emit_fork_event(self.current_user.username, fork_roomid, fork_info, "create")
        data = json.dumps(
            {
                "sessionId": SERVER_SESSION,
                "fork_roomid": fork_roomid,
                "fork_info": fork_info,
            }
        )
        self.set_status(201)
        return self.finish(data)

    @web.authenticated
    @authorized
    async def delete(self, fork_roomid):
        """
        Deletes a forked document, and optionally merges it back in the root document.
        """
        fork_info = FORK_ROOMS.get(fork_roomid, None)
        if fork_info is None:
            self.set_status(404)
            return self.finish({"code": 404, "error": "Fork room not found"})
        root_roomid = fork_info["root_roomid"]
        del FORK_ROOMS[fork_roomid]
        if self.get_query_argument("merge") == "true":
            root_room = await self._websocket_server.get_room(root_roomid)
            root_ydoc = root_room.ydoc
            fork_room = await self._websocket_server.get_room(fork_roomid)
            fork_ydoc = fork_room.ydoc
            fork_update = fork_ydoc.get_update()
            root_ydoc.apply_update(fork_update)
        await self._websocket_server.delete_room(name=fork_roomid)
        self._emit_fork_event(self.current_user.username, fork_roomid, fork_info, "delete")
        self.set_status(200)

    def _emit_fork_event(
        self, username: str, fork_roomid: str, fork_info: dict[str, str], action: str
    ) -> None:
        data = {
            "username": username,
            "fork_roomid": fork_roomid,
            "fork_info": fork_info,
            "action": action,
        }
        self.event_logger.emit(schema_id=JUPYTER_COLLABORATION_FORK_EVENTS_URI, data=data)
