########################################################################
# Copyright (C) 2023 VMware, Inc.
# All Rights Reserved
########################################################################
#

import logging

from .. import Downloader
from .. import Vib
from ..Errors import PartialVibDownloadError
from .ArFile import PartialArFileVib

logger = logging.getLogger(__name__)

class PartialVibDownloader(object):
   """Class to download partial VIB from the URL or zip path to the local
      filesystem.
            Parameters:
            * url - The source download location for the VIB, either a URL or
                    zip path.
            * downloadPath - The path where the VIB will be downloaded
            * payloadsToDownload - List of quickpatch payload names to be
                                   downloaded.
   """
   def __init__(self, url, downloadPath, payloadsToDownload):
      self.url = url
      self.downloadPath = downloadPath
      self.payloadsToDownload = payloadsToDownload

   def GetPartialVib(self):
      """Download the partial VIB from the URL or zip path to the local
         filesystem.
            Returns: The absolute VIB path of the local file.
      """
      logger.info("Downloading partial VIB: %s to %s location", self.url,
                  self.downloadPath)
      fp = None
      try:
         d = Downloader.Downloader(self.url)
         fp = d.Open()
         parArFileVib = PartialArFileVib(fp, self.payloadsToDownload)
         partialVib = Vib.ArFileVib.GetVibFromPartialArFileVib(parArFileVib)
         # Skip the desc.xml and sign and send only payloads
         partialVib.WriteVibFile(self.downloadPath,
                                 parArFileVib.vibContents[2:])
         vib = Vib.ArFileVib.FromFile(self.downloadPath)
         vib._isPartialVib = partialVib._isPartialVib
         # Get the script payload and remove it from self.payloadsToDownload
         # for later comparison.
         scriptPayloadToDownload = {p.name for p in vib.payloads if
                                    p.isquickpatch}
         # Script payload will always be downloaded.
         patchPayloadsToDownload = set(self.payloadsToDownload) - \
                                   scriptPayloadToDownload
         # Check if all patch payloads have files downloaded.
         # Patch payloads will not be needed for scan script execution, and
         # only the script payload will be downloaded.
         vib._hasAllQpPatchPayloadFiles = len(patchPayloadsToDownload) and \
            {p.name for p in vib.payloads if p.overlayorder} == \
            patchPayloadsToDownload
         logger.info("Sucessfully downloaded payload(s) %s into partial VIB",
                     self.payloadsToDownload)
      except Exception as e:
         msg = "Failed to download partial VIB. Error: %s" % (str(e))
         raise PartialVibDownloadError(self.url, msg)
      finally:
         if fp is not None:
            fp.close()
      return vib
