import asyncio
import traceback
from abc import ABC, abstractmethod
from collections.abc import Mapping
from json import JSONDecodeError
from typing import Any, ClassVar, Self, cast

import docker
import docker.errors
from asgiref.sync import sync_to_async
from channels.generic.websocket import AsyncWebsocketConsumer, DenyConnection, json
from django.contrib.auth.models import User
from docker.models.containers import Container
from typing_extensions import override

from vuln_demo.settings import DEBUG
from vuln_demo.util.auth import cast_user
from vuln_demo.util.log import get_logger

logger = get_logger(__name__)


class ContainerUnhealthyError(Exception):
    """Raised when the container healthcheck fails."""


class DockerConsumer(ABC, AsyncWebsocketConsumer):
    """
    Base class for all Docker consumers.

    The communication between the client and the server is done through a WebSocket connection.
    Each message is expected to be a JSON object, with a `type` key, and `message` key.

    Following message types are used:
        - `docker.init`: Sent when the Docker container is being created.
        - `docker.ready`: Sent when the Docker container is ready to be used.
        - `docker.waiting`: Sent when the Docker container is waiting for some action to complete.
        - `consumer.waiting`: Sent when the consumer is waiting for another connection to stop.
        - `docker.error`: Sent when an error occurs during the Docker container creation.
        - `docker.stopped`: Sent when the Docker container has been stopped.
        - `consumer.error`: Sent when an error occurs during the consumer execution.

    The client is expected to handle receiving all of these types.

    If the docker container was started successfully, it will remain running until the
    connection is closed, or until the container exits normally. Child classes can override
    this behavior by calling `self.docker_stop()` to stop the container without also having
    to close the connection.

    Only one connection is allowed per user. If a user tries to connect while already connected,
    from elsewhere, the new connection will be denied. This prevents a DoS attack where the user
    opens multiple connections to the same docker consumer, attempting to start multiple containers.

    Note that the connection
    """

    CONSUMER_ID: ClassVar[str]
    DOCKER_IMAGE: ClassVar[str]
    DOCKER_COMMAND: ClassVar[str | None] = None
    DOCKER_ENVIRONMENT: ClassVar[Mapping[str, str] | None] = None
    DOCKER_HEALTHCHECK: ClassVar[Mapping[str, Any] | None] = None

    DOCKER_STOP_TIMEOUT: ClassVar[int] = 2
    DOCKER_HEALTHCHECK_INTERVAL: ClassVar[float] = 0.5

    ACTIVE_CONSUMERS: ClassVar[dict[User, Self]] = {}

    container: Container

    def __init__(self) -> None:
        if not hasattr(self, "CONSUMER_ID"):
            raise ValueError("CONSUMER_ID must be set on the class.")
        if not hasattr(self, "DOCKER_IMAGE"):
            raise ValueError("CONSUMER_IMAGE must be set on the class.")

        super().__init__()
        self.client = docker.from_env()
        self.docker_prepared = False

    @property
    def user(self) -> User:
        """Return the user associated with this connection."""
        user = self.scope["user"]
        assert cast_user(user, cast=True)
        return user

    @property
    def container_name(self) -> str:
        """Return the name of the associated Docker container."""
        return f"vuln_demo_{self.CONSUMER_ID}_{self.user.username}"

    async def send_message(self, message: str, msg_type: str = "consumer.message", **kwargs: str) -> None:
        """Send a message to the client."""
        await self.send(text_data=json.dumps({"type": msg_type, "message": message, **kwargs}))

    @override
    async def connect(self) -> None:
        """Called when a WebSocket connection is opened."""
        await self.accept()

        # Ensure the user is authenticated
        if not self.user.is_authenticated:
            logger.debug(f"Dropping unauthenticated connection to Docker consumer {self.CONSUMER_ID}.")
            await self.send_message("Not authenticated.", "consumer.error")
            await self.close()
            return

        # Prevent multiple connections from the same user.
        if self.user in self.ACTIVE_CONSUMERS:
            logger.info(
                f"User {self.user.username!r} requested a new connection to docker consumer {self.CONSUMER_ID} "
                "while already connected. Dropping their existing connection in favor of (this) new one."
            )
            other_conn = self.ACTIVE_CONSUMERS[self.user]
            await self.send_message("Waiting for other connection to stop...", "consumer.waiting")
            await other_conn.send_message("Another connection has taken over.", "consumer.error")
            await other_conn.docker_stop()
            await other_conn.close()
        self.ACTIVE_CONSUMERS[self.user] = self

        # In case another container with the same name still wasn't removed
        # we need to first wait for it to get removed, to avoid name conflicts.
        await self.wait_until_removed()

        logger.info(
            f"Docker consumer {self.CONSUMER_ID} connected user {self.user.username!r}. " "Starting Docker container."
        )

        await self.send_message("Starting Docker container...", "docker.init")
        try:
            self.container = await self.docker_run_container()
        except DenyConnection as e:
            raise e from e
        except Exception as e:
            msg = "Error starting Docker container. Report this to DEMO portal maintainers!"
            if DEBUG:
                msg += f"<pre>{e!r}</pre>"
                msg += f"<pre>{traceback.format_exc()}</pre>"
            else:
                msg += " This issue can be found in the logs."

            logger.error(
                "Exception occured while starting the container in docker consumer "
                f"{self.CONSUMER_ID} while preparing new container for {self.user.username!r} ",
                e,
            )

            await self.send_message(msg, "docker.error")
            await self.close()
            raise e from e

        logger.info(
            f"Docker consumer {self.CONSUMER_ID} finished starting new container for {self.user.username!r}. "
            f"Container id: {self.container.id!r}"
        )

        await self.send_message("Preparing your Docker container...", "docker.waiting")
        try:
            prepare_success = await self.docker_prepare()
        except DenyConnection as e:
            raise e from e
        except Exception as e:
            msg = "Error preparing the Docker container. Report this to DEMO portal maintainers!"
            if DEBUG:
                msg += f"<pre>{e!r}</pre>"
                msg += f"<pre>{traceback.format_exc()}</pre>"
                try:
                    logs = await sync_to_async(self.container.logs)()
                    msg += f"<pre>{logs.decode('utf-8')}</pre>"
                except Exception as e2:
                    msg += f"<pre>Failed to get logs: {e2!r}</pre>"
            else:
                msg += " This issue can be found in the logs."

            logger.error(
                "Exception occured during container preparation in docker consumer "
                f"{self.CONSUMER_ID} while preparing new container for {self.user.username!r} "
                f"(Container id: {self.container.id!r})",
                e,
            )

            await self.send_message(msg, "docker.error")
            await self.close()
            raise e from e

        if not prepare_success:
            await self.close()

        await self.send_message("Docker container is ready!", "docker.ready")
        self.docker_prepared = True

    @override
    async def disconnect(self, code: int) -> None:
        """
        Called when the WebSocket connection is closed (either by the client or by us).

        This function is expected to perform any clean up of the resources that were allocated.
        """
        if self.user.is_authenticated:
            logger.info(f"Docker consumer {self.CONSUMER_ID} disconnecting user {self.user.username!r}")

        await self.docker_stop()

    async def docker_run_container(self) -> Container:
        """
        Start the Docker container.

        This function is expected to simply execute `self.client.containers.run()` function,
        starting the container image specified by `self.DOCKER_IMAGE`. The container should
        be started in detached mode.

        The run function can be passed special arguments, to configure volumes, exposed ports,
        health-check config, or other options.

        Once this function exits, the container is expected to be running and healthy.
        """
        extra_kwargs = {}
        if self.DOCKER_ENVIRONMENT is not None:
            extra_kwargs["environment"] = self.DOCKER_ENVIRONMENT
        if self.DOCKER_HEALTHCHECK is not None:
            extra_kwargs["healthcheck"] = self.DOCKER_HEALTHCHECK

        try:
            container = await sync_to_async(self.client.containers.run)(
                self.DOCKER_IMAGE,
                self.DOCKER_COMMAND,
                name=self.container_name,
                hostname=self.container_name,
                detach=True,
                auto_remove=True,
                network_disabled=True,
                mem_limit="512m",
                cpu_shares=128,  # proportional weight for cpu use by the container, default: 1024
                cpu_period=100_000,  # 0.1  seconds
                cpu_quota=50_000,  # 0.5 seconds (50% of cpu_period)
                **extra_kwargs,
            )
        except docker.errors.APIError as exc:
            if exc.status_code == 409:  # Container already exists
                logger.debug(
                    f"Docker consumer {self.CONSUMER_ID} failed with container already exis for user "
                    f"{self.user.username}: {exc!r}",
                )
                await self.wait_until_removed()
                if self.ACTIVE_CONSUMERS[self.user] != self:
                    await self.close()
                    await self.send_message("Another connection has taken over.", "consumer.error")
                    raise DenyConnection

                return await self.docker_run_container()
            raise exc from exc

        return cast(Container, container)

    async def docker_stop(self) -> None:
        """
        Stop the Docker container.

        Once this function ends, the Docker container is expected to be stopped.
        It is crucial that this function performs all necessary cleanup (such as
        removing the subvolumes) to avoid any resource leaks.
        """
        if not hasattr(self, "container"):
            return

        logger.info(
            f"Docker consumer {self.CONSUMER_ID} stopping container {self.container.id!r} "
            f"for user {self.user.username!r}."
        )

        # If the contianer isn't stopped within the given timeout, this will send
        # SIGKILL to stop the container.
        try:
            await sync_to_async(self.container.stop)(timeout=self.DOCKER_STOP_TIMEOUT)
        except docker.errors.NotFound as e:
            logger.debug(f"Skipping stopping non-existent container: {e!r}")
        await self.send_message("Docker container has been stopped.", "docker.stopped")

    async def wait_until_removed(self) -> None:
        """
        Wait until he Docker container has stopped.

        Useful when waiting for another thread to stop the container with the same name.
        This can occur when the user connects from multiple locations, and the old connection
        is being closed in favor of the new one.
        """
        try:
            container = await sync_to_async(self.client.containers.get)(self.container_name)
        except docker.errors.NotFound:
            return
        container = cast(Container, container)

        await self.send_message("Waiting for previous container to stop...", "docker.waiting")
        logger.info(f"Docker consumer {self.CONSUMER_ID} waiting for container {container.name!r} to stop.")
        try:
            res = container.wait(condition="removed")
        except docker.errors.NotFound:
            logger.debug(
                f"Docker consumer {self.CONSUMER_ID} Skipping waiting for non-existent "
                f"container {container.name!r}"
            )
            return
        logger.info(f"Docker consumer {self.CONSUMER_ID} finished waiting for container {container.name!r}: {res!r}")

    async def docker_prepare(self) -> bool:
        """
        Called when the Docker container is ready.

        This function can be used to interact with the container immediately after
        it has became ready. This can include sending any initial messages to the
        client, or performing any initial setup.

        If this function raises an exception, the container will be stopped and removed,
        and the connection will be closed.

        NOTE: For one-shot containers, make sure to call `self.close()`, to end the websocket
        connection once the container exits (`self.container.wait()`), with which the docker
        container will be properly cleaned up. If you need the connection to remain open,
        you can also call `self.docker_stop()` to perform this cleanup manually as soon as
        the container exits.

        By default, this function will do nothing for containers without a healthcheck.

        If the container has a healthcheck, this function will wait until the container
        becomes healthy, or unhealthy (the healthcheck times out). This will use active waiting
        with a blocking :func:`time.sleep` function, sleeping for :attr:`DOCKER_HEALTHCHECK_INTERVAL`
        between each iteration.

        Note that there is no time limit enforced here for this healthcheck, you should instead configure
        the container's healthcheck to have a timeout, which will cause the container status to become
        'unhealthy', and this function will raise an :exc:`ContainerUnhealthyError`.
        """
        logs = (await sync_to_async(self.container.logs)()).decode("utf-8")
        await sync_to_async(self.container.reload)()
        if self.container.health == "unknown":
            return True

        while self.container.health == "starting":
            logs = (await sync_to_async(self.container.logs)()).decode("utf-8")
            await asyncio.sleep(self.DOCKER_HEALTHCHECK_INTERVAL)
            await sync_to_async(self.container.reload)()

        if self.container.health != "healthy":
            await self.send_message(
                f"The docker container didn't pass it's configured healthcheck. <pre>{logs}</pre>", "docker.error"
            )
            raise ContainerUnhealthyError(
                f"Container didn't pass the configured healthcheck: {self.container.health!r}"
            )
        return True

    @override
    async def receive(self, text_data: str | None = None, bytes_data: bytes | None = None) -> None:
        if bytes_data:
            await self.send_message("No binary data allowed.", "consumer.error")
            await self.close()
            return
        if text_data is None:
            await self.send_message("Got empty message.", "consumer.error")
            await self.close()
            return

        try:
            data = json.loads(text_data)
        except JSONDecodeError:
            await self.send_message("Invalid JSON data.", "consumer.error")
            return

        if not self.docker_prepared:
            await self.send_message("Docker container isn't ready yet", "consumer.error")
            return

        logger.info(f"Docker consumer {self.CONSUMER_ID} received message from user {self.user.username!r}: {data!r}")
        await self.process_message(data)

    @abstractmethod
    async def process_message(self, data: Mapping[str, Any]) -> None:
        """
        Process a message received from the client.

        This function takes the message `data`, which is the JSON object sent
        by the client. If the client follows our specification, the message
        should contain a `type` key, and a `message` key.
        """
