# Copyright (c) 2021-2023 VMware, Inc. All rights reserved.
# VMware Confidential

"""
This module manages a consolidated depot, created from the effective micro
depots from vLCM,  for the ESXi host and the ESXio-based DPUs.
"""

import logging
import os
import shutil
import zipfile

from ..Bulletin import ComponentCollection
from ..Depot import DepotIndex, MetadataNode, VendorIndex, VibDownloader
from ..DepotCollection import DepotCollection
from ..Downloader import Downloader
from ..Errors import DpuInfoError, MetadataNotFoundError, VibDownloadError
from ..Utils import XmlUtils
from ..Vib import SoftwarePlatform

HOST_DEPOT = 'hostdepot'
DEPOT_ROOT = os.path.join(os.path.sep, 'usr', 'lib', 'vmware',
                          'hostd', 'docroot', HOST_DEPOT)
OSDATA_PATH = os.path.join(os.path.sep, 'var', 'lib', 'vmware', 'osdata')
TMP_ZIP_FILE = os.path.join(DEPOT_ROOT, 'tmpZipFile')
INDEX_FILE = os.path.join(DEPOT_ROOT, 'index.xml')

etree = XmlUtils.FindElementTree()
log = logging.getLogger(__name__)

def _symlinkDepotDir(depotDir):
   """ Create symlink for host depot at new place.
   """
   if os.path.exists(DEPOT_ROOT):
      os.remove(DEPOT_ROOT)
   os.symlink(depotDir, DEPOT_ROOT)

def _modifyMetaUrl(dc):
   """ Make base name of metadata absurl and url the same so metadata
       path in vendor-index.xml matches the download path.
   """
   for ch in dc._channels.values():
      for meta in ch.metadatas:
         absurl = meta.absurl
         fileName = absurl[absurl.rfind('/') + 1:]
         meta.url = fileName

def _mergeIndexXML(xml, indexXml):
   """ Merge the vendor list to main index XML.
   """
   if indexXml is None:
      return xml
   else:
      vendors = xml.findall('vendor')
      for v in vendors:
         indexXml.append(v)
      return indexXml

def _prefixRelativePath(indexXml, depotName):
   """ Add the depot name into the vendor files' relative path.
   """
   if indexXml is not None:
      for vendor in list(indexXml):
         for elm in list(vendor):
            if elm.tag == 'relativePath':
               if elm.text is None:
                  elm.text = depotName
               else:
                  elm.text = os.path.join(depotName, elm.text)

