from __future__ import annotations

import contextlib
import shlex
import tempfile
from pathlib import Path
from typing import Sequence

import deprecated
from kraken.common import flatten
from kraken.core.api import Project, Property

from . import DockerBuildTask
from .dockerapi import docker_load, docker_run
from .util import render_docker_auth, update_run_commands


class KanikoBuildTask(DockerBuildTask):
    """
    An implementation for building Docker images with Kaniko.

    In order to make secrets available in the Dockerfile under `/run/secrets`, and `RUN` command that is executed
    by the root user will be surrounded by additional shell commands to link `/kaniko/secrets` to `/run/secrets`.
    This is to ensure a degree of compatibility with BuildKit, which mounts secrets to `/run/secrets` when using
    the `--secret` flag.

    However, this compatibility is limited because we can only actually create a at `/run/secrets` if the `RUN`
    command runs as root. Any command that is executed as a non-root user will not be able to access the secrets.
    We use a heuristic to detect this, but it is not perfect. If you need to use secrets in a `RUN` command that
    is executed as a non-root user, you will need to read from `/kaniko/secrets` instead.

    If you start from a base image that doesn't start as the root user, you will need to explicitly add a `USER`
    command to your Dockerfile to either inform the `KanikoBuildTask` that the subsequent commands are not run as
    root, or use it to switch to the root user. Otherwise, the `KanikoBuildTask` will assume that the subsequent
    commands are run as root and it can not know in advance if they are or not.
    """

    kaniko_image: Property[str] = Property.default("gcr.io/kaniko-project/executor:v1.9.0-debug")
    kaniko_context: Property[str] = Property.default("/workspace")
    kaniko_cache_copy_layers: Property[bool] = Property.default(True)
    kaniko_snapshot_mode: Property[str] = Property.default("redo")
    kaniko_secrets_mount_dir: Property[str] = Property.default("/kaniko/secrets")
    kaniko_secrets_from_env: Property[Sequence[str]] = Property.default(())
    kaniko_use_compressed_caching: Property[bool] = Property.default(True)

    def __init__(self, name: str, project: Project) -> None:
        super().__init__(name, project)
        self.preprocess_dockerfile.set(True)

    def render_main_script(self, executor_command: list[str]) -> str:
        """Renders the shell script that will be executed in the Kaniko container."""

        docker_config = render_docker_auth(self.auth.get())

        script = []
        script += [
            "mkdir -p /kaniko/.docker",
            "cat << EOF > /kaniko/.docker/config.json",
            docker_config,
            "EOF",
        ]

        secrets_mount_dir = self.kaniko_secrets_mount_dir.get()
        if self.secrets.get():
            script += [f"mkdir -p {shlex.quote(secrets_mount_dir)}"]
            for secret, value in self.secrets.get().items():
                script += [f"echo {shlex.quote(value)} > {shlex.quote(secrets_mount_dir + '/' + secret)}"]
        if self.kaniko_secrets_from_env.get():
            for env in self.kaniko_secrets_from_env.get():
                script += [f"echo ${env} > {shlex.quote(secrets_mount_dir + '/' + env)}"]

        script += [" ".join(map(shlex.quote, executor_command))]
        return "\n".join(script)

    def get_kaniko_executor_command(self, dockerfile: str | None, tar_path: str | None) -> list[str]:
        if tar_path and not self.tags:
            raise ValueError("Need at least one destination (tag) when exporting to an image tarball")
        executor_command = ["/kaniko/executor"]
        executor_command += flatten(("--build-arg", f"{key}={value}") for key, value in self.build_args.get().items())
        cache_repo = self.cache_repo.get()
        if cache_repo:
            executor_command += ["--cache-repo", cache_repo]
        if self.cache.get_or(False):
            executor_command += ["--cache=true"]
        executor_command += flatten(("--destination", destination) for destination in self.tags.get())
        if dockerfile:
            executor_command += ["--dockerfile", dockerfile]
        if not self.push.get():
            executor_command += ["--no-push"]
        executor_command += ["--snapshotMode", self.kaniko_snapshot_mode.get()]
        if self.squash.get():
            executor_command += ["--single-snapshot"]
        if self.kaniko_cache_copy_layers.get():
            executor_command += ["--cache-copy-layers"]
        if not self.kaniko_use_compressed_caching.get():
            executor_command += ["--compressed-caching=false"]
        target = self.target.get()
        if target:
            executor_command += ["--target", target]
        if tar_path:
            executor_command += ["--tarPath", tar_path]
        executor_command += ["--context", self.kaniko_context.get()]
        return executor_command

    @deprecated.deprecated("use KanikoBuildTask.render_main_script() instead", version="0.3.6")
    def _render_main_script(self, executor_command: list[str]) -> str:
        return self.render_main_script(executor_command)

    @deprecated.deprecated("use KanikoBuildTask.get_kaniko_executor_command() instead", version="0.3.6")
    def _get_kaniko_executor_command(self, dockerfile: str | None, tar_path: str | None) -> list[str]:
        return self.get_kaniko_executor_command(dockerfile, tar_path)

    def _build(
        self,
        exit_stack: contextlib.ExitStack,
    ) -> None:
        volumes = [f"{self.build_context.get().absolute()}:{self.kaniko_context.get()}"]

        # If the Dockerfile is not relative to the build context, we need to mount it explicitly.
        dockerfile = self.dockerfile.get()
        in_container_dockerfile: str | None = None
        try:
            in_container_dockerfile = str(dockerfile.absolute().relative_to(self.build_context.get().absolute()))
        except ValueError:
            in_container_dockerfile = "/kaniko/Dockerfile"
            volumes += [f"{dockerfile.absolute()}:{in_container_dockerfile}"]

        # If the image needs to be loaded into the Docker daemon after building, we need to always
        # export it to a file.
        image_output_file = self.image_output_file.get_or(None)
        if self.load.get() and not image_output_file:
            tempdir = exit_stack.enter_context(tempfile.TemporaryDirectory())
            image_output_file = Path(tempdir) / "image.tgz"

        # Construct the tar path for inside the container.
        tar_path: str | None = None
        if image_output_file:
            volumes += [f"{image_output_file.parent.absolute()}:/kaniko/out"]
            tar_path = f"/kaniko/out/{image_output_file.name}"

        executor_command = self._get_kaniko_executor_command(in_container_dockerfile, tar_path)

        script = self._render_main_script(executor_command)

        result = docker_run(
            image=self.kaniko_image.get(),
            args=["sh", "-c", script],
            entrypoint="",
            remove=True,
            volumes=volumes,
            workdir=self.kaniko_context.get(),
            platform=self.platform.get_or(None),
        )

        if result != 0:
            raise Exception(f"Kaniko build failed with exit code {result}")

        if self.load.get():
            assert image_output_file is not None, "image_output_file is expected to be set when config.load == True"
            result = docker_load(image_output_file)
            if result != 0:
                raise Exception(f"Docker load failed with exit code {result}")

    # DockerBuildTask

    def _preprocess_dockerfile(self, dockerfile: Path) -> str:
        return update_run_commands(
            dockerfile.read_text(),
            prefix="ln -sf /kaniko/secrets /run/secrets && ( ",
            suffix=" ); __ret=$?; unlink /run/secrets; exit $__ret",
            # We can only link to /run/secrets if the user running the command is the root user. In any other
            # case, we cannot link to /run/secrets and reading the secrets from /run/secrets in the Dockerfile
            # will fail.
            only_for_root_user=True,
        )

    # Task

    def finalize(self) -> None:
        if self.cache.get() and not self.push.get() and not self.cache_repo.get():
            self.logger.warning(
                "Disabling cache in Kaniko build %s because it must be combined with push or cache_repo",
                self,
            )
            self.cache.set(False)
        cache_repo = self.cache_repo.get_or(None)
        if cache_repo and ":" in cache_repo:
            raise ValueError(f"Kaniko --cache-repo argument cannot contain `:` (got: {cache_repo!r})")
        return super().finalize()

    def execute(self) -> None:
        with contextlib.ExitStack() as exit_stack:
            self._build(exit_stack)
