#!/usr/bin/env python

"""Ninja toolchain abstraction for Microsoft compiler suite"""

import os
import subprocess

import toolchain
import vslocate

class MSVCToolchain(toolchain.Toolchain):

  def initialize(self, project, archs, configs, includepaths, dependlibs, libpaths, variables, subninja):
    #Local variable defaults
    self.sdkpath = ''
    self.toolchain = ''
    self.includepaths = []
    self.libpaths = libpaths
    self.ccompiler = 'cl'
    self.cxxcompiler = 'cl'
    self.archiver = 'lib'
    self.linker = 'link'
    self.dller = 'dll'

    #Command definitions (to generate assembly, add "/FAs /Fa$out.asm")
    self.cccmd = '$toolchain$cc /showIncludes /I. $includepaths $moreincludepaths $cflags $carchflags $cconfigflags $cmoreflags /c $in /Fo$out /Fd$pdbpath /FS /nologo'
    self.cxxcmd = '$toolchain$cxx /showIncludes /I. $includepaths $moreincludepaths $cxxflags $carchflags $cconfigflags $cmoreflags /c $in /Fo$out /Fd$pdbpath /FS /nologo'
    self.ccdepfile = None
    self.ccdeps = 'msvc'
    self.arcmd = '$toolchain$ar $arflags $ararchflags $arconfigflags /NOLOGO /OUT:$out $in'
    self.linkcmd = '$toolchain$link $libpaths $configlibpaths $linkflags $linkarchflags $linkconfigflags /DEBUG /NOLOGO /SUBSYSTEM:CONSOLE /DYNAMICBASE /NXCOMPAT /MANIFEST /MANIFESTUAC:\"level=\'asInvoker\' uiAccess=\'false\'\" /TLBID:1 /PDB:$pdbpath /OUT:$out $in $libs $archlibs $oslibs'
    self.dllcmd = self.linkcmd + ' /DLL'

    self.cflags = ['/D', '"' + project.upper() + '_COMPILE=1"', '/D', '"_UNICODE"',  '/D', '"UNICODE"', '/std:c17', '/Zi', '/Oi', '/Oy-', '/GS-', '/Gy-', '/Qpar-', '/fp:fast', '/fp:except-', '/Zc:forScope', '/Zc:wchar_t', '/GR-', '/openmp-']
    self.cwarnflags = ['/W4', '/WX', '/wd4201'] #Ignore nameless union/struct which is allowed in C11
    self.cmoreflags = []
    self.arflags = ['/ignore:4221'] #Ignore empty object file warning]
    self.linkflags = ['/DEBUG']
    self.oslibs = ['kernel32', 'user32', 'shell32', 'advapi32']

    self.initialize_subninja(subninja)
    self.initialize_archs(archs)
    self.initialize_configs(configs)
    self.initialize_project(project)
    self.initialize_toolchain()
    self.initialize_depends(dependlibs)

    self.parse_default_variables(variables)
    self.read_build_prefs()

    self.includepaths = self.prefix_includepaths((includepaths or []) + ['.'])

    if self.is_monolithic():
      self.cflags += ['/D', '"BUILD_MONOLITHIC=1"']

    if not 'nowarning' in variables or not variables['nowarning']:
      self.cflags += self.cwarnflags
    self.cxxflags = list(self.cflags)

    #Overrides
    self.objext = '.obj'

    #Builders
    self.builders['c'] = self.builder_cc
    self.builders['cc'] = self.builder_cxx
    self.builders['cpp'] = self.builder_cxx
    self.builders['lib'] = self.builder_lib
    self.builders['multilib'] = self.builder_multicopy
    self.builders['sharedlib'] = self.builder_sharedlib
    self.builders['multisharedlib'] = self.builder_multicopy
    self.builders['bin'] = self.builder_bin
    self.builders['multibin'] = self.builder_multicopy

    #Setup toolchain
    self.build_toolchain()

  def name(self):
    return 'msvc'

  def parse_prefs(self, prefs):
    super(MSVCToolchain, self).parse_prefs(prefs)
    if 'msvc' in prefs:
      msvcprefs = prefs['msvc']
      if 'sdkpath' in msvcprefs:
        self.sdkpath = msvcprefs['sdkpath']
      if 'toolchain' in msvcprefs:
        self.toolchain = msvcprefs['toolchain']
      if 'toolchain_version' in msvcprefs:
        self.toolchain_version = msvcprefs['toolchain_version']

  def write_variables(self, writer):
    super(MSVCToolchain, self).write_variables(writer)
    writer.variable('cc', self.ccompiler)
    writer.variable('cxx', self.cxxcompiler)
    writer.variable('ar', self.archiver)
    writer.variable('link', self.linker)
    writer.variable('dll', self.dller)
    writer.variable('toolchain', self.toolchain)
    writer.variable('includepaths', self.make_includepaths(self.includepaths))
    writer.variable('moreincludepaths', '')
    writer.variable('pdbpath', 'ninja.pdb')
    writer.variable('cflags', self.cflags)
    writer.variable('carchflags', '')
    writer.variable('cconfigflags', '')
    writer.variable('cxxflags', self.cxxflags)
    writer.variable('cmoreflags', self.cmoreflags)
    writer.variable('arflags', self.arflags)
    writer.variable('ararchflags', '')
    writer.variable('arconfigflags', '')
    writer.variable('linkflags', self.linkflags)
    writer.variable('linkarchflags', '')
    writer.variable('linkconfigflags', '')
    writer.variable('libs', '')
    writer.variable('libpaths', self.make_libpaths(self.libpaths))
    writer.variable('configlibpaths', '')
    writer.variable('archlibs', '')
    writer.variable('oslibs', self.make_libs(self.oslibs))
    writer.newline()

  def write_rules(self, writer):
    super(MSVCToolchain, self).write_rules(writer)
    writer.rule('cc', command = self.cccmd, depfile = self.ccdepfile, deps = self.ccdeps, description = 'CC $in')
    writer.rule('cxx', command = self.cxxcmd, depfile = self.ccdepfile, deps = self.ccdeps, description = 'CXX $in')
    writer.rule('ar', command = self.arcmd, description = 'LIB $out')
    writer.rule('link', command = self.linkcmd, description = 'LINK $out')
    writer.rule('dll', command = self.dllcmd, description = 'DLL $out')
    writer.newline()

  def build_toolchain(self):
    if self.toolchain == '':
      installed_versions = vslocate.get_vs_installations()
      for versionstr, installpath in installed_versions:
        major_version = versionstr.split('.')[0]
        if int(major_version) >= 15:
          tools_basepath = os.path.join(installpath, 'VC', 'Tools', 'MSVC')
          tools_list = [item for item in os.listdir(tools_basepath) if os.path.isdir(os.path.join(tools_basepath, item))]
          from distutils.version import StrictVersion
          tools_list.sort(key=StrictVersion)
          self.toolchain = os.path.join(tools_basepath, tools_list[-1])
          self.toolchain_version = major_version + ".0"
          break

      if self.toolchain == '':
        toolchain = ''
        versions = ['17.0', '16.0', '15.0']
        keys = [
          'HKLM\\SOFTWARE\\Microsoft\\VisualStudio\\SxS\\VC7',
          'HKCU\\SOFTWARE\\Microsoft\\VisualStudio\\SxS\\VC7',
          'HKLM\\SOFTWARE\\Microsoft\\VisualStudio\\SxS\\VS7',
          'HKCU\\SOFTWARE\\Microsoft\\VisualStudio\\SxS\\VS7',
          'HKLM\\SOFTWARE\\Wow6432Node\\Microsoft\\VisualStudio\\SxS\\VC7',
          'HKCU\\SOFTWARE\\Wow6432Node\\Microsoft\\VisualStudio\\SxS\\VC7',
          'HKLM\\SOFTWARE\\Wow6432Node\\Microsoft\\VisualStudio\\SxS\\VS7',
          'HKCU\\SOFTWARE\\Wow6432Node\\Microsoft\\VisualStudio\\SxS\\VS7'
        ]
        for version in versions:
          for key in keys:
            try:
              query = subprocess.check_output(['reg', 'query', key, '/v', version ], stderr = subprocess.STDOUT).strip().splitlines()
              if len(query) == 2:
                toolchain = str(query[1]).split('REG_SZ')[-1].strip(" '\"\n\r\t")
            except:
              continue
            if not toolchain == '':
              tools_basepath = os.path.join(toolchain, 'VC', 'Tools', 'MSVC')
              tools_list = [item for item in os.listdir(tools_basepath) if os.path.isdir(os.path.join(tools_basepath, item))]
              from distutils.version import StrictVersion
              tools_list.sort(key=StrictVersion)
              toolchain = os.path.join(tools_basepath, tools_list[-1])
              self.toolchain = toolchain
              self.toolchain_version = version
              break
          if not self.toolchain == '':
            break
    if self.toolchain == '':
      raise Exception("Unable to locate any installed Visual Studio toolchain")
    self.includepaths += [os.path.join(self.toolchain, 'include')]
    if self.sdkpath == '':
      versions = ['v10.0', 'v8.1']
      keys = [
        'HKLM\\SOFTWARE\\Microsoft\\Microsoft SDKs\\Windows',
        'HKCU\\SOFTWARE\\Microsoft\\Microsoft SDKs\\Windows',
        'HKLM\\SOFTWARE\\Wow6432Node\\Microsoft\\Microsoft SDKs\\Windows',
        'HKCU\\SOFTWARE\\Wow6432Node\\Microsoft\\Microsoft SDKs\\Windows'
      ]
      include_path = 'include'
      for version in versions:
        for key in keys:
          sdkpath = ''
          try:
            query = subprocess.check_output(['reg', 'query', key + '\\' + version, '/v', 'InstallationFolder'], stderr = subprocess.STDOUT).strip().splitlines()
            if len(query) == 2:
              sdkpath = str(query[1]).split('REG_SZ')[-1].strip(" '\"\n\r\t")
              if not sdkpath == '' and version == 'v10.0':
                base_path = sdkpath
                sdkpath = ''
                query = subprocess.check_output(['reg', 'query', key + '\\' + version, '/v', 'ProductVersion'], stderr = subprocess.STDOUT).strip().splitlines()
                if len(query) == 2:
                  version_path = str(query[1]).split('REG_SZ')[-1].strip(" '\"\n\r\t")
                  if not version_path == '':
                    sdkpath = base_path
                    self.sdkversionpath = version_path
                    versioned_include_path = os.path.join(include_path, self.sdkversionpath)
                    if not os.path.exists(os.path.join(sdkpath, versioned_include_path)) and os.path.exists(os.path.join(sdkpath, versioned_include_path + '.0')):
                      self.sdkversionpath = self.sdkversionpath + '.0'
                      versioned_include_path = os.path.join(include_path, self.sdkversionpath)
                    include_path = versioned_include_path
          except subprocess.CalledProcessError as e:
            continue
          if not sdkpath == '':
            self.includepaths += [
              os.path.join(sdkpath, include_path, 'shared'),
              os.path.join(sdkpath, include_path, 'um'),
              os.path.join(sdkpath, include_path, 'winrt')
            ]
            if version == 'v10.0':
              self.includepaths += [
                os.path.join(sdkpath, include_path, 'ucrt')
              ]
            self.sdkpath = sdkpath
            self.sdkversion = version
            break
        if not sdkpath == '':
          break
    if self.toolchain != '' and not self.toolchain.endswith('/') and not self.toolchain.endswith('\\'):
      self.toolchain += os.sep

  def make_includepaths(self, includepaths):
    if not includepaths is None:
      return ['/I' + self.path_escape(path) for path in list(includepaths)]
    return []

  def make_libpath(self, path):
    return self.path_escape(path)

  def make_libpaths(self, libpaths):
    if not libpaths is None:
      return ['/LIBPATH:' + self.make_libpath(path) for path in libpaths]
    return []

  def make_arch_toolchain_path(self, arch):
    if arch == 'x86-64':
      return os.path.join(self.toolchain, 'bin', 'HostX64', 'x64\\')
    elif arch == 'x86':
      return os.path.join(self.toolchain, 'bin', 'HostX64', 'x86\\')
    return os.path.join(self.toolchain, 'bin\\')

  def make_carchflags(self, arch, targettype):
    flags = []
    if targettype == 'sharedlib':
      flags += ['/MD', '/D', '"BUILD_DYNAMIC_LINK=1"']
    else:
      flags += ['/MT']
    if arch == 'x86':
      flags += ['/arch:SSE2']
    elif arch == 'x86-64':
      pass
    return flags

  def make_cconfigflags(self, config, targettype):
    flags = ['/Gm-']
    if config == 'debug':
      flags += ['/Od', '/D', '"BUILD_DEBUG=1"', '/GF-']
    else:
      flags += ['/Ob2', '/Ot', '/GT', '/GL', '/GF']
      if config == 'release':
        flags += ['/O2', '/D' '"BUILD_RELEASE=1"']
      elif config == 'profile':
        flags += ['/Ox', '/D', '"BUILD_PROFILE=1"']
      elif config == 'deploy':
        flags += ['/Ox', '/D', '"BUILD_DEPLOY=1"']
    return flags

  def make_ararchflags(self, arch, targettype):
    flags = []
    if arch == 'x86':
      flags += ['/MACHINE:X86']
    elif arch == 'x86-64':
      flags += ['/MACHINE:X64']
    return flags

  def make_arconfigflags(self, config, targettype):
    flags = []
    if config != 'debug':
      flags += ['/LTCG']
    return flags

  def make_linkarchflags(self, arch, targettype):
    flags = []
    if arch == 'x86':
      flags += ['/MACHINE:X86']
    elif arch == 'x86-64':
      flags += ['/MACHINE:X64']
    return flags

  def make_linkconfigflags(self, config, targettype):
    flags = []
    if config == 'debug':
      flags += ['/INCREMENTAL']
    else:
      flags += ['/LTCG', '/INCREMENTAL:NO', '/OPT:REF', '/OPT:ICF']
    return flags

  def make_libs(self, libs):
    if libs != None:
      return [lib + '.lib' for lib in libs]
    return []

  def make_configlibpaths(self, config, arch, extralibpaths):
    libpaths = [
      self.libpath,
      os.path.join(self.libpath, arch),
      os.path.join(self.libpath, config),
      os.path.join(self.libpath, config, arch)
      ]
    if extralibpaths != None:
      libpaths += [os.path.join(libpath, self.libpath) for libpath in extralibpaths]
      libpaths += [os.path.join(libpath, self.libpath, arch) for libpath in extralibpaths]
      libpaths += [os.path.join(libpath, self.libpath, config) for libpath in extralibpaths]
      libpaths += [os.path.join(libpath, self.libpath, config, arch) for libpath in extralibpaths]
    if self.sdkpath != '':
      if arch == 'x86':
        libpaths += [os.path.join(self.toolchain, 'lib', 'x86')]
        if self.sdkversion == 'v8.1':
          libpaths += [os.path.join( self.sdkpath, 'lib', 'winv6.3', 'um', 'x86')]
        if self.sdkversion == 'v10.0':
          libpaths += [os.path.join(self.sdkpath, 'lib', self.sdkversionpath, 'um', 'x86')]
          libpaths += [os.path.join(self.sdkpath, 'lib', self.sdkversionpath, 'ucrt', 'x86')]
      else:
        libpaths += [os.path.join( self.toolchain, 'lib', 'x64')]
        if self.sdkversion == 'v8.1':
          libpaths += [os.path.join( self.sdkpath, 'lib', 'winv6.3', 'um', 'x64')]
        if self.sdkversion == 'v10.0':
          libpaths += [os.path.join( self.sdkpath, 'lib', self.sdkversionpath, 'um', 'x64')]
          libpaths += [os.path.join( self.sdkpath, 'lib', self.sdkversionpath, 'ucrt', 'x64')]
    return self.make_libpaths(libpaths)

  def cc_variables(self, config, arch, targettype, variables):
    localvariables = [('toolchain', self.make_arch_toolchain_path(arch))]
    if 'includepaths' in variables:
      moreincludepaths = self.make_includepaths(variables['includepaths'])
      if not moreincludepaths == []:
        localvariables += [('moreincludepaths', moreincludepaths)]
    if 'modulepath' in variables:
      localvariables += [('pdbpath', os.path.join(variables['modulepath'], 'ninja.pdb'))]
    carchflags = self.make_carchflags(arch, targettype)
    if carchflags != []:
      localvariables += [('carchflags', carchflags)]
    cconfigflags = self.make_cconfigflags(config, targettype)
    if cconfigflags != []:
      localvariables += [('cconfigflags', cconfigflags)]
    if 'defines' in variables:
      definelist = []
      for define in variables['defines']:
        definelist += ['/D', '"' + define + '"']
      localvariables += [('cmoreflags', definelist)]
    return localvariables

  def ar_variables(self, config, arch, targettype, variables):
    localvariables = [('toolchain', self.make_arch_toolchain_path(arch))]
    ararchflags = self.make_ararchflags(arch, targettype)
    if ararchflags != []:
      localvariables += [('ararchflags', ararchflags)]
    arconfigflags = self.make_arconfigflags(config, targettype)
    if arconfigflags != []:
      localvariables += [('arconfigflags', arconfigflags)]
    return localvariables

  def link_variables(self, config, arch, targettype, variables):
    localvariables = [('toolchain', self.make_arch_toolchain_path(arch))]
    linkarchflags = self.make_linkarchflags(arch, targettype)
    if linkarchflags != []:
      localvariables += [('linkarchflags', linkarchflags)]
    linkconfigflags = self.make_linkconfigflags(config, targettype)
    if linkconfigflags != []:
      localvariables += [('linkconfigflags', linkconfigflags)]
    if 'modulepath' in variables:
      localvariables += [('pdbpath', os.path.join(variables['modulepath'], 'ninja.pdb'))]
    if 'libs' in variables:
      libvar = self.make_libs(variables['libs'])
      if libvar != []:
        localvariables += [('libs', libvar)]
    libpaths = []
    if 'libpaths' in variables:
      libpaths = variables['libpaths']
    localvariables += [('configlibpaths', self.make_configlibpaths(config, arch, libpaths))]
    return localvariables

  def builder_cc(self, writer, config, arch, targettype, infile, outfile, variables):
    return writer.build(outfile, 'cc', infile, implicit = self.implicit_deps(config, variables), variables = self.cc_variables(config, arch, targettype, variables))

  def builder_cxx(self, writer, config, arch, targettype, infile, outfile, variables):
    return writer.build(outfile, 'cxx', infile, implicit = self.implicit_deps(config, variables), variables = self.cc_variables(config, arch, targettype, variables))

  def builder_lib(self, writer, config, arch, targettype, infiles, outfile, variables):
    return writer.build(outfile, 'ar', infiles, implicit = self.implicit_deps(config, variables), variables = self.ar_variables(config, arch, targettype, variables))

  def builder_sharedlib(self, writer, config, arch, targettype, infiles, outfile, variables):
    return writer.build(outfile, 'dll', infiles, implicit = self.implicit_deps(config, variables), variables = self.link_variables(config, arch, targettype, variables))

  def builder_bin(self, writer, config, arch, targettype, infiles, outfile, variables):
    return writer.build(outfile, 'link', infiles, implicit = self.implicit_deps(config, variables), variables = self.link_variables(config, arch, targettype, variables))

def create(host, target, toolchain):
  return MSVCToolchain(host, target, toolchain)