import base64
from datetime import datetime
import json
import logging
import uuid
import tempfile

import proton

from .exceptions import SigningError
from .utils.misc import (
    run_entrypoint,
    log_step,
    get_pyxis_ssl_paths,
)
from .quay_client import QuayClient
from .manifest_claims_handler import ManifestClaimsHandler

LOG = logging.getLogger("pubtools.quay")


class SignatureHandler:
    """Base class implementing operations common for container and operator signing."""

    MAX_MANIFEST_DIGESTS_PER_SEARCH_REQUEST = 50

    def __init__(self, hub, task_id, target_settings, target_name):
        """
        Initialize.

        Args:
            hub (HubProxy):
                Instance of XMLRPC pub-hub proxy.
            task_id (str):
                ID of the pub task
            target_settings (dict):
                Target settings.
            target_name (str):
                Name of the target.
        """
        self.hub = hub
        self.task_id = task_id
        self.target_settings = target_settings
        self.target_name = target_name

        # Which URL hostnames will the destination images be accessible by to customers
        self.dest_registries = target_settings["docker_settings"]["docker_reference_registry"]
        self.dest_registries = (
            self.dest_registries
            if isinstance(self.dest_registries, list)
            else [self.dest_registries]
        )

        self.quay_host = self.target_settings.get("quay_host", "quay.io").rstrip("/")
        self._src_quay_client = None
        self._dest_quay_client = None

    @property
    def src_quay_client(self):
        """Create and access QuayClient for source image."""
        if self._src_quay_client is None:
            self._src_quay_client = QuayClient(
                self.target_settings["source_quay_user"],
                self.target_settings["source_quay_password"],
                self.target_settings.get("source_quay_host") or self.quay_host,
            )
        return self._src_quay_client

    @property
    def dest_quay_client(self):
        """Create and access QuayClient for dest image."""
        if self._dest_quay_client is None:
            self._dest_quay_client = QuayClient(
                self.target_settings["dest_quay_user"],
                self.target_settings["dest_quay_password"],
                self.quay_host,
            )
        return self._dest_quay_client

    @classmethod
    def create_manifest_claim_message(
        cls, destination_repo, signature_key, manifest_digest, docker_reference, image_name, task_id
    ):
        """
        Construct a manifest claim (image signature) as well as a message to send to RADAS.

        Constructed signature adheres to the following standard:
        https://github.com/containers/image/blob/master/docs/containers-signature.5.md

        Args:
            destination_repo (str):
                Internal destination repository to send to RADAS.
            signature_key (str):
                Signature key that will be sent to RADAS.
            manifest_digest (str):
                Digest referencing the signed image. Mandatory part of the image signature.
            docker_reference (str):
                Image reference which will be used by customers to pull the image. Mandatory part of
                the image signature.
            image_name (str):
                Name of the image to send to RADAS.
            task_id (str):
                ID of the pub task.
        """
        # container image signature
        manifest_claim = {
            "critical": {
                "type": "atomic container signature",
                "image": {"docker-manifest-digest": manifest_digest},
                "identity": {"docker-reference": docker_reference},
            },
            # NOTE: pub version is no longer written here. I hope that's OK
            "optional": {"creator": "Red Hat RCM Pub"},
        }

        message = {
            "sig_key_id": signature_key,
            # Python 2.6/3 compatibility workaround
            "claim_file": base64.b64encode(json.dumps(manifest_claim).encode("latin1")).decode(
                "latin1"
            ),
            "pub_task_id": task_id,
            "request_id": str(uuid.uuid4()),
            "manifest_digest": manifest_digest,
            "repo": destination_repo,
            "image_name": image_name,
            "docker_reference": docker_reference,
            "created": datetime.utcnow().isoformat() + "Z",
        }
        return message

    def get_tagged_image_digests(self, image_ref, media_type):
        """
        Get all digests referenced by a tagged image.

        There will only be one digest in case of single-arch image (source image), or multiple
        digests for a multi-arch image.

        Args:
            image_ref (str):
                Image reference URL. Must be specified via tag.
            media_type (str):
                Requested media type of the manifest.

        Returns ([str]):
            List of manifest digests referenced by the tag.
        """
        digests = []
        manifest = self.src_quay_client.get_manifest(image_ref, media_type=media_type)
        # If V2S2 manifest, we only want its digest
        if manifest.get("mediaType") != "application/vnd.docker.distribution.manifest.list.v2+json":
            digests.append(
                self.src_quay_client.get_manifest_digest(image_ref, media_type=media_type)
            )
        # If manifest list, we want digests of all its arch images
        else:
            for arch_manifest in manifest["manifests"]:
                digests.append(arch_manifest["digest"])

        return digests

    def get_signatures_from_pyxis(self, manifest_digests=None):
        """
        Get existing signatures from Pyxis based on the specified criteria (currently only digests).

        NOTE: In the current implementation, only manifest digests are being used to search for
        existing signatures. Also, the search is performed in chunks, their size being limited by
        MAX_MANIFEST_DIGESTS_PER_SEARCH_REQUEST.

        Args:
            manifest_digests ([str]|None):
                Digests for which to return signatures.

            Yields (dict):
                Existing signatures as returned by Pyxis based on specified criteria. The returned
                sturcture is an iterator to reduce memory requirements.
        """
        cert, key = get_pyxis_ssl_paths(self.target_settings)
        chunk_size = self.MAX_MANIFEST_DIGESTS_PER_SEARCH_REQUEST
        manifest_digests = sorted(list(set(manifest_digests)))

        for chunk_start in range(0, len(manifest_digests), chunk_size):
            chunk = manifest_digests[chunk_start : chunk_start + chunk_size]  # noqa: E203

            args = ["--pyxis-server", self.target_settings["pyxis_server"]]
            args += ["--pyxis-ssl-crtfile", cert]
            args += ["--pyxis-ssl-keyfile", key]

            with tempfile.NamedTemporaryFile(
                mode="w", prefix="pubtools_quay_get_signatures_"
            ) as signature_fetch_file:
                if manifest_digests:
                    json.dump(chunk, signature_fetch_file)
                    signature_fetch_file.flush()
                    args += ["--manifest-digest", "@{0}".format(signature_fetch_file.name)]

                env_vars = {}
                chunk_results = run_entrypoint(
                    ("pubtools-pyxis", "console_scripts", "pubtools-pyxis-get-signatures"),
                    "pubtools-pyxis-get-signatures",
                    args,
                    env_vars,
                )

            for result in chunk_results:
                yield result

    def remove_duplicate_claim_messages(self, claim_messages):
        """
        Remove claim messages which could be considered 'duplicates', containing the same data.

        Args:
            claim_messages ([dict]):
                All created claim messages.

        Returns ([dict]):
            De-duplicated claim messages.
        """
        # dictionary key is a tuple of all parameters whose combination makes the message unique
        unique_message_mapping = {}
        for message in claim_messages:
            key = (
                message["sig_key_id"],
                message["claim_file"],
                message["pub_task_id"],
                message["manifest_digest"],
                message["repo"],
                message["image_name"],
                message["docker_reference"],
            )
            if key not in unique_message_mapping:
                unique_message_mapping[key] = message

        return list(unique_message_mapping.values())

    def filter_claim_messages(self, claim_messages):
        """
        Filter out the manifest claim messages which are already in the sigstore.

        Args:
            claim_messages ([dict]):
                Messages to be sent to RADAS.

        Returns ([dict]):
            Messages which don't yet exist in Pyxis.
        """
        LOG.info("Removing claim messages which already exist in Pyxis")
        digests = [message["manifest_digest"] for message in claim_messages]
        digests = sorted(list(set(digests)))

        existing_signatures = self.get_signatures_from_pyxis(manifest_digests=digests)

        signatures_by_key = {}
        for signature in existing_signatures:
            # combination of image reference, digest, and signature key makes a signature unique
            key = (signature["reference"], signature["manifest_digest"], signature["sig_key_id"])
            signatures_by_key[key] = signature

        filtered_claim_messages = []
        for message in claim_messages:
            key = (message["docker_reference"], message["manifest_digest"], message["sig_key_id"])
            # New signatures have switched to using long (16B) keys, while old signatures may still
            # contain short (8B) keys. If claim is matched with a shorter key format, it's
            # still considered a duplicate and shouldn't be uploaded again.
            old_key = None
            if len(message["sig_key_id"]) > 8:
                old_key = (
                    message["docker_reference"],
                    message["manifest_digest"],
                    message["sig_key_id"][-8:],
                )

            if key not in signatures_by_key and old_key not in signatures_by_key:
                filtered_claim_messages.append(message)

        LOG.info(
            "{0} claim messages remain after removing duplicates".format(
                len(filtered_claim_messages)
            )
        )
        return filtered_claim_messages

    def get_signatures_from_radas(self, claim_messages):
        """
        Send signature claims to RADAS via UMB and receive signed claims.

        The messaging logic is handled by the ManifestClaimsHandler class.

        Args:
            claim_messages ([dict]):
                Signature claims to be sent to RADAS.
        Returns ([dict]):
            Response messages from RADAS.
        raises MessageHandlerTimeoutException:
            If a message from RADAS hasn't arrived in time.
        """
        LOG.info("Sending claim messages to RADAS and waiting for results")
        # messages will be sent by pub-hub via XMLRPC
        # callback will be utilized by ManifestClaimsHandler, which will decide when to send msgs
        message_sender_callback = (
            lambda messages: self.hub.worker.umb_send_manifest_claim_messages(  # noqa: E731
                self.target_name, self.task_id, messages
            )
        )

        address = (
            "queue://Consumer.msg-producer-pub"
            ".{task_id}.VirtualTopic.eng.robosignatory.container.sign".format(task_id=self.task_id)
        )

        docker_settings = self.target_settings["docker_settings"]
        claims_handler = ManifestClaimsHandler(
            umb_urls=docker_settings["umb_urls"],
            radas_address=docker_settings.get("umb_radas_address", address),
            claim_messages=claim_messages,
            pub_cert=docker_settings.get("umb_pub_cert", "/etc/pub/umb-pub-cert-key.pem"),
            ca_cert=docker_settings.get("umb_ca_cert", "/etc/pki/tls/certs/ca-bundle.crt"),
            timeout=docker_settings.get("umb_signing_timeout", 600),
            throttle=docker_settings.get("umb_signing_throttle", 100),
            retry=docker_settings.get("umb_signing_retry", 3),
            message_sender_callback=message_sender_callback,
        )
        container = proton.reactor.Container(claims_handler)
        container.run()

        return claims_handler.received_messages

    def upload_signatures_to_pyxis(self, claim_mesages, signature_messages):
        """
        Upload signatures to Pyxis by using a pubtools-pyxis entrypoint.

        Data required for a Pyxis POST request:
        - manifest_digest
        - reference
        - repository
        - sig_key_id
        - signature_data

        Signatures are uploaded in batches.

        Args:
            claim_messages ([dict]):
                Signature claim messages constructed for the RADAS service.
            signature_messages ([dict]):
                Messages from RADAS containing image signatures.
        """
        LOG.info("Sending new signatures to Pyxis")

        cert, key = get_pyxis_ssl_paths(self.target_settings)

        signatures = []
        claim_messages_by_id = dict((m["request_id"], m) for m in claim_mesages)
        sorted_signature_messages = sorted(signature_messages, key=lambda msg: msg["request_id"])

        for signature_message in sorted_signature_messages:
            claim_message = claim_messages_by_id[signature_message["request_id"]]

            signatures.append(
                {
                    "manifest_digest": signature_message["manifest_digest"],
                    "reference": claim_message["docker_reference"],
                    "repository": claim_message["image_name"],
                    "sig_key_id": claim_message["sig_key_id"],
                    "signature_data": signature_message["signed_claim"],
                }
            )

        args = ["--pyxis-server", self.target_settings["pyxis_server"]]
        args += ["--pyxis-ssl-crtfile", cert]
        args += ["--pyxis-ssl-keyfile", key]

        with tempfile.NamedTemporaryFile(
            mode="w", prefix="pubtools_quay_upload_signatures_"
        ) as signature_file:
            json.dump(signatures, signature_file)
            signature_file.flush()
            args += ["--signatures", "@{0}".format(signature_file.name)]

            LOG.info("Uploading {0} new signatures".format(len(signatures)))
            env_vars = {}
            run_entrypoint(
                ("pubtools-pyxis", "console_scripts", "pubtools-pyxis-upload-signatures"),
                "pubtools-pyxis-upload-signature",
                args,
                env_vars,
            )

    def validate_radas_messages(self, claim_messages, signature_messages):
        """
        Check if messages received from RADAS contain any errors.

        Args:
            claim_messages ([dict]):
                Messages sent to RADAS.
            signature_messages ([dict]):
                Messages received from RADAS.

        Raises:
            SigningError:
                If RADAS messages contain errors.
        """
        failed_messages = 0
        for message in signature_messages:
            if message["errors"]:
                failed = [m for m in claim_messages if m["request_id"] == message["request_id"]][0]
                LOG.error(
                    "Signing of claim message {0} failed with following errors: {1}".format(
                        failed, message["errors"]
                    )
                )
                failed_messages += 1

        if failed_messages:
            raise SigningError(
                "Signing of {0}/{1} messages has failed".format(
                    failed_messages, len(claim_messages)
                )
            )


