"""Helper library which performs validation of the submission."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from builtins import int  # long in python 2

import csv
import json
import logging
import os
import re
import subprocess

import numpy as np
from six import iteritems
from PIL import Image


EXTRACT_COMMAND = {
    ".zip": ["unzip", "${src}", "-d", "${dst}"],
    ".tar": ["tar", "xvf", "${src}", "-C", "${dst}"],
    ".tar.gz": ["tar", "xvzf", "${src}", "-C", "${dst}"],
}

ALLOWED_SUBMISSION_TYPES = ["attack", "targeted_attack", "defense"]

REQUIRED_METADATA_JSON_FIELDS = ["entry_point", "container", "container_gpu", "type"]

CMD_VARIABLE_RE = re.compile("^\\$\\{(\\w+)\\}$")

BATCH_SIZE = 100
IMAGE_NAME_PATTERN = "IMG{0:04}.png"

ALLOWED_EPS = [4, 8, 12, 16]

MAX_SUBMISSION_SIZE_ZIPPED = 8 * 1024 * 1024 * 1024  # 8 GiB
MAX_SUBMISSION_SIZE_UNPACKED = 16 * 1024 * 1024 * 1024  # 16 GiB
MAX_DOCKER_IMAGE_SIZE = 8 * 1024 * 1024 * 1024  # 8 GiB


def get_extract_command_template(filename):
    """Returns extraction command based on the filename extension."""
    for k, v in iteritems(EXTRACT_COMMAND):
        if filename.endswith(k):
            return v
    return None


def shell_call(command, **kwargs):
    """Calls shell command with parameter substitution.

    Args:
      command: command to run as a list of tokens
      **kwargs: dirctionary with substitutions

    Returns:
      whether command was successful, i.e. returned 0 status code

    Example of usage:
      shell_call(['cp', '${A}', '${B}'], A='src_file', B='dst_file')
    will call shell command:
      cp src_file dst_file
    """
    command = list(command)
    for i in range(len(command)):
        m = CMD_VARIABLE_RE.match(command[i])
        if m:
            var_id = m.group(1)
            if var_id in kwargs:
                command[i] = kwargs[var_id]
    return subprocess.call(command) == 0


def make_directory_writable(dirname):
    """Makes directory readable and writable by everybody.

    Args:
      dirname: name of the directory

    Returns:
      True if operation was successfull

    If you run something inside Docker container and it writes files, then
    these files will be written as root user with restricted permissions.
    So to be able to read/modify these files outside of Docker you have to change
    permissions to be world readable and writable.
    """
    retval = shell_call(
        [
            "docker",
            "run",
            "-v",
            "{0}:/output_dir".format(dirname),
            "busybox:1.27.2",
            "chmod",
            "-R",
            "a+rwx",
            "/output_dir",
        ]
    )
    if not retval:
        logging.error("Failed to change permissions on directory: %s", dirname)
    return retval


def load_defense_output(filename):
    """Loads output of defense from given file."""
    result = {}
    with open(filename) as f:
        for row in csv.reader(f):
            try:
                image_filename = row[0]
                if not image_filename.endswith(".png"):
                    image_filename += ".png"
                label = int(row[1])
            except (IndexError, ValueError):
                continue
            result[image_filename] = label
    return result


class SubmissionValidator(object):
    """Class which performs validation of the submission."""

    def __init__(self, temp_dir, use_gpu):
        """Initializes instance of SubmissionValidator.

        Args:
          temp_dir: temporary working directory
          use_gpu: whether to use GPU
        """
        self._temp_dir = temp_dir
        self._use_gpu = use_gpu
        self._tmp_extracted_dir = os.path.join(self._temp_dir, "tmp_extracted")
        self._extracted_submission_dir = os.path.join(self._temp_dir, "extracted")
        self._sample_input_dir = os.path.join(self._temp_dir, "input")
        self._sample_output_dir = os.path.join(self._temp_dir, "output")

    def _prepare_temp_dir(self):
        """Cleans up and prepare temporary directory."""
        shell_call(["rm", "-rf", os.path.join(self._temp_dir, "*")])
        # NOTE: we do not create self._extracted_submission_dir
        # this is intentional because self._tmp_extracted_dir or it's subdir
        # will be renames into self._extracted_submission_dir
        os.mkdir(self._tmp_extracted_dir)
        os.mkdir(self._sample_input_dir)
        os.mkdir(self._sample_output_dir)
        # make output dir world writable
        shell_call(["chmod", "a+rwX", "-R", self._sample_output_dir])

    def _extract_submission(self, filename):
        """Extracts submission and moves it into self._extracted_submission_dir."""
        # verify filesize
        file_size = os.path.getsize(filename)
        if file_size > MAX_SUBMISSION_SIZE_ZIPPED:
            logging.error(
                "Submission archive size %d is exceeding limit %d",
                file_size,
                MAX_SUBMISSION_SIZE_ZIPPED,
            )
            return False
        # determime archive type
        exctract_command_tmpl = get_extract_command_template(filename)
        if not exctract_command_tmpl:
            logging.error(
                "Input file has to be zip, tar or tar.gz archive; however " "found: %s",
                filename,
            )
            return False
        # extract archive
        submission_dir = os.path.dirname(filename)
        submission_basename = os.path.basename(filename)
        logging.info("Extracting archive %s", filename)
        retval = shell_call(
            [
                "docker",
                "run",
                "--network=none",
                "-v",
                "{0}:/input_dir".format(submission_dir),
                "-v",
                "{0}:/output_dir".format(self._tmp_extracted_dir),
                "busybox:1.27.2",
            ]
            + exctract_command_tmpl,
            src=os.path.join("/input_dir", submission_basename),
            dst="/output_dir",
        )
        if not retval:
            logging.error("Failed to extract submission from file %s", filename)
            return False
        if not make_directory_writable(self._tmp_extracted_dir):
            return False
        # find submission root
        root_dir = self._tmp_extracted_dir
        root_dir_content = [d for d in os.listdir(root_dir) if d != "__MACOSX"]
        if len(root_dir_content) == 1 and os.path.isdir(
            os.path.join(root_dir, root_dir_content[0])
        ):
            logging.info(
                'Looks like submission root is in subdirectory "%s" of ' "the archive",
                root_dir_content[0],
            )
            root_dir = os.path.join(root_dir, root_dir_content[0])
        # Move files to self._extracted_submission_dir.
        # At this point self._extracted_submission_dir does not exist,
        # so following command will simply rename root_dir into
        # self._extracted_submission_dir
        if not shell_call(["mv", root_dir, self._extracted_submission_dir]):
            logging.error("Can" "t move submission files from root directory")
            return False
        return True

    def _verify_submission_size(self):
        submission_size = 0
        for dirname, _, filenames in os.walk(self._extracted_submission_dir):
            for f in filenames:
                submission_size += os.path.getsize(os.path.join(dirname, f))
        logging.info("Unpacked submission size: %d", submission_size)
        if submission_size > MAX_SUBMISSION_SIZE_UNPACKED:
            logging.error(
                "Submission size exceeding limit %d", MAX_SUBMISSION_SIZE_UNPACKED
            )
        return submission_size <= MAX_SUBMISSION_SIZE_UNPACKED

    def _load_and_verify_metadata(self, submission_type):
        """Loads and verifies metadata.

        Args:
          submission_type: type of the submission

        Returns:
          dictionaty with metadata or None if metadata not found or invalid
        """
        metadata_filename = os.path.join(
            self._extracted_submission_dir, "metadata.json"
        )
        if not os.path.isfile(metadata_filename):
            logging.error("metadata.json not found")
            return None
        try:
            with open(metadata_filename, "r") as f:
                metadata = json.load(f)
        except IOError as e:
            logging.error("Failed to load metadata: %s", e)
            return None
        for field_name in REQUIRED_METADATA_JSON_FIELDS:
            if field_name not in metadata:
                logging.error("Field %s not found in metadata", field_name)
                return None
        # Verify submission type
        if submission_type != metadata["type"]:
            logging.error(
                'Invalid submission type in metadata, expected "%s", ' 'actual "%s"',
                submission_type,
                metadata["type"],
            )
            return None
        # Check submission entry point
        entry_point = metadata["entry_point"]
        if not os.path.isfile(
            os.path.join(self._extracted_submission_dir, entry_point)
        ):
            logging.error("Entry point not found: %s", entry_point)
            return None
        if not entry_point.endswith(".sh"):
            logging.warning(
                "Entry point is not an .sh script. "
                "This is not necessarily a problem, but if submission "
                "won"
                "t run double check entry point first: %s",
                entry_point,
            )
        # Metadata verified
        return metadata

    def _verify_docker_image_size(self, image_name):
        """Verifies size of Docker image.

        Args:
          image_name: name of the Docker image.

        Returns:
          True if image size is withing the limits, False otherwise.
        """
        shell_call(["docker", "pull", image_name])
        try:
            image_size = subprocess.check_output(
                ["docker", "inspect", "--format={{.Size}}", image_name]
            ).strip()
            image_size = int(image_size)
        except (ValueError, subprocess.CalledProcessError) as e:
            logging.error("Failed to determine docker image size: %s", e)
            return False
        logging.info("Size of docker image %s is %d", image_name, image_size)
        if image_size > MAX_DOCKER_IMAGE_SIZE:
            logging.error("Image size exceeds limit %d", MAX_DOCKER_IMAGE_SIZE)
        return image_size <= MAX_DOCKER_IMAGE_SIZE

    def _prepare_sample_data(self, submission_type):
        """Prepares sample data for the submission.

        Args:
          submission_type: type of the submission.
        """
        # write images
        images = np.random.randint(
            0, 256, size=[BATCH_SIZE, 299, 299, 3], dtype=np.uint8
        )
        for i in range(BATCH_SIZE):
            Image.fromarray(images[i, :, :, :]).save(
                os.path.join(self._sample_input_dir, IMAGE_NAME_PATTERN.format(i))
            )
        # write target class for targeted attacks
        if submission_type == "targeted_attack":
            target_classes = np.random.randint(1, 1001, size=[BATCH_SIZE])
            target_class_filename = os.path.join(
                self._sample_input_dir, "target_class.csv"
            )
            with open(target_class_filename, "w") as f:
                for i in range(BATCH_SIZE):
                    f.write(
                        (IMAGE_NAME_PATTERN + ",{1}\n").format(i, target_classes[i])
                    )

    def _run_submission(self, metadata):
        """Runs submission inside Docker container.

        Args:
          metadata: dictionary with submission metadata

        Returns:
          True if status code of Docker command was success (i.e. zero),
          False otherwise.
        """
        if self._use_gpu:
            docker_binary = "nvidia-docker"
            container_name = metadata["container_gpu"]
        else:
            docker_binary = "docker"
            container_name = metadata["container"]
        if metadata["type"] == "defense":
            cmd = [
                docker_binary,
                "run",
                "--network=none",
                "-m=24g",
                "-v",
                "{0}:/input_images:ro".format(self._sample_input_dir),
                "-v",
                "{0}:/output_data".format(self._sample_output_dir),
                "-v",
                "{0}:/code".format(self._extracted_submission_dir),
                "-w",
                "/code",
                container_name,
                "./" + metadata["entry_point"],
                "/input_images",
                "/output_data/result.csv",
            ]
        else:
            epsilon = np.random.choice(ALLOWED_EPS)
            cmd = [
                docker_binary,
                "run",
                "--network=none",
                "-m=24g",
                "-v",
                "{0}:/input_images:ro".format(self._sample_input_dir),
                "-v",
                "{0}:/output_images".format(self._sample_output_dir),
                "-v",
                "{0}:/code".format(self._extracted_submission_dir),
                "-w",
                "/code",
                container_name,
                "./" + metadata["entry_point"],
                "/input_images",
                "/output_images",
                str(epsilon),
            ]
        logging.info("Command to run submission: %s", " ".join(cmd))
        return shell_call(cmd)

    def _verify_output(self, submission_type):
        """Verifies correctness of the submission output.

        Args:
          submission_type: type of the submission

        Returns:
          True if output looks valid
        """
        result = True
        if submission_type == "defense":
            try:
                image_classification = load_defense_output(
                    os.path.join(self._sample_output_dir, "result.csv")
                )
                expected_keys = [
                    IMAGE_NAME_PATTERN.format(i) for i in range(BATCH_SIZE)
                ]
                if set(image_classification.keys()) != set(expected_keys):
                    logging.error("Classification results are not saved for all images")
                    result = False
            except IOError as e:
                logging.error("Failed to read defense output file: %s", e)
                result = False
        else:
            for i in range(BATCH_SIZE):
                image_filename = os.path.join(
                    self._sample_output_dir, IMAGE_NAME_PATTERN.format(i)
                )
                try:
                    img = np.array(Image.open(image_filename).convert("RGB"))
                    if list(img.shape) != [299, 299, 3]:
                        logging.error(
                            "Invalid image size %s for image %s",
                            str(img.shape),
                            image_filename,
                        )
                        result = False
                except IOError as e:
                    result = False
        return result

    def validate_submission(self, filename, submission_type):
        """Validates submission.

        Args:
          filename: submission filename
          submission_type: type of the submission,
            one of 'attack', 'targeted_attack' or 'defense'

        Returns:
          whether submission is valid
        """
        if submission_type not in ALLOWED_SUBMISSION_TYPES:
            logging.error("Invalid submission type: %s", submission_type)
            return False
        self._prepare_temp_dir()
        # Convert filename to be absolute path,
        # relative path might cause problems when monting directory in Docker
        filename = os.path.abspath(filename)
        # extract submission
        if not self._extract_submission(filename):
            return False
        # verify submission size
        if not self._verify_submission_size():
            return False
        # Load metadata
        metadata = self._load_and_verify_metadata(submission_type)
        if not metadata:
            return False
        # verify docker container size
        if not self._verify_docker_image_size(metadata["container_gpu"]):
            return False
        # Try to run submission on sample data
        self._prepare_sample_data(submission_type)
        if not self._run_submission(metadata):
            logging.error("Failure while running submission")
            return False
        if not self._verify_output(submission_type):
            logging.warning(
                "Some of the outputs of your submission are invalid or "
                "missing. You submission still will be evaluation "
                "but you might get lower score."
            )
        return True
