#!/usr/bin/python
########################################################################
# Copyright (C) 2010-2011, 2015-2021 VMware, Inc.                      #
# All Rights Reserved                                                  #
########################################################################

import os
import tempfile
import tarfile

from .. import Depot, Downloader, Errors, Vib
from ..ImageManager import ReservedVibTar
from ..Utils import BootCfg, EsxGzip
from ..Utils.HashedStream import HashedStream
from ..Utils.Misc import isString, seekable

"""This module provides an ancestor class for creating different image types
   based on an image profile."""

def resetFObj(fObj):
   """Resets a file object or an HashedStream object.
   """
   if isinstance(fObj, HashedStream):
      fObj.reset()
   else:
      fObj.seek(0)

def getSeekableFObj(fObj):
   """Returns a seekable file object based on the given one, a temporary
      file will be created if it is not seekable.
   """
   PAYLOAD_READ_CHUNKSIZE = 1024 * 1024

   if seekable(fObj):
      return fObj

   tfp = tempfile.TemporaryFile()
   bytesRead = fObj.read(PAYLOAD_READ_CHUNKSIZE)
   while bytesRead:
      tfp.write(bytesRead)
      bytesRead = fObj.read(PAYLOAD_READ_CHUNKSIZE)
   fObj.close()
   tfp.seek(0)

   return tfp

def createTgz(payload, payloadName, tarDest):
   """Creates a tgz containing a single file.
         Parameters:
            * payload     - File object of the payload to be tgz'ed.
            * payloadName - Name to be given to the payload in the tgz.
            * tarDest     - Complete path with file name of the tgz to
                            be created.
   """
   if isString(tarDest):
      fobj = EsxGzip.GzipFile(tarDest, "wb")
   else:
      fobj = EsxGzip.GzipFile(fileobj=tarDest, mode="wb")

   with tarfile.open(fileobj=fobj, mode="w") as depotTar:
      tarinfo = tarfile.TarInfo(payloadName)

      payload = getSeekableFObj(payload)
      payload.seek(0, 2)
      tarinfo.size = payload.tell()
      resetFObj(payload)

      depotTar.addfile(tarinfo, payload)

   fobj.close()

