"""
 Copyright 2021-2022 VMware, Inc.  All rights reserved.
 -- VMware Confidential

 """
__author__ = "VMware, Inc"

"""
builder.py

This script is used to generate the vSAN Java SDK and sample codes. Moreoever,
this script contains the following functionalities:

   1. Copies and prepares the environment for building vSAN Java SDK;
   2. Compiles the vSAN Java SDK from the WSDL file;
   3. Compiles the vSAN Java sample code;

The vSAN SDK is generated into the <current_path>/lib/vsanmgmt-sdk.jar file
and the sample code will be generated into vsan-samples.jar.

usage: Build utility for generating the vSAN Management SDK Stubs and compiling the samples
       [-h] [-c | -s | -w] [-j JAVA_HOME] [-v VSPHERE_SDK_FOLDER]

"""

import re
import os
import shutil
import argparse
import sys
from os import listdir
from glob import glob
from os.path import isfile, join
import subprocess
from subprocess import check_output

VSAN_BUILD_DIR = os.path.dirname(os.path.realpath(__file__))
VSAN_SDK_JAVA_PACKAGE = 'com.vmware.vsan.sdk'
VSAN_SAMPLE_DIR = join(VSAN_BUILD_DIR, 'samples')
VSAN_STUB_DIR = join(VSAN_SAMPLE_DIR, *(VSAN_SDK_JAVA_PACKAGE.split('.')))
VSAN_STUB_PATCH_ROOT = join(VSAN_BUILD_DIR, 'vsan_sdk')
LIB_DIR = join(VSAN_BUILD_DIR, 'lib')
WSDLDIR = join(VSAN_BUILD_DIR, os.pardir, 'bindings', 'wsdl')
WSDLNAME = 'vsanService.wsdl'
VSANMGMT_SDK_JAR = 'vsanmgmt-sdk.jar'
VSANMGMT_SAMPLES_JAR = 'vsan-samples.jar'
VSAN_VMODL_DEV = os.environ.get('VSAN_VMODL_DEV', False)

PRODUCT_NAME = 'vSAN Management SDK'
VSANHEALTH_VIM_PORT_TYPE = "VsanhealthVimPortType"