class ContainerSignatureHandler(SignatureHandler):
    """Class for handling the signing of container images."""

    def construct_item_claim_messages(self, push_item):
        """
        Construct all the signature claim messages for RADAS for one push item.

        push_item (ContainerPushItem):
            Container push item whose claim messages will be created.
        Returns ([dict]):
            Claim messages for a given push item.
        """
        LOG.info("Constructing claim messages for push item '{0}'".format(push_item))
        claim_messages = []

        if push_item.claims_signing_key:
            if (
                QuayClient.MANIFEST_LIST_TYPE
                in push_item.metadata["build"]["extra"]["image"]["media_types"]
            ):
                mtype = QuayClient.MANIFEST_LIST_TYPE
            elif (
                QuayClient.MANIFEST_V2S2_TYPE
                in push_item.metadata["build"]["extra"]["image"]["media_types"]
            ):
                mtype = QuayClient.MANIFEST_V2S2_TYPE
            else:
                mtype = QuayClient.MANIFEST_V2S1_TYPE
            digests = self.get_tagged_image_digests(push_item.metadata["pull_url"], mtype)
            # each image digest needs its own signature
            for digest in digests:
                # each destination image reference needs its own signature
                for repo, tags in sorted(push_item.metadata["tags"].items()):
                    for tag in tags:
                        claim_messages += self.construct_variant_claim_messages(
                            repo, tag, digest, [push_item.claims_signing_key]
                        )

        return claim_messages

    def construct_variant_claim_messages(self, repo, tag, digest, signing_keys):
        """
        Construct claim messages for all specified variations of a given image.

        The variations are customer visible destination registry and signing key.

        Args:
            repo (str):
                Destination external repository  of a pushed image.
            tag: (str):
                Destination tag of a pushed image
            digest (str):
                Digest of the pushed image.
            signing_keys ([str]):
                Signing keys to construct the signatures with.

        Returns ([dict]):
            Signature claim messages to send to RADAS.
        """
        claim_messages = []
        image_schema = "{host}/{repository}:{tag}"

        for registry in self.dest_registries:
            reference = image_schema.format(host=registry, repository=repo, tag=tag)

            for signing_key in signing_keys:
                claim_message = self.create_manifest_claim_message(
                    destination_repo=repo,
                    signature_key=signing_key,
                    manifest_digest=digest,
                    docker_reference=reference,
                    image_name=repo,
                    task_id=self.task_id,
                )
                claim_messages.append(claim_message)

        return claim_messages

    @log_step("Sign container images")
    def sign_container_images(self, push_items):
        """
        Perform all the steps needed to sign the images of specified push items.

        The workflow can be summarized as:
        - create manifest claim messages for all items, registries, keys, digests, repos, and tags
        - filter out requests for signatures which are already in Pyxis
        - send messages to RADAS and receive signatures (ManifestClaimsHandler class)
        - Upload new signatures to Pyxis

        Args:
            push_items (([ContainerPushItem])):
                Container push items whose images will be signed.
        Returns (list(dict)):
            List of claim message sent for signing
        """
        if not self.target_settings["docker_settings"].get(
            "docker_container_signing_enabled", False
        ):
            LOG.info("Container signing not allowed in target settings, skipping.")
            return []

        claim_messages = []
        for item in push_items:
            claim_messages += self.construct_item_claim_messages(item)
        claim_messages = self.remove_duplicate_claim_messages(claim_messages)
        claim_messages = self.filter_claim_messages(claim_messages)
        if not claim_messages:
            LOG.info("No new claim messages will be uploaded")
            return []

        LOG.info("{0} claim messages will be uploaded".format(len(claim_messages)))
        signature_messages = self.get_signatures_from_radas(claim_messages)
        self.validate_radas_messages(claim_messages, signature_messages)
        self.upload_signatures_to_pyxis(claim_messages, signature_messages)
        return claim_messages

    @log_step("Sign container images")
    def sign_container_images_new_digests(self, push_items):
        """
        Sign digests for manifest types which haven't existed before the push.

        Manifest types which are not pushed in the push items can be fetched
        from quay and then created dynamically on the fly. This method is meant
        to create signatures for those.

        Args:
            push_items (([ContainerPushItem])):
                Container push items whose images will be signed.
        Returns (list(dict)):
            List of claim message sent for signing
        """
        if not self.target_settings["docker_settings"].get(
            "docker_container_signing_enabled", False
        ):
            LOG.info("Container signing not allowed in target settings, skipping.")
            return []

        claim_messages = []
        for item in push_items:
            for (repo, tag), digests_per_type in sorted(item.metadata["new_digests"].items()):
                for digest in digests_per_type.values():
                    claim_messages += self.construct_variant_claim_messages(
                        repo, tag, digest, [item.claims_signing_key]
                    )
        claim_messages = self.remove_duplicate_claim_messages(claim_messages)
        claim_messages = self.filter_claim_messages(claim_messages)
        if not claim_messages:
            LOG.info("No new claim messages will be uploaded")
            return []

        LOG.info("{0} claim messages will be uploaded".format(len(claim_messages)))
        signature_messages = self.get_signatures_from_radas(claim_messages)
        self.validate_radas_messages(claim_messages, signature_messages)
        self.upload_signatures_to_pyxis(claim_messages, signature_messages)
        return claim_messages


