import contextlib
import os
from collections import namedtuple

import yaml
from dagster import __version__ as current_dagster_version
from dagster import check

from .ecr import ecr_image, get_aws_account_id, get_aws_region
from .utils import (
    execute_docker_build,
    execute_docker_push,
    execute_docker_tag,
    python_version_image_tag,
)

# Default repository prefix used for local images
DEFAULT_LOCAL_PREFIX = "dagster"

# Location of the template assets used here
IMAGES_PATH = os.path.join(os.path.dirname(__file__), "images")


@contextlib.contextmanager
def do_nothing(_cwd):
    yield


class DagsterDockerImage(namedtuple("_DagsterDockerImage", "image build_cm path")):
    """Represents a Dagster image.

    Properties:
        image (str): Name of the image
        build_cm (function): function that is a context manager for build (e.g. for populating a
            build cache)
        path (Optional(str)): The path to the image's path. Defaults to docker/images/<IMAGE NAME>
    """

    def __new__(cls, image, build_cm=do_nothing, path=None):
        return super(DagsterDockerImage, cls).__new__(
            cls,
            check.str_param(image, "image"),
            check.callable_param(build_cm, "build_cm"),
            check.opt_str_param(
                path, "path", default=os.path.join(os.path.dirname(__file__), "images", image)
            ),
        )

    @property
    def python_versions(self):
        """List of Python versions supported for this image."""
        with open(os.path.join(self.path, "versions.yaml"), "r") as f:
            versions = yaml.safe_load(f.read())
        return list(versions.keys())

    def _get_last_updated_for_python_version(self, python_version):
        """Retrieve the last_updated timestamp for a particular python_version of this image."""
        check.str_param(python_version, "python_version")
        with open(os.path.join(self.path, "last_updated.yaml"), "r") as f:
            last_updated = yaml.safe_load(f.read())
            return last_updated[python_version]

    def _set_last_updated_for_python_version(self, timestamp, python_version):
        """Update the last_updated timestamp for a particular python_version of this image."""
        check.str_param(timestamp, "timestamp")
        check.str_param(python_version, "python_version")

        last_updated = {}

        last_updated_path = os.path.join(self.path, "last_updated.yaml")
        if os.path.exists(last_updated_path):
            with open(last_updated_path, "r") as f:
                last_updated = yaml.safe_load(f.read())

        last_updated[python_version] = timestamp

        with open(os.path.join(self.path, "last_updated.yaml"), "w") as f:
            yaml.dump(last_updated, f, default_flow_style=False)

    def local_image(self, python_version):
        """Generates the local image name, like: "dagster/foo:some-tag" """
        check.str_param(python_version, "python_version")

        last_updated = self._get_last_updated_for_python_version(python_version)
        tag = python_version_image_tag(python_version, last_updated)
        return "{}/{}:{}".format(DEFAULT_LOCAL_PREFIX, self.image, tag)

    def aws_image(self, python_version=None, custom_tag=None):
        """Generates the AWS ECR image name, like:
        "1234567890.dkr.ecr.us-west-1.amazonaws.com/foo:some-tag"
        """
        check.invariant(not (python_version and custom_tag))
        check.opt_str_param(python_version, "python_version")
        check.opt_str_param(custom_tag, "custom_tag")

        if python_version:
            last_updated = self._get_last_updated_for_python_version(python_version)
            tag = python_version_image_tag(python_version, last_updated)
        else:
            tag = custom_tag

        return ecr_image(
            self.image,
            tag,
            aws_account_id=get_aws_account_id(),
            aws_region=get_aws_region(),
        )

    def _get_docker_args(self, python_version):
        """Retrieve Docker arguments from this image's versions.yaml, and update with latest Dagster
        version.

        Also, we allow references in the image versions.yaml to another Dagster image to use as a
        base image. If defined, set the BASE_IMAGE Docker arg from the full name of the parent
        image.
        """
        with open(os.path.join(self.path, "versions.yaml"), "r") as f:
            versions = yaml.safe_load(f.read())
            image_info = versions.get(python_version, {})

        docker_args = image_info.get("docker_args", {})

        if "base_image" in image_info:
            check.invariant(
                "BASE_IMAGE" not in docker_args, "Cannot override an existing BASE_IMAGE"
            )

            base_image = DagsterDockerImage(image_info["base_image"]["name"])
            source = image_info["base_image"]["source"]

            if source == "aws":
                docker_args["BASE_IMAGE"] = base_image.aws_image(python_version)
            elif source == "local":
                docker_args["BASE_IMAGE"] = base_image.local_image(python_version)
            else:
                raise Exception("Unrecognized source {}".format(source))

        # Set Dagster version
        docker_args["DAGSTER_VERSION"] = current_dagster_version
        return docker_args

    def build(self, timestamp, dagster_version, python_version):
        check.str_param(timestamp, "timestamp")
        check.str_param(python_version, "python_version")
        check.invariant(
            dagster_version == current_dagster_version,
            desc="Current dagster version ({}) does not match provided arg ({})".format(
                current_dagster_version, dagster_version
            ),
        )
        with self.build_cm(self.path):
            self._set_last_updated_for_python_version(timestamp, python_version)

            execute_docker_build(
                self.local_image(python_version),
                docker_args=self._get_docker_args(python_version),
                cwd=self.path,
            )

    def push(self, python_version, custom_tag=None):
        """Push this image to ECR."""

        if custom_tag:
            execute_docker_tag(
                self.local_image(python_version),
                self.aws_image(python_version=None, custom_tag=custom_tag),
            )
            execute_docker_push(self.aws_image(python_version=None, custom_tag=custom_tag))
        else:
            execute_docker_tag(self.local_image(python_version), self.aws_image(python_version))
            execute_docker_push(self.aws_image(python_version))

