import aiohttp
import asyncio
import base64
import fnmatch
import hashlib
import json
import logging

from gettext import gettext as _
from urllib.parse import urljoin, urlparse, urlunparse

from asgiref.sync import sync_to_async
from django.db import IntegrityError
from pulpcore.plugin.models import Artifact, ProgressReport, Remote
from pulpcore.plugin.stages import DeclarativeArtifact, DeclarativeContent, Stage
from pulpcore.plugin.constants import TASK_STATES

from pulp_container.constants import (
    MEDIA_TYPE,
    SIGNATURE_API_EXTENSION_VERSION,
    SIGNATURE_HEADER,
    SIGNATURE_SOURCE,
    SIGNATURE_TYPE,
)
from pulp_container.app.models import (
    Blob,
    BlobManifest,
    Manifest,
    ManifestListManifest,
    ManifestSignature,
    Tag,
)
from pulp_container.app.utils import (
    determine_media_type,
    extract_data_from_signature,
    urlpath_sanitize,
)

log = logging.getLogger(__name__)


V2_ACCEPT_HEADERS = {
    "Accept": ",".join(
        [
            MEDIA_TYPE.MANIFEST_V2,
            MEDIA_TYPE.MANIFEST_LIST,
            MEDIA_TYPE.INDEX_OCI,
            MEDIA_TYPE.MANIFEST_OCI,
        ]
    )
}


def _save_artifact_blocking(artifact_attributes):
    saved_artifact = Artifact(**artifact_attributes)
    try:
        saved_artifact.save()
    except IntegrityError:
        del artifact_attributes["file"]
        saved_artifact = Artifact.objects.get(**artifact_attributes)
        saved_artifact.touch()
    return saved_artifact


