import json
import os
import random
import string
import subprocess
import tempfile
from contextlib import ExitStack, contextmanager

from osbuild.util.mnt import MountGuard, MountPermissions

# use `/run/osbuild/containers/storage` for the host's containers-storage bind mount
HOST_CONTAINERS_STORAGE = os.path.join(os.sep, "run", "osbuild", "containers", "storage")
# use `/run/osbuild/containers/storage2` for and empty containers-storage
HOST_CONTAINERS_STORAGE2 = os.path.join(os.sep, "run", "osbuild", "containers", "storage2")


def is_manifest_list(data):
    """Inspect a manifest determine if it's a multi-image manifest-list."""
    media_type = data.get("mediaType")
    #  Check if mediaType is set according to docker or oci specifications
    if media_type in ("application/vnd.docker.distribution.manifest.list.v2+json",
                      "application/vnd.oci.image.index.v1+json"):
        return True

    # According to the OCI spec, setting mediaType is not mandatory. So, if it is not set at all, check for the
    # existence of manifests
    if media_type is None and data.get("manifests") is not None:
        return True

    return False


def parse_manifest_list(manifests):
    """Return a map with single-image manifest digests as keys and the manifest-list digest as the value for each"""
    manifest_files = manifests["data"]["files"]
    manifest_map = {}
    for fname in manifest_files:
        filepath = os.path.join(manifests["path"], fname)
        with open(filepath, mode="r", encoding="utf-8") as mfile:
            data = json.load(mfile)

        for manifest in data["manifests"]:
            digest = manifest["digest"]  # single image manifest digest
            manifest_map[digest] = fname

    return manifest_map


def manifest_digest(path):
    """Get the manifest digest for a container at path, stored in dir: format"""
    return subprocess.check_output(["skopeo", "manifest-digest", os.path.join(path, "manifest.json")]).decode().strip()


def parse_containers_input(inputs):
    manifests = inputs.get("manifest-lists")
    manifest_map = {}
    manifest_files = {}
    if manifests:
        manifest_files = manifests["data"]["files"]
        # reverse map manifest-digest -> manifest-list path
        manifest_map = parse_manifest_list(manifests)

    images = inputs["images"]
    archives = images["data"]["archives"]

    res = {}
    for checksum, data in archives.items():
        filepath = os.path.join(images["path"], checksum)
        list_path = None
        if data["format"] == "dir":
            digest = manifest_digest(filepath)

            # get the manifest list path for this image
            list_digest = manifest_map.get(digest)
            if list_digest:
                # make sure all manifest files are used
                del manifest_files[list_digest]
                list_path = os.path.join(manifests["path"], list_digest)

        if data["format"] == "containers-storage":
            # filepath is the storage bindmount
            filepath = os.path.join(images["path"], "storage")

        res[checksum] = {
            "filepath": filepath,
            "manifest-list": list_path,
            "data": data,
            "checksum": checksum,  # include the checksum in the value
        }

    if manifest_files:
        raise RuntimeError(
            "The following manifest lists specified in the input did not match any of the container images: " +
            ", ".join(manifest_files)
        )

    return res


def merge_manifest(list_manifest, destination):
    """
    Merge the list manifest into the image directory. This preserves the manifest list with the image in the registry so
    that users can run or inspect a container using the original manifest list digest used to pull the container.

    See https://github.com/containers/skopeo/issues/1935
    """
    # calculate the checksum of the manifest of the container image in the destination
    dest_manifest = os.path.join(destination, "manifest.json")
    manifest_checksum = subprocess.check_output(["skopeo", "manifest-digest", dest_manifest]).decode().strip()
    parts = manifest_checksum.split(":")
    assert len(parts) == 2, f"unexpected output for skopeo manifest-digest: {manifest_checksum}"
    manifest_checksum = parts[1]

    # rename the manifest to its checksum
    os.rename(dest_manifest, os.path.join(destination, manifest_checksum + ".manifest.json"))

    # copy the index manifest into the destination
    subprocess.run(["cp", "--reflink=auto", "-a", list_manifest, dest_manifest], check=True)