class OperatorSignatureHandler(SignatureHandler):
    """Class for handling the signing of index images."""

    def construct_index_image_claim_messages(self, index_image, tags, signing_keys):
        """
        Construct signature claim messages for RADAS for the specified index image.

        index_image (str):
            Reference to a new index image constructed by IIB.
        tags (str):
            List of tags to be signed.
        signing_keys (str):
            Signing keys to be used for signing.

        Returns ([dict]):
            Structured messages to be sent to UMB.
        """
        LOG.info("Constructing claim messages for index image '{0}'".format(index_image))
        claim_messages = []
        image_schema = "{host}/{repository}:{tag}"

        # Get digests of all archs this index image was build for
        index_image_credential = self.target_settings["iib_overwrite_from_index_token"].split(":")
        index_image_quay_client = QuayClient(
            index_image_credential[0],
            index_image_credential[1],
            self.quay_host,
        )
        manifest_list = index_image_quay_client.get_manifest(
            index_image, media_type=QuayClient.MANIFEST_LIST_TYPE
        )
        digests = [m["digest"] for m in manifest_list["manifests"]]
        for registry in self.dest_registries:
            for signing_key in signing_keys:
                if not signing_key:
                    continue
                for digest in digests:
                    repo = self.target_settings["quay_operator_repository"]
                    for _tag in tags:
                        reference = image_schema.format(host=registry, repository=repo, tag=_tag)
                        claim_message = self.create_manifest_claim_message(
                            destination_repo=repo,
                            signature_key=signing_key,
                            manifest_digest=digest,
                            docker_reference=reference,
                            image_name=repo,
                            task_id=self.task_id,
                        )
                        claim_messages.append(claim_message)

        return claim_messages

    @log_step("Sign operator images")
    def sign_operator_images(self, iib_results, tag_suffix):
        """
        Perform all the steps needed to sign the newly constructed index images.

        Sigstore is not checked for existing signatures, as there's no way any could exist for a
        newly constructed image.

        Args:
            iib_results ({str:dict}):
                IIB results for each version the push was performed for.
        """
        image_schema = "{host}/{namespace}/{repo}:{tag}"
        if not self.target_settings["docker_settings"].get(
            "docker_container_signing_enabled", False
        ):
            LOG.info("Container signing not allowed in target settings, skipping.")
            return

        claim_messages = []
        for version, iib_details in sorted(iib_results.items()):
            iib_result = iib_details["iib_result"]
            signing_keys = iib_details["signing_keys"]
            # Index image used to fetch manifest list. This image will never be overwritten
            iib_namespace = iib_result.index_image_resolved.split("/")[1]
            permanent_index_image = image_schema.format(
                host=self.target_settings.get("quay_host", "quay.io").rstrip("/"),
                namespace=iib_namespace,
                repo="iib",
                tag=iib_result.build_tags[0],
            )
            # Version acts as a tag of the index image
            claim_messages += self.construct_index_image_claim_messages(
                permanent_index_image, [version, "%s-%s" % (version, tag_suffix)], signing_keys
            )

        if not claim_messages:
            LOG.info("No new claim messages will be uploaded")
            return

        signature_messages = self.get_signatures_from_radas(claim_messages)
        self.validate_radas_messages(claim_messages, signature_messages)

        self.upload_signatures_to_pyxis(claim_messages, signature_messages)
        return claim_messages

    def sign_task_index_image(self, signing_keys, index_image, tags):
        """
        Perform an alternatve signing workflow used by IIB methods in pub.

        This workflow is used by methods 'PushAddIIBBundles', 'PushRemoveIIBOperators',
        'PushIIBBuildFromScratch'.

        Args:
            signing_keys ([str]):
                Signing key to be used.
            index_image (str):
                Index image pointing to the new manifest list.
            tags ([str]):
                List of tags to sign.
        Returns ([dict]):
            Constructed claim messages.
        """
        claim_messages = self.construct_index_image_claim_messages(index_image, tags, signing_keys)
        if not claim_messages:
            LOG.info("No new claim messages will be uploaded")
            return

        signature_messages = self.get_signatures_from_radas(claim_messages)
        self.validate_radas_messages(claim_messages, signature_messages)

        self.upload_signatures_to_pyxis(claim_messages, signature_messages)

        return claim_messages