class ContainerFirstStage(Stage):
    """
    The first stage of a pulp_container sync pipeline.

    In this stage all the content is discovered, including the nested one.

    """

    def __init__(self, remote, signed_only):
        """Initialize the stage."""
        super().__init__()
        self.remote = remote
        self.deferred_download = self.remote.policy != Remote.IMMEDIATE
        self.signed_only = signed_only

    async def run(self):
        """
        ContainerFirstStage.
        """

        async def get_signature_source(self):
            """
            Find out where signatures come from: sigstore, extension API or not available at all.
            """
            if self.remote.sigstore:
                return SIGNATURE_SOURCE.SIGSTORE

            registry_v2_url = urljoin(self.remote.url, "v2")
            extension_check_downloader = self.remote.get_noauth_downloader(url=registry_v2_url)
            response_headers = {}
            try:
                result = await extension_check_downloader.run()
                response_headers = result.headers
            except aiohttp.client_exceptions.ClientResponseError as exc:
                if exc.status == 401:
                    response_headers = dict(exc.headers)
            if response_headers.get(SIGNATURE_HEADER) == "1":
                return SIGNATURE_SOURCE.API_EXTENSION

        tag_list = []
        tag_dcs = []
        to_download = []
        man_dcs = {}
        signature_dcs = []
        signature_source = await get_signature_source(self)

        if signature_source is None and self.signed_only:
            raise ValueError(
                _(
                    "It is requested to sync only signed content but no sigstore URL is "
                    "provided. Please configure a `sigstore` on your Remote or set "
                    "`signed_only` to `False` for your sync request."
                )
            )

        async with ProgressReport(
            message="Downloading tag list", code="sync.downloading.tag_list", total=1
        ) as pb:
            repo_name = self.remote.namespaced_upstream_name
            relative_url = "/v2/{name}/tags/list".format(name=repo_name)
            tag_list_url = urljoin(self.remote.url, relative_url)
            list_downloader = self.remote.get_downloader(url=tag_list_url)
            await list_downloader.run(extra_data={"repo_name": repo_name})

            with open(list_downloader.path) as tags_raw:
                tags_dict = json.loads(tags_raw.read())
                tag_list = tags_dict["tags"]

            # check for the presence of the pagination link header
            link = list_downloader.response_headers.get("Link")
            await self.handle_pagination(link, repo_name, tag_list)
            tag_list = self.filter_tags(tag_list)
            await pb.aincrement()

        for tag_name in tag_list:
            relative_url = "/v2/{name}/manifests/{tag}".format(
                name=self.remote.namespaced_upstream_name, tag=tag_name
            )
            url = urljoin(self.remote.url, relative_url)
            downloader = self.remote.get_downloader(url=url)
            to_download.append(downloader.run(extra_data={"headers": V2_ACCEPT_HEADERS}))

        async with ProgressReport(
            message="Processing Tags",
            code="sync.processing.tag",
            state=TASK_STATES.RUNNING,
            total=len(tag_list),
        ) as pb_parsed_tags:

            for download_tag in asyncio.as_completed(to_download):
                tag = await download_tag
                with open(tag.path, "rb") as content_file:
                    raw_data = content_file.read()
                tag.artifact_attributes["file"] = tag.path
                saved_artifact = await sync_to_async(_save_artifact_blocking)(
                    tag.artifact_attributes
                )

                tag_name = tag.url.split("/")[-1]
                tag_dc = DeclarativeContent(Tag(name=tag_name))

                content_data = json.loads(raw_data)
                media_type = determine_media_type(content_data, tag)
                if media_type in (MEDIA_TYPE.MANIFEST_LIST, MEDIA_TYPE.INDEX_OCI):
                    list_dc = self.create_tagged_manifest_list(
                        tag_name, saved_artifact, content_data, media_type
                    )
                    for manifest_data in content_data.get("manifests"):
                        man_dc = self.create_manifest(list_dc, manifest_data)
                        if signature_source is not None:
                            man_sig_dcs = await self.create_signatures(man_dc, signature_source)
                            if self.signed_only and not man_sig_dcs:
                                log.info(
                                    _(
                                        "The unsigned image {img_digest} which is a part of the "
                                        "manifest list {ml_digest} (tagged as `{tag}`) can't be "
                                        "synced due to a requirement to sync signed content only. "
                                        "The whole manifest list is skipped.".format(
                                            img_digest=man_dc.content.digest,
                                            ml_digest=list_dc.content.digest,
                                            tag=tag_name,
                                        )
                                    )
                                )
                                # do not pass down the pipeline a manifest list with unsigned
                                # manifests.
                                break
                            signature_dcs.extend(man_sig_dcs)
                        man_dcs[man_dc.content.digest] = man_dc
                        await self.put(man_dc)

                    else:
                        # only pass the manifest list and tag down the pipeline if there were no
                        # issues with signatures (no `break` in the `for` loop)
                        tag_dc.extra_data["tagged_manifest_dc"] = list_dc
                        await self.put(list_dc)
                        tag_dcs.append(tag_dc)
                        await pb_parsed_tags.aincrement()

                else:
                    man_dc = self.create_tagged_manifest(
                        tag_name, saved_artifact, content_data, raw_data, media_type
                    )
                    if signature_source is not None:
                        man_sig_dcs = await self.create_signatures(man_dc, signature_source)
                        if self.signed_only and not man_sig_dcs:
                            # do not pass down the pipeline unsigned manifests
                            continue
                        signature_dcs.extend(man_sig_dcs)
                    await self.put(man_dc)
                    tag_dc.extra_data["tagged_manifest_dc"] = man_dc
                    await man_dc.resolution()
                    await self.handle_blobs(man_dc, content_data)
                    tag_dcs.append(tag_dc)
                    await pb_parsed_tags.aincrement()

        for digest, man_dc in man_dcs.items():
            man = await man_dc.resolution()
            artifact = await sync_to_async(man._artifacts.get)()
            with artifact.file.open() as content_file:
                raw = content_file.read()
            content_data = json.loads(raw)
            await self.handle_blobs(man_dc, content_data)

        for tag_dc in tag_dcs:
            tagged_manifest_dc = tag_dc.extra_data["tagged_manifest_dc"]
            tag_dc.content.tagged_manifest = await tagged_manifest_dc.resolution()
            await self.put(tag_dc)

        for sig_dc in signature_dcs:
            signed_manifest_dc = sig_dc.extra_data["signed_manifest_dc"]
            sig_dc.content.signed_manifest = await signed_manifest_dc.resolution()
            await self.put(sig_dc)

    def filter_tags(self, tag_list):
        """
        Filter tags by a list of included and excluded tags.
        """
        include_tags = self.remote.include_tags
        if include_tags:
            tag_list = [
                tag
                for tag in tag_list
                if any(fnmatch.fnmatch(tag, pattern) for pattern in include_tags)
            ]

        exclude_tags = self.remote.exclude_tags
        if exclude_tags:
            tag_list = [
                tag
                for tag in tag_list
                if not any(fnmatch.fnmatch(tag, pattern) for pattern in exclude_tags)
            ]

        return tag_list

    async def handle_pagination(self, link, repo_name, tag_list):
        """
        Handle registries that have pagination enabled.
        """
        while link:
            # according RFC5988 URI-reference can be relative or absolute
            _, _, path, params, query, fragm = urlparse(link.split(";")[0].strip(">, <"))
            rel_link = urlunparse(("", "", path, params, query, fragm))
            link = urljoin(self.remote.url, rel_link)
            list_downloader = self.remote.get_downloader(url=link)
            await list_downloader.run(extra_data={"repo_name": repo_name})
            with open(list_downloader.path) as tags_raw:
                tags_dict = json.loads(tags_raw.read())
                tag_list.extend(tags_dict["tags"])
            link = list_downloader.response_headers.get("Link")

    async def handle_blobs(self, manifest_dc, content_data):
        """
        Handle blobs.
        """
        for layer in content_data.get("layers") or content_data.get("fsLayers"):
            if not self._include_layer(layer):
                continue
            blob_dc = self.create_blob(layer)
            blob_dc.extra_data["blob_relation"] = manifest_dc
            await self.put(blob_dc)
        layer = content_data.get("config", None)
        if layer:
            blob_dc = self.create_blob(layer, deferred_download=False)
            blob_dc.extra_data["config_relation"] = manifest_dc
            await self.put(blob_dc)

    def create_tagged_manifest_list(self, tag_name, saved_artifact, manifest_list_data, media_type):
        """
        Create a ManifestList.

        Args:
            tag_name (str): A name of a tag
            saved_artifact (pulpcore.plugin.models.Artifact): A saved manifest's Artifact
            manifest_list_data (dict): Data about a ManifestList
            media_type (str): The type of a manifest

        """
        digest = f"sha256:{saved_artifact.sha256}"
        manifest_list = Manifest(
            digest=digest,
            schema_version=manifest_list_data["schemaVersion"],
            media_type=media_type,
        )

        return self._create_manifest_declarative_content(
            manifest_list, saved_artifact, tag_name, digest
        )

    def create_tagged_manifest(self, tag_name, saved_artifact, manifest_data, raw_data, media_type):
        """
        Create an Image Manifest.

        Args:
            tag_name (str): A name of a tag
            saved_artifact (pulpcore.plugin.models.Artifact): A saved manifest's Artifact
            manifest_data (dict): Data about a single new ImageManifest.
            raw_data: (str): The raw JSON representation of the ImageManifest.
            media_type (str): The type of a manifest

        """
        if media_type in (MEDIA_TYPE.MANIFEST_V2, MEDIA_TYPE.MANIFEST_OCI):
            digest = f"sha256:{saved_artifact.sha256}"
        else:
            digest = self._calculate_digest(raw_data)

        manifest = Manifest(
            digest=digest, schema_version=manifest_data["schemaVersion"], media_type=media_type
        )

        return self._create_manifest_declarative_content(manifest, saved_artifact, tag_name, digest)

    def _create_manifest_declarative_content(self, manifest, saved_artifact, tag_name, digest):
        relative_url = f"/v2/{self.remote.namespaced_upstream_name}/manifests/"
        da_digest = self._create_manifest_declarative_artifact(
            relative_url + digest, saved_artifact, digest
        )
        da_tag = self._create_manifest_declarative_artifact(
            relative_url + tag_name, saved_artifact, digest
        )

        man_dc = DeclarativeContent(content=manifest, d_artifacts=[da_digest, da_tag])
        return man_dc

    def _create_manifest_declarative_artifact(self, relative_url, saved_artifact, digest):
        url = urljoin(self.remote.url, relative_url)
        da = DeclarativeArtifact(
            artifact=saved_artifact,
            url=url,
            relative_path=digest,
            remote=self.remote,
            extra_data={"headers": V2_ACCEPT_HEADERS},
        )
        return da

    def _create_signature_declarative_content(
        self, signature_raw, man_dc, name=None, signature_b64=None
    ):
        signature_json = extract_data_from_signature(signature_raw, man_dc.content.digest)
        if signature_json is None:
            return

        sig_digest = hashlib.sha256(signature_raw).hexdigest()
        signature = ManifestSignature(
            name=name or f"{man_dc.content.digest}@{sig_digest[:32]}",
            digest=f"sha256:{sig_digest}",
            type=SIGNATURE_TYPE.ATOMIC_SHORT,
            key_id=signature_json["signing_key_id"],
            timestamp=signature_json["signature_timestamp"],
            creator=signature_json["optional"].get("creator"),
            data=signature_b64 or base64.b64encode(signature_raw).decode(),
        )
        sig_dc = DeclarativeContent(
            content=signature,
            extra_data={"signed_manifest_dc": man_dc},
        )
        return sig_dc

    def create_manifest(self, list_dc, manifest_data):
        """
        Create an Image Manifest from manifest data in a ManifestList.

        Args:
            list_dc (pulpcore.plugin.stages.DeclarativeContent): dc for a ManifestList
            manifest_data (dict): Data about a single new ImageManifest.

        """
        digest = manifest_data["digest"]
        relative_url = "/v2/{name}/manifests/{digest}".format(
            name=self.remote.namespaced_upstream_name, digest=digest
        )
        manifest_url = urljoin(self.remote.url, relative_url)
        da = DeclarativeArtifact(
            artifact=Artifact(),
            url=manifest_url,
            relative_path=digest,
            remote=self.remote,
            extra_data={"headers": V2_ACCEPT_HEADERS},
        )
        manifest = Manifest(
            digest=manifest_data["digest"],
            schema_version=2
            if manifest_data["mediaType"] in (MEDIA_TYPE.MANIFEST_V2, MEDIA_TYPE.MANIFEST_OCI)
            else 1,
            media_type=manifest_data["mediaType"],
        )
        platform = {}
        p = manifest_data["platform"]
        platform["architecture"] = p["architecture"]
        platform["os"] = p["os"]
        platform["features"] = p.get("features", "")
        platform["variant"] = p.get("variant", "")
        platform["os.version"] = p.get("os.version", "")
        platform["os.features"] = p.get("os.features", "")
        man_dc = DeclarativeContent(
            content=manifest,
            d_artifacts=[da],
            extra_data={"relation": list_dc, "platform": platform},
        )
        return man_dc

    def create_blob(self, blob_data, deferred_download=True):
        """
        Create blob.

        Args:
            blob_data (dict): Data about a blob
            deferred_download (bool): boolean that indicates whether not to download a blob
                immediatly. Config blob is downloaded regardless of the remote's settings

        """
        digest = blob_data.get("digest") or blob_data.get("blobSum")
        blob_artifact = Artifact(sha256=digest[len("sha256:") :])
        blob = Blob(digest=digest)
        relative_url = "/v2/{name}/blobs/{digest}".format(
            name=self.remote.namespaced_upstream_name, digest=digest
        )
        blob_url = urljoin(self.remote.url, relative_url)
        da = DeclarativeArtifact(
            artifact=blob_artifact,
            url=blob_url,
            relative_path=digest,
            remote=self.remote,
            deferred_download=deferred_download and self.deferred_download,
        )
        blob_dc = DeclarativeContent(content=blob, d_artifacts=[da])

        return blob_dc

    async def create_signatures(self, man_dc, signature_source):
        """
        Create signature declarative artifacts.

        Signatures are currently supported only for image manifests and not manifest lists.

        For sigstore, signatures are downloaded from a specified url. The number of signatures is
        unknown upfront, need to download them in order by incrementing the index until hit 404.

        Args:
            man_dc: Declarative content instance of a related Manifest
            signature_source: Source where to get signatures from
        """
        signature_dcs = []
        if signature_source == SIGNATURE_SOURCE.SIGSTORE:
            man_digest_reformatted = man_dc.content.digest.replace(":", "=")
            sig_relative_baseurl = f"{self.remote.upstream_name}@{man_digest_reformatted}"

            signature_counter = 1
            while True:
                signature_url = urlpath_sanitize(
                    self.remote.sigstore, sig_relative_baseurl, f"signature-{signature_counter}"
                )
                signature_downloader = self.remote.get_noauth_downloader(url=signature_url)
                try:
                    signature_download_result = await signature_downloader.run()
                except aiohttp.client_exceptions.ClientResponseError as exc:
                    if exc.status != 404:
                        log.info(
                            _(
                                "{} is not accessible, can't sync an image signature. "
                                "Error: {} {}".format(signature_url, exc.status, exc.message)
                            )
                        )
                    # 404 is fine, it means there are no or no more signatures available
                    break

                with open(signature_download_result.path, "rb") as f:
                    signature_raw = f.read()

                signature_counter += 1
                sig_dc = self._create_signature_declarative_content(signature_raw, man_dc)
                if sig_dc:
                    signature_dcs.append(sig_dc)

        elif signature_source == SIGNATURE_SOURCE.API_EXTENSION:
            signatures_url = urlpath_sanitize(
                self.remote.url,
                "extensions/v2",
                self.remote.upstream_name,
                "signatures",
                man_dc.content.digest,
            )
            signatures_downloader = self.remote.get_downloader(url=signatures_url)
            await signatures_downloader.run()
            with open(signatures_downloader.path) as signatures_fd:
                api_extension_signatures = json.loads(signatures_fd.read())
            for signature in api_extension_signatures.get("signatures", []):
                if (
                    signature.get("schemaVersion") == SIGNATURE_API_EXTENSION_VERSION
                    and signature.get("type") == SIGNATURE_TYPE.ATOMIC_SHORT
                ):
                    signature_base64 = signature.get("content")
                    if signature_base64 is None:
                        continue
                    signature_raw = base64.b64decode(signature_base64)
                    sig_dc = self._create_signature_declarative_content(
                        signature_raw, man_dc, signature.get("name"), signature_base64
                    )
                    if sig_dc:
                        signature_dcs.append(sig_dc)

        return signature_dcs

    def _include_layer(self, layer):
        """
        Decide whether to include a layer.

        Args:
            layer (dict): Layer reference.

        Returns:
            bool: True when the layer should be included.

        """
        foreign_excluded = not self.remote.include_foreign_layers
        layer_type = layer.get("mediaType", MEDIA_TYPE.REGULAR_BLOB)
        is_foreign = layer_type in (MEDIA_TYPE.FOREIGN_BLOB, MEDIA_TYPE.FOREIGN_BLOB_OCI)
        if is_foreign and foreign_excluded:
            log.debug(_("Foreign Layer: %(d)s EXCLUDED"), dict(d=layer))
            return False
        return True

    def _calculate_digest(self, manifest):
        """
        Calculate the requested digest of the ImageManifest, given in JSON.

        Args:
            manifest (str):  The raw JSON representation of the Manifest.

        Returns:
            str: The digest of the given ImageManifest

        """
        decoded_manifest = json.loads(manifest)
        if "signatures" in decoded_manifest:
            # This manifest contains signatures. Unfortunately, the Docker manifest digest
            # is calculated on the unsigned version of the Manifest so we need to remove the
            # signatures. To do this, we will look at the 'protected' key within the first
            # signature. This key indexes a (malformed) base64 encoded JSON dictionary that
            # tells us how many bytes of the manifest we need to keep before the signature
            # appears in the original JSON and what the original ending to the manifest was after
            # the signature block. We will strip out the bytes after this cutoff point, add back the
            # original ending, and then calculate the sha256 sum of the transformed JSON to get the
            # digest.
            protected = decoded_manifest["signatures"][0]["protected"]
            # Add back the missing padding to the protected block so that it is valid base64.
            protected = self._pad_unpadded_b64(protected)
            # Now let's decode the base64 and load it as a dictionary so we can get the length
            protected = base64.b64decode(protected)
            protected = json.loads(protected)
            # This is the length of the signed portion of the Manifest, except for a trailing
            # newline and closing curly brace.
            signed_length = protected["formatLength"]
            # The formatTail key indexes a base64 encoded string that represents the end of the
            # original Manifest before signatures. We will need to add this string back to the
            # trimmed Manifest to get the correct digest. We'll do this as a one liner since it is
            # a very similar process to what we've just done above to get the protected block
            # decoded.
            signed_tail = base64.b64decode(self._pad_unpadded_b64(protected["formatTail"]))
            # Now we can reconstruct the original Manifest that the digest should be based on.
            manifest = manifest[:signed_length] + signed_tail

        return "sha256:{digest}".format(digest=hashlib.sha256(manifest).hexdigest())

    def _pad_unpadded_b64(self, unpadded_b64):
        """
        Fix bad padding.

        Docker has not included the required padding at the end of the base64 encoded
        'protected' block, or in some encased base64 within it. This function adds the correct
        number of ='s signs to the unpadded base64 text so that it can be decoded with Python's
        base64 library.

        Args:
            unpadded_b64 (str): The unpadded base64 text.

        Returns:
            str: The same base64 text with the appropriate number of ='s symbols.

        """
        # The Pulp team has not observed any newlines or spaces within the base64 from Docker, but
        # Docker's own code does this same operation so it seemed prudent to include it here.
        # See lines 167 to 168 here:
        # https://github.com/docker/libtrust/blob/9cbd2a1374f46905c68a4eb3694a130610adc62a/util.go
        unpadded_b64 = unpadded_b64.replace("\n", "").replace(" ", "")
        # It is illegal base64 for the remainder to be 1 when the length of the block is
        # divided by 4.
        if len(unpadded_b64) % 4 == 1:
            raise ValueError(_("Invalid base64: {t}").format(t=unpadded_b64))
        # Add back the missing padding characters, based on the length of the encoded string
        paddings = {0: "", 2: "==", 3: "="}
        return unpadded_b64 + paddings[len(unpadded_b64) % 4]