@contextmanager
def containers_storage_source(image, image_filepath, container_format):
    storage_conf = image["data"]["storage"]
    driver = storage_conf.get("driver", "overlay")

    storage_path = HOST_CONTAINERS_STORAGE
    storage_path_empty = HOST_CONTAINERS_STORAGE2
    os.makedirs(storage_path, exist_ok=True)

    with MountGuard() as mg:
        mg.mount(image_filepath, storage_path, permissions=MountPermissions.READ_WRITE)
        # NOTE: the ostree.deploy.container needs explicit `rw` access to
        # the containers-storage store even when bind mounted. Remounting
        # the bind mount is a pretty dirty fix to get us up and running with
        # containers-storage in `bootc-image-builder`. We could maybe check
        # if we're inside a bib-continaer and only run this conidtionally.
        mg.mount(image_filepath, storage_path, remount=True, permissions=MountPermissions.READ_WRITE)

        image_id = image["checksum"].split(":")[1]

        # Only the overlayfs backend supoorts additional image store
        use_additional_image_store = driver == "overlay"

        if use_additional_image_store:
            # If additional image store is available, then we use a setup where the base graphroot is an
            # empty storage directory, and then we access the host via a separate directory (--imagestore
            # or additional image store). This allows us to support accessing the host storage via virtiofsd
            # mounts. Virtiofs mounts don't supoprt being used as an overlayfs upper dir, but when we set it
            # up like this, the upper directory ends up being in the empty graphroot which is on tmpfs, so
            # things work.
            podman_opts = [f"--root={storage_path_empty}", f"--imagestore={storage_path}"]
            image_source = f"{container_format}:[{driver}@{storage_path_empty}+/run/containers/storage:additionalimagestore={storage_path}]{image_id}"
        else:
            # On vfs backends, we use the traditional setup
            podman_opts = [f"--imagestore={storage_path}"]
            image_source = f"{container_format}:[{driver}@{storage_path}+/run/containers/storage]{image_id}"

        yield image_source, podman_opts

        if driver == "overlay":
            # NOTE: the overlay sub-directory isn't always released,
            # so we need to force unmount it
            ret = subprocess.run(["umount", "-f", "--lazy", os.path.join(storage_path, "overlay")], check=False)
            if ret.returncode != 0:
                print(f"WARNING: umount of overlay dir failed with an error: {ret}")


@contextmanager
def dir_oci_archive_source(image, image_filepath, container_format):
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp_source = os.path.join(tmpdir, "image")

        if container_format == "dir" and image["manifest-list"]:
            # copy the source container to the tmp source so we can merge the manifest into it
            subprocess.run(["cp", "-a", "--reflink=auto", image_filepath, tmp_source], check=True)
            merge_manifest(image["manifest-list"], tmp_source)
        else:
            # We can't have special characters like ":" in the source names because containers/image
            # treats them special, like e.g. /some/path:tag, so we make a symlink to the real name
            # and pass the symlink name to skopeo to make it work with anything
            os.symlink(image_filepath, tmp_source)

        image_source = f"{container_format}:{tmp_source}"
        yield image_source, None


@contextmanager
def container_source(image):
    image_filepath = image["filepath"]
    container_format = image["data"]["format"]
    image_name = image["data"]["name"]

    if container_format not in ("dir", "oci-archive", "containers-storage"):
        raise RuntimeError(f"Unknown container format {container_format}")

    if container_format == "containers-storage":
        container_source_fn = containers_storage_source
    elif container_format in ("dir", "oci-archive"):
        container_source_fn = dir_oci_archive_source
    else:
        raise RuntimeError(f"Unknown container format {container_format}")

    # pylint: disable=contextmanager-generator-missing-cleanup
    # thozza: As far as I can tell, the problematic use case is when the ctx manager is used inside a generator.
    # However, this is not the case here. The ctx manager is used inside another ctx manager with the expectation
    # that the inner ctx manager won't be cleaned up until the execution returns to this ctx manager.
    with container_source_fn(image, image_filepath, container_format) as (image_source, podman_opts):
        yield (image_name, image_source, podman_opts)


@contextmanager
def container_mount(image):
    # Helper function for doing the `podman image mount`
    @contextmanager
    def _mount_container(img, podman_opts):
        cmd = ["podman"] + (podman_opts or [])

        result = subprocess.run(cmd + ["image", "mount", img], encoding="utf-8",
                                check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.returncode != 0:
            code = result.returncode
            msg = result.stderr.strip()
            raise RuntimeError(f"Failed to mount image ({code}): {msg}")
        try:
            yield result.stdout.strip()
        finally:
            subprocess.run(cmd + ["image", "umount", img], check=True)

    with container_source(image) as (_, source, podman_opts):
        with ExitStack() as cm:
            img = ""
            if image["data"]["format"] == 'containers-storage':
                # In the case where we are container storage we don't need to
                # skopeo copy. We already have access to a mounted container storage
                # that has the image ready to use.
                image_id = image["checksum"].split(":")[1]
                img = image_id
            else:
                # We cannot use a tmpdir as storage here because of
                # https://github.com/containers/storage/issues/1779 so instead
                # just pick a random suffix. This runs inside bwrap which gives a
                # tmp /var so it does not really matter much.
                tmp_image_name = "tmp-container-mount-" + "".join(random.choices(string.digits, k=14))
                cm.callback(subprocess.run, ["podman", "rmi", tmp_image_name], check=True)
                # skopeo needs /var/tmp but the bwrap env is minimal and may not have it
                os.makedirs("/var/tmp", mode=0o1777, exist_ok=True)
                cmd = ["skopeo", "copy", source, f"containers-storage:{tmp_image_name}"]
                subprocess.run(cmd, check=True)
                img = tmp_image_name

            with _mount_container(img, podman_opts) as container_mountpoint:
                yield container_mountpoint