class ImageBuilder(object):
   """This class is a skeleton to be inherited for the different methods of
   generating an image with the contents of an image profile."""

   DATABASE_NAME = "imgdb.tgz"
   PAYLOADTAR_NAME = "imgpayld.tgz"
   # Payloads of esx-base that are not directly stored in bootbanks. They are
   # tar'ed in this file and mounted as one.
   BASE_MISC_PAYLOADTAR_NAME = 'basemisc.tgz'
   RESERVED_VIBS_TAR_NAME = 'resvibs.tgz'
   ESXIO_DEPOT_TAR_NAME = 'esxio-depot.tgz'

   def __init__(self, imageprofile):
      """Class constructor.
            Parameters:
               * imageprofile - An instance of ImageProfile. The 'vibs'
                                attribute of the object must contain valid
                                references to the VIBs in the 'vibIDs'
                                property. Those references must include either
                                a file object or a valid remote location.
      """
      self.imageprofile = imageprofile

   @staticmethod
   def _GetSeekableFileObj(url):
      # Each VIB can potentially contain more than one payload. It is likely
      # that the order of payloads within the VIB will not be the order in
      # which they will be written to the target. Therefore, the file-like
      # object from which a VIB's payloads are read must be seekable. This
      # method takes a URL or path as input and returns a seek-able file object.
      # For remote URLs, it accomplishes that by downloading to a temporary
      # file.
      if os.path.exists(url):
         # Avoid using Downloader and all the python/ssl dependencies for
         # a local path, especially during build.
         d = open(url, 'rb')
      else:
         from .. import Downloader
         d = Downloader.Downloader(url).Open()
      # We only care if something is seekable.
      if hasattr(d, "seek"):
         return d
      # Else, copy it locally.
      t = tempfile.TemporaryFile()
      data = d.read(512)
      while data:
         t.write(data)
         data = d.read(512)
      d.close()
      t.seek(0)
      return t

   @classmethod
   def _CheckVibFile(cls, vib, checkacceptance=True):
      # This method just ensures that a single VIB has an appropriate file-like
      # object behind it.

      # Should probably make public attributes/properties/methods to implement
      # this check.
      if vib._arfile is not None and vib._arfile.seekable():
         pass
      elif not vib.remotelocations:
         msg = ("VIB %s has neither a seek-able file object nor a URL "
                "location. This may indicate a problem with the depot "
                "metadata. " % vib.id)
         raise Errors.VibIOError(None, msg)
      else:
         problems = list()
         success = False
         for url in vib.remotelocations:
            try:
               vib.OpenFile(cls._GetSeekableFileObj(url))
               # We only need one. :-)
               success = True
               break
            except Exception as e:
               problems.append(str(e))

         if not success:
            msg = ("Error retrieving file for VIB '%s': %s."
                   % (vib.id, "; ".join(problems)))
            raise Errors.VibDownloadError(None, None, msg)

      if checkacceptance:
         vib.VerifyAcceptanceLevel()

   def _CheckVibFiles(self, checkacceptance=True):
      # This method makes sure that every VIB has an appropriate file object
      # behind it, so that we'll be able to read payloads into the target object.
      # It also checks to make sure that things in vibIDs have corresponding
      # objects in vibs. This is the only place where we make this check, so
      # this method should be called before _AddPayloads() or _AddDatabase().
      for vibid in self.imageprofile.vibIDs:
         try:
            vib = self.imageprofile.vibs[vibid]
         except KeyError:
            msg = "Could not find object for VIB '%s'." % vibid
            raise Errors.ProfileFormatError(self.imageprofile.name, msg)
         self._CheckVibFile(vib, checkacceptance)

      # We need to make sure that every reserved VIB has an appropriate
      # file object behind it.
      for vibid in self.imageprofile.reservedVibIDs:
         try:
            vib = self.imageprofile.reservedVibs[vibid]
         except KeyError:
            msg = "Could not find object for reserved VIB '%s'." % vibid
            raise Errors.ProfileFormatError(self.imageprofile.name, msg)
         self._CheckVibFile(vib, checkacceptance)

   def _GetBootCfg(self, installer=True, moduleroot='', isoImage=False,
                   kernelopts=None, bootbankVibOnly=False,
                   appendResVibsTgz=True, esxiodepot=None):
      '''Return BootCfg instance if boot modules is not zero, otherwise return
         None
         Parameters:
            * installer  - True if the bootcfg is for installer
            * moduleroot - root for module files
            * isoImage   - True if the bootcfg is used to build an iso
            * kernelopts - Additional kernel boot options other than
                           feature states
            * bootbankVibOnly - if True, only include bootbank VIB modules.
            * appendResVibsTgz - if True, resvibs.tgz is appended to modules.
            * esxiodepot - If not None, esxio-depot.tgz is appended to modules.
      '''
      payload_types = [Vib.Payload.TYPE_TGZ, Vib.Payload.TYPE_VGZ,
                       Vib.Payload.TYPE_BOOT]
      if installer:
         payload_types.append(Vib.Payload.TYPE_INSTALLER_VGZ)

      vib_types = [Vib.BaseVib.TYPE_BOOTBANK] if bootbankVibOnly else None
      bootorder = self.imageprofile.GetBootOrder(payload_types, vib_types)

      modules = [p.localname for (vibid, p) in bootorder]
      if not modules:
         return None

      bootcfg = BootCfg.BootCfg()

      # ESX Database
      modules.append(self.DATABASE_NAME)
      # Misc esx-base payloads.
      modules.append(self.BASE_MISC_PAYLOADTAR_NAME)
      if appendResVibsTgz:
         modules.append(self.RESERVED_VIBS_TAR_NAME)
      if esxiodepot:
         modules.append(self.ESXIO_DEPOT_TAR_NAME)
      if installer:
         bootcfg.kernelopt["runweasel"] = None
         modules.append(self.PAYLOADTAR_NAME)
         bootcfg.title = "Loading ESXi installer"
      else:
         bootcfg.kernelopt["autoPartition"] = "TRUE"
         bootcfg.title = "Loading ESXi"

      if isoImage:
         bootcfg.kernelopt["cdromBoot"] = None

      if kernelopts is not None:
         bootcfg.kernelopt.update(kernelopts)

      if moduleroot:
         modules = [os.path.join(moduleroot, module) for module in modules]

      bootcfg.kernel = modules[0]
      bootcfg.modules = modules[1:]
      bootcfg.build = self.imageprofile.GetEsxVersion()

      return bootcfg

   def _AddReservedVibs(self, reservedVibTarPath):
      """This method generates a tar file that contains reserved vibs.
         The method accepts a path as string that it uses to create the
         tar file.
      """
      reservedVibTar = ReservedVibTar.ReservedVibTar(reservedVibTarPath)
      try:
         for vibid in self.imageprofile.reservedVibIDs:
            vib = self.imageprofile.reservedVibs[vibid]
            # TODO: Not skipping tools-light vib breaks server build ISO/PXE
            # generation.
            # locker VIB is specially reserved for cases where
            # bootbank and locker are not synced, a full copy is not needed.
            if vib.vibtype == vib.TYPE_LOCKER or not vib.relativepath:
               continue

            # Windows does not allow double opening of a file, close the
            # initial descriptor and cleanup after use.
            with tempfile.NamedTemporaryFile(delete=False) as tmpfd:
               localPath = tmpfd.name
            try:
               Depot.VibDownloader(localPath, vib)
               reservedVibTar.AddVib(localPath, vib.relativepath)
            finally:
               os.remove(localPath)
      except Exception as e:
         msg = "Could not download and package reserved VIBs. %s" % str(e)
         raise Errors.VibDownloadError(None, None, msg)
      finally:
         reservedVibTar.close()

   def _AddPayloads(self, target, checkdigests=True):
      # This method adds each payload to <target>, while doing the right things
      # for various payload types.
      raise NotImplementedError("_AddPayloads is not implemented in the child"
                                " class.")

   def _AddDatabase(self, target):
      # This method generates a tar database from the image profile, writes it
      # to a temp file, then adds the temp file to <target>.
      raise NotImplementedError("_AddDatabase is not implemented in the child"
                                " class.")

   def _AddMetadataZip(self, target):
      # This method generates a metadata.zip from the image profile.
      raise NotImplementedError("_AddMetadataZip is not implemented in the"
                                " child class.")

   def _AddProfileXml(self, target):
      # Adds profile.xml to the upgrade directory in the target.
      raise NotImplementedError("_AddProfileXml is not implemented in the"
                                " child class.")

   # Not necessary for a pxeboot, but could be generated either way just incase
   # it is needed.
   def _AddBootCfg(self, target, installer=True):
      # This method populates efi/boot/boot.cfg. It is required for EFI.
      raise NotImplementedError("_AddBootCfg is not implemented in the child"
                                " class")

