# Copyright (c) 2021-2024 Broadcom. All Rights Reserved.
# Broadcom Confidential. The term "Broadcom" refers to Broadcom Inc.
# and/or its subsidiaries.

"""
This module contains utilities for managing files that are exported from a
VIB into the depot metdata.
"""

import os
import bisect
from functools import total_ordering
import shutil
import itertools

# The location of the platform locale directories.
CONFIG_LOCALE_DIR = "usr/lib/vmware/configmanager/locale"
MODULES_DIR = "usr/lib/vmware/configmanager/apply_modules"


@total_ordering
class VibExport(object):
   """
   Base class for exports.

   Attributes:
      * path - The path under the vibExports directory in the metadata archive.
               The first component of the path is the ID of the VIB.
   """
   def __init__(self, path):
      self.path = path
      self._splitPath = self.path.split('/')
      if len(self._splitPath) < 2:
         raise ValueError("path must have a VIB ID prefix")

   def __eq__(self, other):
      return self.path == other.path

   def __ne__(self, other):
      return not self.__eq__(other)

   def __lt__(self, other):
      return self.path < other.path

   @property
   def schemaId(self):
      """
      Returns: The config schema ID of the VIB this export is from.  The
               format of the ID is the same as used with config schemas: <vib-name>-<vib-version-with-underscores>.
      """
      return self._splitPath[0]


class _OpaqueExport(VibExport):
   """
   Base class for exports that are an opaque bag of bytes.
   """
   def __init__(self, path, content):
       super().__init__(path)
       self.content = content

   @classmethod
   def FromFile(cls, path, content):
      """
      Creates an export from the given path and file content.

      Parameters:
         * path    - The path to the export under the vibExports metadata directory
         * content - The file content
      """
      return cls(path, content)

   def WriteFile(self, root):
      """
      Write this export to a file at the given location.

      Parameters:
         * root - The path to the vibExports metadata directory where the file
                  should be written to.
      """
      filePath = os.path.join(root, self.path)
      os.makedirs(os.path.dirname(filePath), exist_ok=True)
      with open(filePath, 'wb') as fp:
         fp.write(self.content)


class VmsgExport(_OpaqueExport):
   """
   Represents a vmsg VIB export.
   """
   pass


def getLoaderForVibPath(schemaId, path):
   """
   If the given path from a VIB is an export, return a loader function.

   Parameters:
      * schemaId - The schema ID of the VIB this path is from.
      * path  - The path to the file in the VIB.
   Returns: If the path is an export, a loader function is returned that
            accepts the file content and returns a VibExport object.  If
            the path is not an export, None is returned.
   """
   if path.startswith(CONFIG_LOCALE_DIR) and path.endswith('.vmsg'):
      def vmsgLoader(content):
         langPath = os.path.relpath(path, CONFIG_LOCALE_DIR)
         metaPath = os.path.join(schemaId, 'config', 'locale', langPath)
         return VmsgExport.FromFile(metaPath, content)
      return vmsgLoader
   if path.startswith(MODULES_DIR) and path.endswith('.vmsg'):
      def pluginVmsgLoader(content):
         langPath = os.path.relpath(path, MODULES_DIR)
         metaPath = os.path.join(schemaId, 'config', langPath)
         return VmsgExport.FromFile(metaPath, content)
      return pluginVmsgLoader
   return None


def getLoaderForMetadataPath(path):
   """
   If the given path from the metadata is an export, return a loader function.

   Parameters:
      * path  - The path to the file in metadata archive.
   Returns: If the path is an export, a loader function is returned that
            accepts the file content and returns a VibExport object.  If
            the path is not an export, None is returned.
   """
   splitPath = path.split('/')
   if splitPath[0] == "vibExports" and splitPath[2] == 'config' and path.endswith('.vmsg'):
      def vmsgLoader(content):
         return VmsgExport.FromFile('/'.join(splitPath[1:]), content)
      return vmsgLoader
   return None


class VibExportCollection(dict):
   """
   A dictionary of VIB IDs to a list of exports from the corresponding VIB.
   """
   def __add__(self, other):
      """Merge two objects and return a new one.
      """
      new = self.__class__()
      for export in itertools.chain.from_iterable(self.values()):
         new.AddVibExport(export)
      new += other
      return new

   def __iadd__(self, other):
      for export in itertools.chain.from_iterable(other.values()):
         self.AddVibExport(export)
      return self

   def AddVibExport(self, export):
      """
      Adds a VibExport object to this collection.  The export is added to the
      list of exports for the VIB the export is from.

      Parameters:
         * export - The export to add to this collection.
      """
      bisect.insort(self.setdefault(export.schemaId, list()), export)

   def FromDirectory(self, path):
      """
      Populate this collection from the given metadata directory.

      Parameters:
         * path - The path to the extracted vibExports metadata directory.
      """
      self.clear()
      for root, _, files in os.walk(path, topdown=True):
         for name in files:
            filePath = os.path.join(root, name)
            exportPath = os.path.relpath(filePath, path)
            with open(filePath, 'rb') as fp:
               if filePath.endswith('.vmsg'):
                  export = VmsgExport.FromFile(exportPath, fp.read())
               else:
                  raise NotImplementedError("Unhandled vib export: %s" % exportPath)
            self.AddVibExport(export)

   def ToDirectory(self, path):
      """
      Write this collection of exports to disk.

      Parameters:
         * path - The path to the vibExports metadata directory.
      """
      if os.path.isdir(path):
         shutil.rmtree(path)
      os.makedirs(path)
      for export in itertools.chain.from_iterable(self.values()):
         export.WriteFile(path)