class InterrelateContent(Stage):
    """
    Stage for relating Content to other Content.
    """

    async def run(self):
        """
        Relate each item in the input queue to objects specified on the DeclarativeContent.
        """
        async for batch in self.batches():
            manifest_to_list_list = []
            blob_list = []
            config_blob_list = []
            for dc in batch:
                if dc.extra_data.get("relation"):
                    manifest_to_list = self.relate_manifest_to_list(dc)
                    manifest_to_list_list.append(manifest_to_list)
                elif dc.extra_data.get("blob_relation"):
                    blob = self.relate_blob(dc)
                    blob_list.append(blob)
                elif dc.extra_data.get("config_relation"):
                    config_blob = self.relate_config_blob(dc)
                    config_blob_list.append(config_blob)

            await sync_to_async(ManifestListManifest.objects.bulk_create)(
                objs=manifest_to_list_list, ignore_conflicts=True, batch_size=1000
            )
            await sync_to_async(BlobManifest.objects.bulk_create)(
                objs=blob_list, ignore_conflicts=True, batch_size=1000
            )
            await sync_to_async(Manifest.objects.bulk_update)(
                objs=config_blob_list, fields=["config_blob"], batch_size=1000
            )

            for dc in batch:
                await self.put(dc)

    def relate_config_blob(self, dc):
        """
        Relate a Blob to a Manifest as a config layer.

        Args:
            dc (pulpcore.plugin.stages.DeclarativeContent): dc for a Blob

        Returns:
            pulp_container.app.models.Manifest: An existing Manifest object with the updated
                config layer

        """
        configured_dc = dc.extra_data.get("config_relation")
        configured_dc.content.config_blob = dc.content
        return configured_dc.content

    def relate_blob(self, dc):
        """
        Relate a Blob to a Manifest.

        Args:
            dc (pulpcore.plugin.stages.DeclarativeContent): dc for a Blob

        Returns:
            pulp_container.app.models.BlobManifest: A new BlobManifest object with the added
                reference to a Manifest object

        """
        related_dc = dc.extra_data.get("blob_relation")
        return BlobManifest(manifest=related_dc.content, manifest_blob=dc.content)

    def relate_manifest_to_list(self, dc):
        """
        Relate an ImageManifest to a ManifestList.

        Args:
            dc (pulpcore.plugin.stages.DeclarativeContent): dc for a ImageManifest

        Returns:
            pulp_container.app.models.ManifestListManifest: A new ManifestListManifest object
                with the associated ImageManifest object and corresponding platform information

        """
        related_dc = dc.extra_data.get("relation")
        platform = dc.extra_data.get("platform")
        return ManifestListManifest(
            manifest_list=dc.content,
            image_manifest=related_dc.content,
            architecture=platform["architecture"],
            os=platform["os"],
            features=platform.get("features"),
            variant=platform.get("variant"),
            os_version=platform.get("os.version"),
            os_features=platform.get("os.features"),
        )