class HostDepot(object):
   """ The consolidated depot for this ESXi host and its ESXio-based DPUs.
   """

   def _createHostDepotDir(self):
      """ Helper function to create host depot dir.
      """

      # Use OSData as storage
      depotDir = os.path.join(OSDATA_PATH, HOST_DEPOT)
      if os.path.exists(depotDir):
         if os.path.isdir(depotDir):
            _symlinkDepotDir(depotDir)
            return
         else:
            os.remove(depotDir)

      try:
         os.mkdir(depotDir)
      except Exception as e:
         log.exception('Failed to create %s: %s', depotDir, str(e))
      _symlinkDepotDir(depotDir)

   def __init__(self):
      """ Construct host depot: create an empty host depot dir.
      """
      self._extraVibSpecs = [] # name:version of standalone VIBs
      self._createHostDepotDir()

   @property
   def extraVibSpecs(self):
      """name:version specs of standalone VIBs.
      """
      return self._extraVibSpecs

   def cleanHostDepot(self):
      """ Remove all content from host depot. Recreate the dir.
      """
      try:
         realDepotPath = os.path.realpath(DEPOT_ROOT)
         if os.path.exists(realDepotPath):
            if os.path.isdir(realDepotPath):
               shutil.rmtree(realDepotPath)
            else:
               os.remove(realDepotPath)
      except Exception as e:
         log.exception('Failed to remove dir %s: %s', DEPOT_ROOT, str(e))
      self._createHostDepotDir()

   def createHostDepot(self, zipLocation):
      """ Create a clean host depot for the depot content from the provided
          depot zip file.
      """
      log.info('Creating host depot from: %s', zipLocation)
      self.cleanHostDepot()
      oldCwd = os.getcwd()
      try:
         os.chdir(os.path.realpath(DEPOT_ROOT))
         tmpZipFile = shutil.copyfile(zipLocation,
                                      os.path.join(os.path.realpath(DEPOT_ROOT),
                                                   'tmpZip'))

         tmpZip = zipfile.ZipFile(tmpZipFile)
         tmpZip.extractall()
         os.remove(tmpZipFile)
      except Exception as e:
         log.exception('Failed to unzip file %s: %s', tmpZipFile, str(e))
      finally:
         os.chdir(oldCwd)
      log.info('Host depot is created at: %s', os.path.realpath(DEPOT_ROOT))

   def _downloadBundledDepot(self, url, depotName):
      """ Download the content of a local or remote zip bundle. Return the XML
          object of index.xml.

          Parameters:
             url: The bundle file path or url.
             depotName: The depot name.
      """
      depotRoot = os.path.join(DEPOT_ROOT, depotName)
      if not os.path.exists(url):
         # Remote zip file.
         depotZipFile = os.path.join(depotRoot, os.path.basename(url))
         d = Downloader(url, depotZipFile)
         depotZipFile = d.Get()
      else:
         depotZipFile = url

      try:
         depotZip = zipfile.ZipFile(depotZipFile)
         depotZip.extractall(depotRoot)
      except Exception as e:
         log.exception('Failed to unzip file %s: %s', TMP_ZIP_FILE, str(e))
         raise

      if depotZipFile != url:
         try:
            os.remove(depotZipFile)
         except Exception as e:
            log.warning('Failed to remove tmp zip file %s: %s',
                        depotZipFile, str(e))

      indexFile = os.path.join(depotRoot, 'index.xml')
      if os.path.exists(indexFile):
         with open(indexFile) as f:
            indexXml = XmlUtils.ParseXMLFromString(f.read())
            _prefixRelativePath(indexXml, depotName)
            return indexXml
      raise MetadataNotFoundError('Depot has no index.xml.')

   def _downloadUnbundledDepot(self, url, depotName, compSpecs=None):
      """ Download an online depot or a unbundle local depot defined
          by its index.xml.

          Parameters:
             url      : The index xml file path or url.
             depotName: The depot name.
             compSpecs: The selected components whose ESXio vib payload will
                        be downloaded.
      """
      dc = DepotCollection()
      dc.ConnectDepots([url])
      _modifyMetaUrl(dc)

      depotRoot = os.path.join(DEPOT_ROOT, depotName)
      for ch in dc.channels.values():
         chDir = os.path.join(depotRoot, ch.vendorindex.relativePath)
         os.makedirs(chDir, exist_ok=True)

         # Download vendor index files.
         fileName = os.path.basename(ch.vendorIndexUrl)
         fileName = os.path.join(chDir, fileName)
         with open(fileName, 'wb') as vendorFile:
            vendorFile.write(ch.vendorindex.ToString())

         # Download metadata and collect components.
         selectedComps = ComponentCollection()
         for meta in ch.metadatas:
            fileName = os.path.join(chDir, meta.url)
            meta.WriteMetadataZip(os.path.join(chDir, fileName))

            comps = ComponentCollection(meta.bulletins, True)
            if compSpecs:
               for cid in compSpecs:
                  if isinstance(cid, str):
                     if comps.HasComponent(cid):
                        cs = comps.GetComponents(cid)
                        for c in cs:
                           selectedComps.AddComponent(c)
                  elif comps.HasComponent(*cid):
                     selectedComps.AddComponent(comps.GetComponent(*cid))
            else:
               selectedComps += comps

         selectedVibs = selectedComps.GetVibCollection(dc.vibs)
         selectedVibs = selectedVibs.GetVibsForSoftwarePlatform(
                           SoftwarePlatform.PRODUCT_ESXIO_ARM)

         # Download vibs.
         for vibid in selectedVibs:
            localfn = os.path.join(chDir, selectedVibs[vibid].GetRelativePath())
            try:
               VibDownloader(localfn, selectedVibs[vibid])
            except Exception as e:
               vibLocation = ', '.join(selectedVibs[vibid].remotelocations)
               log.error("Failed to download the vib %s: %s", vibLocation, str(e))
               raise VibDownloadError('', vibLocation, str(e))

      # Create index XML.
      indexXml = None
      for depot in dc.depots:
         xml = depot.ToXml()
         indexXml = _mergeIndexXML(xml, indexXml)

      _prefixRelativePath(indexXml, depotName)
      return indexXml

   def _getStandaloneVibDepot(self, vibUrls, depotName):
      """ Form a depot to contain all standalone VIBs.
      """
      METADATA = 'metadata.zip'
      VENDOR_INDEX = 'vendor-index.xml'

      depotRoot = os.path.join(DEPOT_ROOT, depotName)
      os.makedirs(depotRoot, exist_ok=True)
      self._extraVibSpecs.clear()

      # Create metadata
      meta = MetadataNode(url=METADATA)
      # Platform and version of depot do not matter when the depot is consumed
      # only by DPU.
      meta.AddPlatform([SoftwarePlatform.PRODUCT_ESXIO_ARM], '8.0')
      vibNum = 1
      for vibUrl in vibUrls:
         log.info('Downloading standalone VIB %s into host depot', vibUrl)

         tmpVibPath = os.path.join(depotRoot, 'vib%u.vib' % vibNum)
         vibNum += 1

         # Download VIB and move to proper location.
         try:
            d = Downloader(vibUrl, tmpVibPath)
            vibPath = d.Get()
            vibObj = meta.vibs.AddVibFromVibfile(vibPath)
            vibObj.relativepath = vibObj.GetRelativePath()

            self._extraVibSpecs.append('%s:%s'
                  % (vibObj.name, vibObj.version.versionstring))

            vibFinalPath = os.path.join(depotRoot, vibObj.relativepath)
            os.makedirs(os.path.dirname(vibFinalPath), exist_ok=True)
            shutil.copy2(vibPath, vibFinalPath)
         finally:
            if os.path.isfile(tmpVibPath):
               # Remove original file only when tmp location was used.
               os.remove(tmpVibPath)
      metaPath = os.path.join(depotRoot, METADATA)
      meta.WriteMetadataZip(metaPath)

      # Create and write vendor index
      vendorIndex = VendorIndex(name='VMware', code='vmw',
                                indexfile=VENDOR_INDEX,
                                relativePath=depotName + os.sep,
                                children=[meta])
      xmlPath = os.path.join(depotRoot, VENDOR_INDEX)
      with open(xmlPath, 'wb') as f:
         f.write(vendorIndex.ToString())

      # Create vendor index object, no need to write
      depotIndex = DepotIndex(children=[vendorIndex])
      indexXml = XmlUtils.ParseXMLFromString(depotIndex.ToString())
      return indexXml

   def _downloadDepots(self, depotList, compSpecs=None, extraVibs=None):
      """ Download the content of all depots provided by URLs or local zip file
          into host depot.

          Parameters:
             depotList: The list of depot/bundle path or url.
             compSpecs: The selected components whose ESXio vib payload will
                        be downloaded.
             extraVibs: Standalone VIB URLs/paths, to be included into the
                        depot.
      """
      indexXml = None
      Downloader.setEsxupdateFirewallRule('true')
      try:
         depotNum = 0
         for url in depotList:
            depotNum += 1
            depotName = 'depot' + str(depotNum)
            log.info('Downloading %s into host depot', url)

            if url.lower().endswith('.zip'):
               indexXmlForDepot = self._downloadBundledDepot(url, depotName)
            else:
               indexXmlForDepot = self._downloadUnbundledDepot(url, depotName,
                                                               compSpecs)

            indexXml = _mergeIndexXML(indexXmlForDepot, indexXml)

         if extraVibs:
            # Create a depot for standalone VIBs.
            depotNum = len(depotList) + 1
            indexXmlForVibs = self._getStandaloneVibDepot(extraVibs,
                  'depot%u' % depotNum)
            indexXml = _mergeIndexXML(indexXmlForVibs, indexXml)
      finally:
         Downloader.setEsxupdateFirewallRule('false')

      if indexXml is not None:
         try:
            with open(INDEX_FILE, 'wb') as f:
               f.write(etree.tostring(indexXml))
         except Exception as e:
            log.exception('Failed to write consolidated index.xml file: %s',
                          str(e))
            raise
      else:
         log.warning('Nothing downloaded into consolidated host depot.')

   def createHostDepotFromList(self, depotList, compSpecs=None, extraVibs=None):
      """ Create a clean host depot for depot content from the provided
          depot URLs.

          Parameters:
             depotList: The list of depot/bundle path or url.
             compSpecs: The selected components whose ESXio vib payload will
                        be downloaded.
             extraVibs: Standalone VIB URLs/paths, to be included into the
                        depot.
      """
      self.cleanHostDepot()
      self._downloadDepots(depotList, compSpecs, extraVibs)