[((29, 25, 29, 81), 'collections.namedtuple', 'namedtuple', ({(29, 36, 29, 57): '"""_DagsterDockerImage"""', (29, 59, 29, 80): '"""image build_cm path"""'}, {}), "('_DagsterDockerImage', 'image build_cm path')", False, 'from collections import namedtuple\n'), ((21, 27, 21, 52), 'os.path.dirname', 'os.path.dirname', ({(21, 43, 21, 51): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((58, 8, 58, 57), 'dagster.check.str_param', 'check.str_param', ({(58, 24, 58, 38): 'python_version', (58, 40, 58, 56): '"""python_version"""'}, {}), "(python_version, 'python_version')", False, 'from dagster import check\n'), ((65, 8, 65, 47), 'dagster.check.str_param', 'check.str_param', ({(65, 24, 65, 33): 'timestamp', (65, 35, 65, 46): '"""timestamp"""'}, {}), "(timestamp, 'timestamp')", False, 'from dagster import check\n'), ((66, 8, 66, 57), 'dagster.check.str_param', 'check.str_param', ({(66, 24, 66, 38): 'python_version', (66, 40, 66, 56): '"""python_version"""'}, {}), "(python_version, 'python_version')", False, 'from dagster import check\n'), ((70, 28, 70, 72), 'os.path.join', 'os.path.join', ({(70, 41, 70, 50): 'self.path', (70, 52, 70, 71): '"""last_updated.yaml"""'}, {}), "(self.path, 'last_updated.yaml')", False, 'import os\n'), ((71, 11, 71, 44), 'os.path.exists', 'os.path.exists', ({(71, 26, 71, 43): 'last_updated_path'}, {}), '(last_updated_path)', False, 'import os\n'), ((82, 8, 82, 57), 'dagster.check.str_param', 'check.str_param', ({(82, 24, 82, 38): 'python_version', (82, 40, 82, 56): '"""python_version"""'}, {}), "(python_version, 'python_version')", False, 'from dagster import check\n'), ((92, 8, 92, 60), 'dagster.check.invariant', 'check.invariant', ({(92, 24, 92, 59): '(not (python_version and custom_tag))'}, {}), '(not (python_version and custom_tag))', False, 'from dagster import check\n'), ((93, 8, 93, 61), 'dagster.check.opt_str_param', 'check.opt_str_param', ({(93, 28, 93, 42): 'python_version', (93, 44, 93, 60): '"""python_version"""'}, {}), "(python_version, 'python_version')", False, 'from dagster import check\n'), ((94, 8, 94, 53), 'dagster.check.opt_str_param', 'check.opt_str_param', ({(94, 28, 94, 38): 'custom_tag', (94, 40, 94, 52): '"""custom_tag"""'}, {}), "(custom_tag, 'custom_tag')", False, 'from dagster import check\n'), ((143, 8, 143, 47), 'dagster.check.str_param', 'check.str_param', ({(143, 24, 143, 33): 'timestamp', (143, 35, 143, 46): '"""timestamp"""'}, {}), "(timestamp, 'timestamp')", False, 'from dagster import check\n'), ((144, 8, 144, 57), 'dagster.check.str_param', 'check.str_param', ({(144, 24, 144, 38): 'python_version', (144, 40, 144, 56): '"""python_version"""'}, {}), "(python_version, 'python_version')", False, 'from dagster import check\n'), ((42, 12, 42, 43), 'dagster.check.str_param', 'check.str_param', ({(42, 28, 42, 33): 'image', (42, 35, 42, 42): '"""image"""'}, {}), "(image, 'image')", False, 'from dagster import check\n'), ((43, 12, 43, 54), 'dagster.check.callable_param', 'check.callable_param', ({(43, 33, 43, 41): 'build_cm', (43, 43, 43, 53): '"""build_cm"""'}, {}), "(build_cm, 'build_cm')", False, 'from dagster import check\n'), ((78, 12, 78, 64), 'yaml.dump', 'yaml.dump', (), '', False, 'import yaml\n'), ((124, 12, 126, 13), 'dagster.check.invariant', 'check.invariant', ({(125, 16, 125, 47): "('BASE_IMAGE' not in docker_args)", (125, 49, 125, 89): '"""Cannot override an existing BASE_IMAGE"""'}, {}), "('BASE_IMAGE' not in docker_args,\n    'Cannot override an existing BASE_IMAGE')", False, 'from dagster import check\n'), ((52, 18, 52, 58), 'os.path.join', 'os.path.join', ({(52, 31, 52, 40): 'self.path', (52, 42, 52, 57): '"""versions.yaml"""'}, {}), "(self.path, 'versions.yaml')", False, 'import os\n'), ((59, 18, 59, 62), 'os.path.join', 'os.path.join', ({(59, 31, 59, 40): 'self.path', (59, 42, 59, 61): '"""last_updated.yaml"""'}, {}), "(self.path, 'last_updated.yaml')", False, 'import os\n'), ((77, 18, 77, 62), 'os.path.join', 'os.path.join', ({(77, 31, 77, 40): 'self.path', (77, 42, 77, 61): '"""last_updated.yaml"""'}, {}), "(self.path, 'last_updated.yaml')", False, 'import os\n'), ((117, 18, 117, 58), 'os.path.join', 'os.path.join', ({(117, 31, 117, 40): 'self.path', (117, 42, 117, 57): '"""versions.yaml"""'}, {}), "(self.path, 'versions.yaml')", False, 'import os\n'), ((45, 51, 45, 76), 'os.path.dirname', 'os.path.dirname', ({(45, 67, 45, 75): '__file__'}, {}), '(__file__)', False, 'import os\n')]