""".. Ignore pydocstyle D400.

==================
List Docker images
==================

"""

import functools
import logging
import operator
import shlex
import subprocess
import threading
import time

import yaml
from django.conf import settings
from django.core.management.base import BaseCommand, CommandError

from resolwe.flow.models import Process

logger = logging.getLogger(__name__)

# Global list of images already pulled when running in the same process. This
# is used to avoid pulling images more than once when executed from tests.
PULLED_IMAGES = set()
PULLED_IMAGES_LOCK = threading.Lock()


class Command(BaseCommand):
    """List Docker images used by processes.  Optionally also pull them."""

    help = "List Docker images used by processes in either plain text or YAML (for Ansible)"

    def add_arguments(self, parser):
        """Add an argument to specify output format."""
        parser.add_argument(
            "--format",
            dest="format",
            default="plain",
            help="Set output format ('plain' [default] or 'yaml')",
        )
        parser.add_argument(
            "--pull",
            dest="pull",
            default=False,
            action="store_true",
            help="Pull all images with Docker",
        )
        parser.add_argument(
            "--ignore-pull-errors",
            dest="ignore_pull_errors",
            default=getattr(settings, "FLOW_DOCKER_IGNORE_PULL_ERRORS", False),
            action="store_true",
            help="Don't fail whenever a Docker image can't be pulled",
        )

    def _map_image(self, docker_image: str, mapper: dict[str, str]) -> str:
        """Transform the given Docker image according to the mapper.

        If image starts with the prefix in the mapper keys, it is replaced with the
        corresponding value. If no prefix matches, the image is returned unchanged.
        """
        for prefix, replacement in mapper.items():
            if docker_image.startswith(prefix):
                return replacement + docker_image[len(prefix) :]
        return docker_image

    def _map_images(self, docker_images: set[str]) -> set[str]:
        """Map the docker images in the given set according to the settings."""
        mapper = getattr(settings, "FLOW_CONTAINER_IMAGE_MAP", {})
        return {self._map_image(image, mapper) for image in docker_images}

    def handle(self, *args, **options):
        """Handle command list_docker_images."""
        verbosity = int(options.get("verbosity"))

        # Check that the specified output format is valid
        if options["format"] != "plain" and options["format"] != "yaml":
            raise CommandError("Unknown output format: %s" % options["format"])

        # Gather only unique latest custom Docker requirements that the processes are using
        # The 'image' field is optional, so be careful about that as well
        unique_docker_images = set(
            p.requirements["executor"]["docker"]["image"]
            for p in Process.objects.filter(is_active=True)
            .only("requirements")
            .filter(requirements__icontains="docker")
            if "image" in p.requirements.get("executor", {}).get("docker", {})
        )

        # Add the default image.
        unique_docker_images.add(
            getattr(
                settings,
                "FLOW_DOCKER_DEFAULT_PROCESSING_CONTAINER_IMAGE",
                "public.ecr.aws/s4q6j6e8/resolwe/base:ubuntu-20.04",
            )
        )

        # Map the Docker images according to the settings.
        unique_docker_images = self._map_images(unique_docker_images)

        # Pull images if requested or just output the list in specified format
        if options["pull"]:
            # Remove set of already pulled images.
            with PULLED_IMAGES_LOCK:
                unique_docker_images.difference_update(PULLED_IMAGES)

            # Get the desired 'docker' command from settings or use the default
            docker = getattr(settings, "FLOW_DOCKER_COMMAND", "docker")

            # Pull each image
            for img in unique_docker_images:
                max_retries = 10
                for _ in range(max_retries):
                    ret = subprocess.run(
                        shlex.split("{} pull {}".format(docker, img)),
                        stdout=subprocess.PIPE,
                        stderr=subprocess.STDOUT,
                    )
                    if ret.returncode == 0:
                        break
                    time.sleep(1)

                # Update set of pulled images.
                with PULLED_IMAGES_LOCK:
                    PULLED_IMAGES.add(img)

                if ret.returncode != 0:
                    errmsg = "Failed to pull Docker image '{}': '{}'.".format(
                        img, ret.stdout.decode()
                    )

                    if not options["ignore_pull_errors"]:
                        # Print error and stop execution
                        raise CommandError(errmsg)
                    else:
                        # Print error, but keep going
                        logger.error(errmsg)
                        if verbosity > 0:
                            self.stderr.write(errmsg)
                else:
                    msg = "Docker image '{}' pulled successfully!".format(img)
                    logger.info(msg)
                    if verbosity > 0:
                        self.stdout.write(msg)
        else:
            # Sort the set of unique Docker images for nicer output.
            unique_docker_images = sorted(unique_docker_images)

            # Convert the set of unique Docker images into a list of dicts for easier output
            imgs = [
                dict(name=s[0], tag=s[1] if len(s) == 2 else "latest")
                for s in (img.split(":") for img in unique_docker_images)
            ]

            # Output in YAML or plaintext (one image per line), as requested
            if options["format"] == "yaml":
                out = yaml.safe_dump(imgs, default_flow_style=True, default_style="'")
            else:
                out = functools.reduce(
                    operator.add, ("{name}:{tag}\n".format(**i) for i in imgs), ""
                )

            self.stdout.write(out, ending="")