def _getLocalEndpointAddress(dpu):
   """Returns local endpoint address to the DPU.
   """
   if dpu and isinstance(dpu, dict):
      return dpu.get('Local Endpoint Address')
   return None

def getHostDepotURLs(dpus):
   """ Generate host depot URLs with the IP for the DPU endpoint address.
   """
   depotURLs = []
   for dpu in dpus:
      endPoint = _getLocalEndpointAddress(dpu)
      if endPoint:
         url = ''.join(['http://', endPoint, '/', HOST_DEPOT, '/index.xml'])
         depotURLs.append(url)
   return depotURLs or None

def hostTextFile(text, fileName, dpu):
   """Hosts a single text file in host depot folder, returns URL for DPU.
   """
   if dpu is None:
      raise ValueError('DPU instance must be provided')
   endPoint = _getLocalEndpointAddress(dpu)
   if endPoint:
      filePath = os.path.join(DEPOT_ROOT, fileName)

      # Check if the file with the spec already exists to avoid writing the same
      # content for each DPU.
      if not os.path.exists(filePath):
         with open(filePath, 'w') as f:
            f.write(text)
      return ''.join(['http://', endPoint, '/', HOST_DEPOT, '/', fileName])
   raise DpuInfoError('Local endpoint address is not present in DPU info')