class BasicSignatureHandler(SignatureHandler):
    """Class that handles signing claims which were constructed by user."""

    def __init__(self, hub, target_settings, target_name):
        """
        Initialize.

        NOTE: "task_id" is not needed for this workflow

        Args:
            hub (HubProxy):
                Instance of XMLRPC pub-hub proxy.
            target_settings (dict):
                Target settings.
            target_name (str):
                Name of the target.
        """
        SignatureHandler.__init__(self, hub, "1", target_settings, target_name)

    def sign_claim_messages(self, claim_messages, remove_duplicates=True, filter_existing=True):
        """
        Sign claim messages that were provided by the user and upload them to Pyxis.

        Args:
            claim_messages ([dict]):
                Claim messages to be signed and uploaded.
            remove_duplicates (bool):
                Whether to check if there are any duplicates among the messages and remove them.
            filter_existing (bool):
                Whether to check if the signatures already exist in Pyxis, and only upload
                those that aren't.
        """
        if not self.target_settings["docker_settings"].get(
            "docker_container_signing_enabled", False
        ):
            LOG.info("Container signing not allowed in target settings, skipping.")
            return

        if remove_duplicates:
            claim_messages = self.remove_duplicate_claim_messages(claim_messages)
        if filter_existing:
            claim_messages = self.filter_claim_messages(claim_messages)
        if not claim_messages:
            LOG.info("No new claim messages will be uploaded")
            return

        LOG.info("{0} claim messages will be uploaded".format(len(claim_messages)))
        signature_messages = self.get_signatures_from_radas(claim_messages)
        self.validate_radas_messages(claim_messages, signature_messages)
        self.upload_signatures_to_pyxis(claim_messages, signature_messages)