class VsanSdkBuilder:
   def __init__(self):
      self.java_home = None
      self.java_bin_path = None
      self.tool_jar_path = None
      self.vsphere_sdk_path = None
      self.result = None

   def run(self):
      args = self.parse_args()
      self.set_build_env(args)

      do_all = not args.copy_only and \
               not args.build_stub_only and \
               not args.build_sample_only

      copy = args.copy_only or args.build_stub_only or args.build_sample_only or do_all
      build_stub = args.build_stub_only or do_all
      build_sample = args.build_sample_only or do_all

      if copy:
         self.copy_vsphere_java_sdk()
         self.copy_JAXWS_lib()
      if build_stub:
         self.generate_vsan_sdk_stub()
      if build_sample:
         self.compile_vsan_sdk_sample()

      self.print_result(build_stub, build_sample)

   def parse_args(self):
      title = 'Build utility for generating the {0} Stubs and compiling the samples'
      parser = argparse.ArgumentParser(title.format(PRODUCT_NAME))
      group = parser.add_mutually_exclusive_group()
      group.add_argument("-c", "--copy-only",
                         help="copy vSphere SDK jar files only.",
                         action="store_true")
      group.add_argument("-s", "--build-stub-only",
                         help="build vSAN SDK Stubs only.",
                         action="store_true")
      group.add_argument("-w", "--build-sample-only",
                         help="build vSAN Samples code only.",
                         action="store_true")
      parser.add_argument("-j", "--java-home", type=str,
                          default=os.environ['JAVA_HOME'],
                          help="Java home path (optional).")
      parser.add_argument("-v", "--vsphere-sdk-folder", type=str,
                          default=join(VSAN_BUILD_DIR, os.pardir, os.pardir),
                          help="vSphere SDK folder (optional).")
      return parser.parse_args()

   def set_build_env(self, args):
      if not args.java_home:
         print("Please set the JAVA_HOME through environment variable or command line option")
         sys.exit(1)

      if not os.path.isdir(LIB_DIR):
         os.mkdir(LIB_DIR)

      self.vsphere_sdk_path = args.vsphere_sdk_folder
      self.java_home = args.java_home
      version = subprocess.check_output(\
            [join(self.java_home , 'bin', 'java'), '-version'], \
            stderr=subprocess.STDOUT)
      pattern = '\"(\d+\.\d+).*\"'.encode('utf-8')
      results = re.search(pattern, version).groups()
      if len(results) == 0:
         print("Parse Java Version Error. Is JAVA_HOME valid?")
         sys.exit(1)
      self.java_version = results[0].decode()
      self.jaxws_home = join(self.vsphere_sdk_path, 'libs', 'jaxws-ri')
      self.java_bin_path = join(self.java_home, 'bin', 'java')
      self.tool_jar_path = join(self.java_home, 'lib', 'tools.jar')

      print("JAVA_HOME        : %s" % self.java_home)
      print("JDK Version      : %s" % str(self.getJDKVersion()))
      print("vSphere SDK path : %s" % self.vsphere_sdk_path)
      print("")

   def print_result(self, build_stub, build_sample):
      print("")
      if not build_stub and not build_sample:
         return

      print("Build result:")
      if build_stub:
         print("vSAN Java SDK    : %s " % join(LIB_DIR, VSANMGMT_SDK_JAR))
      if build_sample and not VSAN_VMODL_DEV:
         print("vSAN Java sample : %s " % join(LIB_DIR, VSANMGMT_SAMPLES_JAR))
      return

   def add_prerequisite(self, source, target):
      filename = os.path.basename(source)

      if isfile(target):
         print("File %s already exists in the install directory, skip."
               % filename)
      elif not isfile(source):
         print("File %s not found. Is %s a valid vSphere SDK path?"
               % (filename, self.vsphere_sdk_path))
      else:
         shutil.copyfile(source, target)

   def copy_vsphere_java_sdk(self):
      '''
      Copy vSphere Java SDK libraries to the sample code base.
      '''
      print("Preparing vSphere SDK sample library...")
      prerequisite_list = [
            'ssoclient/java/JAXWS/lib/ssoclient.jar',
            'ssoclient/java/JAXWS/lib/ssosamples.jar',
            'vsphere-ws/java/JAXWS/lib/vim25.jar',
            'vsphere-ws/java/JAXWS/lib/samples-annotations-1.0.0.jar',
            'vsphere-ws/java/JAXWS/lib/samples-core-1.0.0.jar',
            'vsphere-ws/java/JAXWS/lib/wssamples.jar',
            'spbm/java/JAXWS/lib/pbm.jar',
            'vsphere-ws/java/JAXWS/lib/sblim-cim-client2-2.1.1.jar',
            'vsphere-ws/java/JAXWS/lib/sblim-cim-client2.properties',
      ]

      for prerequ in prerequisite_list:
         source = join(self.vsphere_sdk_path,
                       *(prerequ.split('/')))
         target_filename = os.path.basename(source)
         target = join(LIB_DIR, target_filename)
         self.add_prerequisite(source, target)


      if os.name == 'nt':
         shutil.copyfile(join(self.vsphere_sdk_path,
                           *('vsphere-ws/java/JAXWS/lcp.bat'.split('/'))),
                         join(VSAN_BUILD_DIR, 'lcp.bat'))

   def copy_JAXWS_lib(self):
      if os.path.exists(join(self.jaxws_home, 'lib')):
         for file in os.listdir(join(self.jaxws_home, 'lib')):
            if file[-4:].lower() == '.jar':
               self.add_prerequisite(join(self.jaxws_home, 'lib', file), join(LIB_DIR, file))


   def getJDKVersion(self):
      version = self.java_version.split('.')[0]
      return int(version)

   def generate_vsan_sdk_stub(self):
      '''
      Generate and compiles vSAN SDK stub files by parsing WSDL files.
      '''
      print("Generating vSAN stubs from WSDL...")

      stub_src_root = join(VSAN_STUB_PATCH_ROOT, 'source')
      stub_class_root = join(VSAN_STUB_PATCH_ROOT, 'target')
      stub_src_dir = join(stub_src_root, *(VSAN_SDK_JAVA_PACKAGE.split('.')))
      stub_class_dir = join(stub_class_root, *(VSAN_SDK_JAVA_PACKAGE.split('.')))

      if not os.path.isdir(stub_src_dir):
         os.makedirs(stub_src_dir)
      if not os.path.isdir(stub_class_dir):
         os.makedirs(stub_class_dir)

      wsimport_cmd_list = []

      if os.name == 'nt':
         '''
         wsimport on Windows:
         ${SDK_PATH}/libs/jaxws-ri/bin/wsimport.bat -wsdllocation $WSDLNAME \
                               -p com.vmware.vsan.sdk -s ${VSAN_SAMPLE_DIR} \
                               ${WSDLDIR}/${WSDLNAME}
         '''
         wsimport_cmd_list = [join(self.jaxws_home, 'bin', 'wsimport.bat'),
                              '-wsdllocation', WSDLNAME, '-p', VSAN_SDK_JAVA_PACKAGE, \
                              '-s', VSAN_SAMPLE_DIR, join(WSDLDIR, WSDLNAME), \
                              '-Xnocompile', '-B-npa']
      else:
         if self.getJDKVersion() >= 11:
            '''
            wsimport on Linux:
            sh ${SDK_PATH}/libs/jaxws-ri/bin/wsimport.sh -wsdllocation $WSDLNAME \
                               -p com.vmware.vsan.sdk -s ${VSAN_SAMPLE_DIR} \
                               ${WSDLDIR}/${WSDLNAME}
            '''
            wsimport_cmd_list = ['sh', join(self.jaxws_home, 'bin', 'wsimport.sh'),
                              '-wsdllocation', WSDLNAME, '-p', VSAN_SDK_JAVA_PACKAGE, \
                              '-s', VSAN_SAMPLE_DIR, join(WSDLDIR, WSDLNAME), \
                              '-Xnocompile', '-B-npa']
         else:
            '''
            wsimport on Linux:
            ${JAVA_HOME}/bin/java -Xmx512m -classpath ${JAVA_HOME}/lib/tools.jar \
                               com.sun.tools.internal.ws.WsImport \
                               -wsdllocation $WSDLNAME -p com.vmware.vsan.sdk \
                               -s ${VSAN_SAMPLE_DIR} ${WSDLDIR}/${WSDLNAME}
                               -Xnocompile -B-npa
            '''
            wsimport_cmd_list = [self.java_bin_path, '-Xmx512m', '-classpath',
                              self.tool_jar_path, 'com.sun.tools.internal.ws.WsImport',
                              '-wsdllocation', WSDLNAME, '-p',
                              VSAN_SDK_JAVA_PACKAGE, '-s', VSAN_SAMPLE_DIR,
                              join(WSDLDIR, WSDLNAME), '-Xnocompile', '-B-npa']

      check_output(wsimport_cmd_list)

      '''
      ${JAVA_HOME}/bin/java -classpath ${LIB_DIR} FixJaxWsWsdlResource \
                  "${VSAN_STUB_DIR}/VsanhealthService.java" VsanhealthService
      '''
      fix_wsdl_resource_cmd = [self.java_bin_path, '-classpath', LIB_DIR,
                              'FixJaxWsWsdlResource',
                              join(VSAN_STUB_DIR, 'VsanhealthService.java'),
                              'VsanhealthService']
      check_output(fix_wsdl_resource_cmd)

      # Patch the generated vSAN stubs.
      self.process_generated_vsan_stub_src(VSAN_STUB_DIR, stub_src_dir)

      print("Compiling vSAN stubs...")
      '''
      ${JAVA_HOME}/bin/javac -J-Xms1024M -J-Xmx1024M -classpath ${LIB_DIR}/vim25.jar @sdk_sources.txt
      '''
      sdk_sources_txt = join(stub_src_dir, 'sdk_sources.txt')
      self.write_files_with_ext_to_file(stub_src_root, '*.java', sdk_sources_txt)

      compile_src_cmd = [join(self.java_home, 'bin', 'javac'), '-J-Xms1024M',
                         '-J-Xmx1024M', '-classpath',
                         join(LIB_DIR, '*'),
                         '-d', stub_class_root, '@' + sdk_sources_txt]
      check_output(compile_src_cmd, cwd=stub_src_root)

      print("Generating vSAN stub Jar...")
      shutil.copy(join(WSDLDIR, WSDLNAME), join(stub_class_dir, WSDLNAME))
      shutil.copy(join(WSDLDIR, 'vsan.wsdl'),
                  join(stub_class_dir, 'vsan.wsdl'))

      '''
      ${JAVA_HOME}/bin/jar cf ${LIB_DIR}/vsanmgmt-sdk.jar -C ${stub_class_root} .
      '''
      jar_cmd = [join(self.java_home, 'bin', 'jar'), 'cf',
                 join(LIB_DIR, VSANMGMT_SDK_JAR), '-C', stub_class_root, '.']
      check_output(jar_cmd)

      shutil.rmtree(VSAN_STUB_DIR)
      shutil.rmtree(VSAN_STUB_PATCH_ROOT)

   def write_files_with_ext_to_file(self, root_dir, extension, output_file):
      '''
      Find all files in root_dir recursively with extension(*.java,
      *.class etc.) to , and write them output_file with each file path
      relative to root_dir.
      '''
      import fnmatch
      match_files = []
      for root, dirnames, filenames in os.walk(root_dir):
         for filename in fnmatch.filter(filenames, extension):
            path = join(root, filename)
            length = len(root_dir)
            if not root_dir.endswith(os.sep):
               length = length + 1
            match_files.append(path[length:] + os.linesep)

      with open(join(output_file), 'w+') as ofile:
         ofile.writelines(match_files)


   def compile_vsan_sdk_sample(self):
      '''
      Compile vSAN SDK sample code.
      '''
      if VSAN_VMODL_DEV:
         print("\nYou are compiling with dev version and " + \
               "so far we don't support sample code compilation")
         return

      print("Compiling vSAN SDK samples...")

      package_path = VSAN_SDK_JAVA_PACKAGE.replace('.', os.sep)
      source_txt = join(VSAN_SAMPLE_DIR, 'sources.txt')
      class_txt = join(VSAN_SAMPLE_DIR, 'classes.txt')

      self.write_files_with_ext_to_file(VSAN_SAMPLE_DIR, '*.java', source_txt)
      '''
      ${JAVA_HOME}/bin/javac -classpath "../lib/*" @source.txt
      '''
      class_path_sep = ';' if os.name == 'nt' else ':'
      jar_deps = class_path_sep.join(glob(join(LIB_DIR, '*')))
      compile_src_cmd = [join(self.java_home, 'bin', 'javac'), '-J-Xms1024M',
                        '-J-Xmx1024M', '-classpath', jar_deps, '@' + source_txt]
      check_output(compile_src_cmd, cwd=VSAN_SAMPLE_DIR)

      print("Generating vSAN SDK samples Jar...")
      self.write_files_with_ext_to_file(VSAN_SAMPLE_DIR, '*.class', class_txt)
      '''
      ${JAVA_HOME}/bin/jar cf ../lib/vsan-samples.jar ${sample_class_files}
      '''
      jar_cmd = [join(self.java_home, 'bin', 'jar'), 'cf',
                 join(LIB_DIR, VSANMGMT_SAMPLES_JAR),
                 '@' + class_txt]

      check_output(jar_cmd, cwd=VSAN_SAMPLE_DIR)

      os.remove(source_txt)
      os.remove(class_txt)


   def process_port_type_file(self, sdk_src_file, output_dir):
      '''
      Process the VsanhealthPortType.java to generate the port type for the
      namespace of vsan and vim25 respectively.

      After processing:

       * VsanhealthPortType.java having all methods with ns urn:vsan.
       * VsanhealthVimPortType.java having all methods with ns urn:vim25.
      '''
      filename = os.path.basename(sdk_src_file)
      with open(sdk_src_file, 'r') as ifile, \
            open(join(output_dir, VSANHEALTH_VIM_PORT_TYPE + ".java"), 'w+') as vim_port_file, \
            open(join(output_dir, filename), 'w+') as vsan_port_file:
         while True:
            line = ifile.readline()
            if not line:
               break
            if line.startswith('public interface VsanhealthPortType'):
               vim_port_file.write(re.sub(r'VsanhealthPortType',
                                          VSANHEALTH_VIM_PORT_TYPE+" extends VsanhealthPortType",
                                          line))
               vsan_port_file.write(line)
               continue

            if line.startswith("    @RequestWrapper") or \
               line.startswith("    @ResponseWrapper"):
               vim_line = re.sub(r'com.vmware.vsan.sdk.(\w+)(RequestType|Response)',
                                 r'com.vmware.vsan.sdk.\g<1>Vim\g<2>',
                                 line.replace('"urn:vsan"', '"urn:vim25"'))
               vsan_line = re.sub(r'com.vmware.vsan.sdk.(\w+)(RequestType|Response)',
                                  r'com.vmware.vsan.sdk.\g<1>Vsan\g<2>',
                                  line)

               vim_port_file.write(vim_line)
               vsan_port_file.write(vsan_line)
               continue

            vim_line = re.sub(r'"urn:vsan"', '"urn:vim25"', line)
            vim_line = re.sub(r'"urn:vsan/.*"', '"urn:vim25/6.5"', vim_line)
            vim_port_file.write(vim_line)
            vsan_port_file.write(line)

   def process_api_types(self, wsdlname, sdk_src_file, type_suffix, output_dir):
      '''
      Process <wsdlName>RequestType or <wsdlName>Response file, generating
      2 files for each of them, for request and response with namespace vim25
      and vsan respectively.
      '''
      assert sdk_src_file.endswith(('RequestType.java', 'Response.java'))
      class_vsan = "{0}{1}{2}".format(wsdlname, 'Vsan', type_suffix)
      class_vim = "{0}{1}{2}".format(wsdlname, 'Vim', type_suffix)
      # This is an old(urn:vim25/6.5) vSAN method, so we need generate 2
      # RequestType(or Response), one for vSAN, one for vim25
      with open(sdk_src_file, 'r') as ifile, \
            open(join(output_dir, class_vim + '.java'), 'w+') as vimfile, \
            open(join(output_dir, class_vsan + '.java'), 'w+') as vsanfile:
         def write_line(ofile, line, typ, class_name):
            if typ == 'vim':
               if 'namespace = "urn:vsan"' in line:
                  line = re.sub(r'namespace = "urn:vsan"',
                                r'namespace = "urn:vim25"', line)
            line = re.sub(r'public (?P<type>class|enum) (?P<class_name>\w+) {',
                          r'public \g<1> {0} {{'.format(class_name), line)
            line = line.replace(wsdlname + type_suffix, class_name)
            line = line.replace(wsdlname[0].lower() + wsdlname[1:] + type_suffix, \
                  class_name)
            ofile.write(line)
         for line in ifile:
            write_line(vimfile, line, 'vim', class_vim)
            write_line(vsanfile, line, 'vsan', class_vsan)

   def process_service_file(self, sdk_src_file, output_dir):
      '''
      Process VsanhealthService.java file and reflects the method into
      getVsanhealthVimPort() and getVsanhealthVimPort() for the vim25 and
      vsan namespace.
      '''
      func_def_vim_port = '''
   /**
     *
     * @return
     *     returns VsanhealthVimPortType
     */
    @WebEndpoint(name = "VsanhealthPort")
    public VsanhealthPortType getVsanhealthVimPort() {
        return super.getPort(new QName("urn:vsanService", "VsanhealthPort"), VsanhealthVimPortType.class);
    }

    /**
     *
     * @param features
     *     A list of {@link javax.xml.ws.WebServiceFeature} to configure on the proxy.  Supported features not in the <code>features</code> parameter will have their default values.
     * @return
     *     returns VsanhealthVimPortType
     */
    @WebEndpoint(name = "VsanhealthPort")
    public VsanhealthPortType getVsanhealthVimPort(WebServiceFeature... features) {
        return super.getPort(new QName("urn:vsanService", "VsanhealthPort"), VsanhealthVimPortType.class, features);
    }
   '''
      filename = os.path.basename(sdk_src_file)
      with open(sdk_src_file, 'r') as ifile, \
           open(join(output_dir, filename), 'w+') as vimfile:
         for line in ifile:
            if line.startswith('    private static URL __getWsdlLocation() {'):
               vimfile.write(func_def_vim_port)
            vimfile.write(line)


   def process_factory_file(self, sdk_src_file, output_dir):
      '''
      Process ObjectFactory.java file. As we just add 2 types for both
      <MethodName>RequestType and <MethodName>Response Java files,
      we need to update this file by which a list of factory methods are impacted.
      '''

      filename = os.path.basename(sdk_src_file)
      func_pattern1 = re.compile(r'     \* Create an instance of {@link (?P<wsdlName>\w+)(?P<rType>Response|RequestType) }')
      func_pattern2 = re.compile(r'     \* Create an instance of {@link JAXBElement }{@code <}{@link (?P<wsdlName>\w+)(?P<rType>RequestType) }{@code >}')
      FUNC_START = '    /**'
      func_format1  = '''     * Create an instance of {{@link {0}{1}{2} }}
        *
        */
       public {0}{1}{2} create{0}{1}{2}() {{
           return new {0}{1}{2}();
       }}
   '''
      FUNC_END = '    }'

      func_format2 = '''     * Create an instance of {{@link JAXBElement }}{{@code <}}{{@link {0}{1}{2} }}{{@code >}}
        *
        */
       @XmlElementDecl(namespace = "urn:vsan", name = "{0}")
       public JAXBElement<{0}{1}RequestType> create{0}({0}{1}RequestType value) {{
           return new JAXBElement<{0}{1}RequestType>(new QName("urn:vsan", "{0}{1}"), {0}{1}RequestType.class, null, value);
       }}
   '''

      recompile = True
      def process_method(match, func_format, ofile):
         wsdl_name = match.group('wsdlName')
         rtype = match.group('rType')
         ofile.write(func_format.format(wsdl_name, 'Vsan', rtype))
         ofile.write(os.linesep)
         # Remove the vim25 vmodl API due to large ObjectFactory.java file
         # It might impact the compatibility with vSAN SDK prior to 6.5
         # ofile.write(FUNC_START + os.linesep)
         # vim_func_main = func_format.format(wsdl_name, 'Vim', rtype)
         # ofile.write(vim_func_main.replace('"urn:vsan"', '"urn:vim25"'))

      with open(sdk_src_file, 'r') as ifile, \
           open(join(output_dir, filename), 'w+') as dest_file:
         while True:
            line = ifile.readline()
            if not line:
               break
            if (line.startswith("}")):
               dest_file.write("}")
               break
            if line.startswith(FUNC_START):
               recompile = True
            match1 = func_pattern1.search(line)
            if match1:
               process_method(match1, func_format1, dest_file)
               recompile = False
               continue
            match2 = func_pattern2.search(line)
            if match2:
               process_method(match2, func_format2, dest_file)
               recompile = False
               continue
            if recompile:
               dest_file.write(line.replace('namespace = "urn:vsan", ', ''))

   def process_src(self, sdk_src_file, output_dir):
      '''
      Process a single Java source file to generate vSAN SDK stub.
      '''
      filename = os.path.basename(sdk_src_file)
      if filename == 'VsanhealthService.java':
         self.process_service_file(sdk_src_file, output_dir)
         return
      elif filename == 'VsanhealthPortType.java':
         self.process_port_type_file(sdk_src_file, output_dir)
         return
      elif filename == 'ObjectFactory.java':
         self.process_factory_file(sdk_src_file, output_dir)
         return

      # If this is an API RequestType or Response class, we need to fork this
      # class into two to fit in vim25 and vsan namespace.
      class_name, _ = os.path.splitext(filename)
      method_names = class_name.split('RequestType')[0].split('Response')
      if (method_names[0] != class_name):
         method_name = method_names[0]
         type_suffix = class_name.replace(method_name, "")
         self.process_api_types(method_name, sdk_src_file, type_suffix,
                           output_dir)
         return

      # For plain Data Object, keep as it is.
      shutil.copyfile(sdk_src_file, join(output_dir, filename))

   def process_generated_vsan_stub_src(self, input_stub_dir, output_stub_dir):
      '''
      Patch generated original Java source files that are compatible for
      both vCenter 6.5 and vSAN 6.6 environment.
      '''
      for source_file in listdir(input_stub_dir):
         if isfile(join(input_stub_dir, source_file)) and \
               source_file.endswith('.java'):
            self.process_src(join(input_stub_dir, source_file), output_stub_dir)

if __name__ == '__main__':
   VsanSdkBuilder().run()
