from collections import defaultdict
import logging
from contextvars import Context, ContextVar, copy_context
from asyncio.subprocess import create_subprocess_exec
import subprocess

task_id_var = ContextVar("task_id_var")
task_id_var.set(None)


class AsyncTaskLogger(logging.Logger):
    """Custom logger for async tasks that ensures logs grouped and ordered by task.

    Allows for logging in concurrent async tasks without the logs getting mixed by
    specifying a task id in the `task_id_var` contextvar.

    Saves logs created during running of an async task separated per task and
    only writes the saved logs when calling `AsyncTaskLogger.flush_async_logs()`
    When flushed, the logs are written task by task and ordered by the order of the
    task ids.

    If no `task_id_var` is set this acts as a normal `logging.Logger`.

    Note that this logger does not automatically flush the logs when an exception is
    raised. The user is responsible for that.
    """

    logs_mapping: defaultdict[int, list]
    level_method_name_mapping = {
        logging.INFO: "info",
        logging.DEBUG: "debug",
    }

    def __init__(self, name: str, level: int | str = 0) -> None:
        super().__init__(name, level)

        self.logs_mapping = defaultdict(list)

    def log_async(self, level, message, *args, **kwargs):
        """Either log message as eagerly as usual or save it to be logged
        later when flushed.
        """
        task_id = task_id_var.get()

        # if no context with a task_id is set use normal logging
        if task_id is None:
            log_method = getattr(super(), self.level_method_name_mapping[level])
            return log_method(message, *args, **kwargs)

        self.logs_mapping[task_id].append((level, message, args, kwargs))

    def info(self, message, *args, **kwargs):
        self.log_async(logging.INFO, message, *args, **kwargs)

    def debug(self, message, *args, **kwargs):
        self.log_async(logging.DEBUG, message, *args, **kwargs)

    def flush_async_logs(self, prepend_task_id: bool = False):
        """Flush logs by logging the saved messages as usual but in
        the correct order.

        The task id can optionally be prepended to the log message when
        `prepend_task_id` is True. This can especially be useful when tasks
        don't finish due to an exception.
        """
        ordered_task_ids = sorted(self.logs_mapping.keys())

        for task_id in ordered_task_ids:
            for log_level, message, args, kwargs in self.logs_mapping[task_id]:
                if log_level not in [logging.INFO, logging.DEBUG]:
                    raise NotImplementedError

                log_method = getattr(super(), self.level_method_name_mapping[log_level])
                log_method(
                    message if not prepend_task_id else f"Task {task_id} -- {message}",
                    *args,
                    **kwargs,
                )

        self.logs_mapping = defaultdict(list)


def create_task_id_context(task_id: int) -> Context:
    """Creates a context with `task_id_var` set to the given task id."""
    task_id_var.set(task_id)
    new_context = copy_context()
    task_id_var.set(None)
    return new_context


def get_async_logger(name: str) -> AsyncTaskLogger:
    logging.setLoggerClass(AsyncTaskLogger)
    return logging.getLogger(name)


async def run_process_async(command: list[str]) -> tuple[int, str, str]:
    """Run a command as a subprocess asyncronously.

    Returns a tuple of the return code, stdout and stderr.
    """
    process = await create_subprocess_exec(
        *command,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    stdout, stderr = await process.communicate()

    if process.returncode is None:
        raise RuntimeError("Subprocess has not yet terminated, but it should have.")

    return process.returncode, stdout.decode(), stderr.decode()
