#!/usr/bin/env python3

# This script has been gratefully sourced from Martins Mozeiko of HMN
# https://gist.github.com/mmozeiko/7f3162ec2988e81e56d5c4e22cde9977
#
# Further modifications by https://github.com/doy-lee with the primary purpose
# of facilitating multiple versions to be stored in the same root directory.
# ('Redist' in the SDK was unversioned, store it versioned like all other
# folders, skip the downloading of MSVC or the SDK if we only need one of them).
# We have some other useful features ..
#
# - Cache the downloads to ./msvc/Installers instead of a temporary directory
#   meaning downloads that are interrupted can be resumed/reused.
# - Allow downloading of _just_ the SDK or _just_ the VC toolchain.
# - Create scripts for each individual component,
#
#     msvc-{version}.bat
#     win-sdk-{version}.bat
#
# - Allow multiple versions to co-exist in the same installation root
#
# Changelog
# 2024-03-14
# - Fold in mmozeiko's latest changes from revision 2024-03-14
#   a6303b82561a4061c71e8838976143f06f295cc9a8a60cf53fc170cb5b9f3f1d
#   - SSL checks
#   - Removal of vctip.exe (telemetry)
#   - DIA SDK DLL
#
# 2024-02-21
# - Add more environment variables that Windows may use, `WindowsSDKVersion`
#   `WindowsSDKDir`, `VSCMD_ARG_TGT_ARCH`
# - Cache the downloads into the `msvc/Installers` directory to allow reusing
#   or resuming an installation that has failed or been terminated prematurely.
# - Fix output scripts using MSVC/SDK version arguments which could be empty
#   if defaulted.
#
# 2023-04-15
# - Fix "msvc-{version}.bat" script generating trailing "\;" on
#   VCToolsInstallDir environment variable. clang-cl relies on this variable to
#   identify the location of the Visual Studio toolchain or otherwise reports a
#   program invocation error during the linking stage.
#
# 2023-01-30
# - Generate the short-hand version of the msvc-{version}.bat and
#   win-sdk-{version}.bat using the versions passed as the argument parameter.
# - Fix win-sdk-{version}.bat overwriting the INCLUDE and LIB environment
#   variables instead of appending.
#
# 2023-01-28
# - Inital revision from mmozeiko
#   https://gist.github.com/mmozeiko/7f3162ec2988e81e56d5c4e22cde9977/6863f19cb98b933c7535acf3d59ac64268c6bd1b
# - Add individual scripts to source variables for MSVC and Windows 10
#   separately "msvc-{version}.bat" and "win-sdk-{version}.bat"
# - Add '--no-sdk' and '--no-msvc' to prevent the download and installation of
#   the Windows SDK and MSVC respectively.
# - Installation used to create 'Windows Kit/10/Redist' and unpack a D3D and MBN
#   folder without being versioned. These folders are now placed under
#   a versioned sub-directory to preserve the binaries and allow subsequent
#   side-by-side installation of other versions of the SDK.

import io
import os
import sys
import json
import shutil
import hashlib
import zipfile
import tempfile
import argparse
import subprocess
import urllib.request
import urllib.error
from pathlib import Path

OUTPUT = Path("msvc") # output folder

# other architectures may work or may not - not really tested
HOST   = "x64" # or x86
TARGET = "x64" # or x86, arm, arm64

MANIFEST_URL = "https://aka.ms/vs/17/release/channel"

ssl_context = None

def download(url):
  with urllib.request.urlopen(url, context=ssl_context) as res:
    return res.read()

def download_progress(url, check, name, f):
  data = io.BytesIO()
  with urllib.request.urlopen(url, context=ssl_context) as res:
    total = int(res.headers["Content-Length"])
    size = 0
    while True:
      block = res.read(1<<20)
      if not block:
        break
      f.write(block)
      data.write(block)
      size += len(block)
      perc = size * 100 // total
      print(f"\r{name} ... {perc}%", end="")
  print()
  data = data.getvalue()
  digest = hashlib.sha256(data).hexdigest()
  if check.lower() != digest:
    exit(f"Hash mismatch for f{pkg}")
  return data

# super crappy msi format parser just to find required .cab files
def get_msi_cabs(msi):
  index = 0
  while True:
    index = msi.find(b".cab", index+4)
    if index < 0:
      return
    yield msi[index-32:index+4].decode("ascii")

def first(items, cond):
  return next(item for item in items if cond(item))
  

### parse command-line arguments

ap = argparse.ArgumentParser()
ap.add_argument("--show-versions", const=True, action="store_const", help="Show available MSVC and Windows SDK versions")
ap.add_argument("--accept-license", const=True, action="store_const", help="Automatically accept license")
ap.add_argument("--msvc-version", help="Get specific MSVC version")
ap.add_argument("--sdk-version", help="Get specific Windows SDK version")
ap.add_argument("--no-msvc", const=True, action="store_const", help="Skip download and installing of msvc")
ap.add_argument("--no-sdk", const=True, action="store_const", help="Skip download and installing of Windows SDK")
args = ap.parse_args()

### get main manifest

try:
  manifest = json.loads(download(MANIFEST_URL))
except urllib.error.URLError as err:
  import ssl
  if isinstance(err.args[0], ssl.SSLCertVerificationError):
    # for more info about Python & issues with Windows certificates see https://stackoverflow.com/a/52074591
    print("ERROR: ssl certificate verification error")
    try:
      import certifi
    except ModuleNotFoundError:
      print("ERROR: please install 'certifi' package to use Mozilla certificates")
      print("ERROR: or update your Windows certs, see instructions here: https://woshub.com/updating-trusted-root-certificates-in-windows-10/#h2_3")
      exit()
    print("NOTE: retrying with certifi certificates")
    ssl_context = ssl.create_default_context(cafile=certifi.where())
    manifest = json.loads(download(MANIFEST_URL))
  else:
    raise


### download VS manifest

vs = first(manifest["channelItems"], lambda x: x["id"] == "Microsoft.VisualStudio.Manifests.VisualStudio")
payload = vs["payloads"][0]["url"]

vsmanifest = json.loads(download(payload))


### find MSVC & WinSDK versions

packages = {}
for p in vsmanifest["packages"]:
  packages.setdefault(p["id"].lower(), []).append(p)

msvc = {}
sdk = {}

for pid,p in packages.items():
  if pid.startswith("Microsoft.VisualStudio.Component.VC.".lower()) and pid.endswith(".x86.x64".lower()):
    pver = ".".join(pid.split(".")[4:6])
    if pver[0].isnumeric():
      msvc[pver] = pid
  elif pid.startswith("Microsoft.VisualStudio.Component.Windows10SDK.".lower()) or \
       pid.startswith("Microsoft.VisualStudio.Component.Windows11SDK.".lower()):
    pver = pid.split(".")[-1]
    if pver.isnumeric():
      sdk[pver] = pid

if args.show_versions:
  print("MSVC versions:", " ".join(sorted(msvc.keys())))
  print("Windows SDK versions:", " ".join(sorted(sdk.keys())))
  exit(0)

install_sdk  = not args.no_sdk
install_msvc = not args.no_msvc
if args.no_sdk and args.no_msvc:
    exit()

msvc_ver_key = args.msvc_version or max(sorted(msvc.keys()))
sdk_ver_key  = args.sdk_version or max(sorted(sdk.keys()))
msvc_ver     = msvc_ver_key
sdk_ver      = sdk_ver_key

info_line = "Downloading"
if install_msvc:
    if msvc_ver in msvc:
      msvc_pid = msvc[msvc_ver]
      msvc_ver = ".".join(msvc_pid.split(".")[4:-2])
    else:
      exit(f"Unknown MSVC version: {args.msvc_version}")
    info_line += f" MSVC v{msvc_ver}"

if install_sdk:
    if sdk_ver in sdk:
      sdk_pid = sdk[sdk_ver]
    else:
      exit(f"Unknown Windows SDK version: {args.sdk_version}")
    info_line += f" Windows SDK v{sdk_ver}"

print(info_line)


### agree to license

tools = first(manifest["channelItems"], lambda x: x["id"] == "Microsoft.VisualStudio.Product.BuildTools")
resource = first(tools["localizedResources"], lambda x: x["language"] == "en-us")
license = resource["license"]

if not args.accept_license:
  accept = input(f"Do you accept Visual Studio license at {license} [Y/N] ? ")
  if not accept or accept[0].lower() != "y":
    exit(0)

OUTPUT.mkdir(exist_ok=True)
total_download = 0

OUTPUT_CACHE_DIR = OUTPUT / "Installers"
OUTPUT_CACHE_DIR.mkdir(exist_ok=True)

OUTPUT_CACHE_MSVC_DIR    = OUTPUT_CACHE_DIR / f"msvc-{msvc_ver}-{HOST}-{TARGET}"
OUTPUT_CACHE_WIN_SDK_DIR = OUTPUT_CACHE_DIR / f"win-sdk-{sdk_ver}-{HOST}-{TARGET}"

### download MSVC

if install_msvc:
    msvc_packages = [
      # MSVC binaries
      f"microsoft.vc.{msvc_ver}.tools.host{HOST}.target{TARGET}.base",
      f"microsoft.vc.{msvc_ver}.tools.host{HOST}.target{TARGET}.res.base",
      # MSVC headers
      f"microsoft.vc.{msvc_ver}.crt.headers.base",
      # MSVC libs
      f"microsoft.vc.{msvc_ver}.crt.{TARGET}.desktop.base",
      f"microsoft.vc.{msvc_ver}.crt.{TARGET}.store.base",
      # MSVC runtime source
      f"microsoft.vc.{msvc_ver}.crt.source.base",
      # ASAN
      f"microsoft.vc.{msvc_ver}.asan.headers.base",
      f"microsoft.vc.{msvc_ver}.asan.{TARGET}.base",
      # MSVC redist
      #f"microsoft.vc.{msvc_ver}.crt.redist.x64.base",
    ]

    OUTPUT_CACHE_MSVC_DIR.mkdir(exist_ok=True)

    for pkg in msvc_packages:
      p = first(packages[pkg], lambda p: p.get("language") in (None, "en-US"))
      for payload in p["payloads"]:
        pkg_path = OUTPUT_CACHE_MSVC_DIR / pkg
        if not os.path.exists(pkg_path):
          with tempfile.NamedTemporaryFile(delete=False) as f:
            data = download_progress(payload["url"], payload["sha256"], pkg, f)
            total_download += len(data)
          shutil.move(f.name, pkg_path)

        with zipfile.ZipFile(pkg_path) as z:
          for name in z.namelist():
            if name.startswith("Contents/"):
              out = OUTPUT / Path(name).relative_to("Contents")
              out.parent.mkdir(parents=True, exist_ok=True)
              out.write_bytes(z.read(name))


### download Windows SDK

if install_sdk:
    sdk_packages = [
      # Windows SDK tools (like rc.exe & mt.exe)
      f"Windows SDK for Windows Store Apps Tools-x86_en-us.msi",
      # Windows SDK headers
      f"Windows SDK for Windows Store Apps Headers-x86_en-us.msi",
      f"Windows SDK Desktop Headers x86-x86_en-us.msi",
      # Windows SDK libs
      f"Windows SDK for Windows Store Apps Libs-x86_en-us.msi",
      f"Windows SDK Desktop Libs {TARGET}-x86_en-us.msi",
      # CRT headers & libs
      f"Universal CRT Headers Libraries and Sources-x86_en-us.msi",
      # CRT redist
      #"Universal CRT Redistributable-x86_en-us.msi",
    ]

    sdk_pkg       = packages[sdk_pid][0]
    sdk_pkg       = packages[first(sdk_pkg["dependencies"], lambda x: True).lower()][0]
    OUTPUT_CACHE_WIN_SDK_DIR.mkdir(exist_ok=True)

    msi = []
    cabs = []

    # download msi files
    for pkg in sdk_packages:
      payload = first(sdk_pkg["payloads"], lambda p: p["fileName"] == f"Installers\\{pkg}")
      msi_path = OUTPUT_CACHE_WIN_SDK_DIR / pkg
      msi.append(msi_path)
      if not os.path.exists(msi_path):
        with tempfile.NamedTemporaryFile(delete=False) as f:
          data = download_progress(payload["url"], payload["sha256"], pkg, f)
          total_download += len(data)
        shutil.move(f.name, msi_path)

      with open(msi_path, "rb") as f:
        cabs += list(get_msi_cabs(f.read()))

    # download .cab files
    for pkg in cabs:
      payload = first(sdk_pkg["payloads"], lambda p: p["fileName"] == f"Installers\\{pkg}")
      cab_path = OUTPUT_CACHE_WIN_SDK_DIR / pkg
      if not os.path.exists(cab_path):
        with tempfile.NamedTemporaryFile(delete=False) as f:
          download_progress(payload["url"], payload["sha256"], pkg, f)
        shutil.move(f.name, cab_path)

    print("Unpacking msi files...")

    # run msi installers
    for m in msi:
      subprocess.check_call(["msiexec.exe", "/a", m, "/quiet", "/qn", f"TARGETDIR={OUTPUT.resolve()}"])


### versions

msvcv = ""
sdkv  = ""

if install_msvc:
    msvcv = list((OUTPUT / "VC/Tools/MSVC").glob(f"{msvc_ver_key}*"))[0].name

if install_sdk:
    sdkv = list((OUTPUT / "Windows Kits/10/bin").glob(f"*{sdk_ver_key}*"))[0].name


if install_msvc:
    # place debug CRT runtime into MSVC folder (not what real Visual Studio installer does... but is reasonable)
    dst = OUTPUT / "VC/Tools/MSVC" / msvcv / f"bin/Host{HOST}/{TARGET}"

    pkg = "microsoft.visualcpp.runtimedebug.14"
    dbg = first(packages[pkg], lambda p: p["chip"] == HOST)
    for payload in dbg["payloads"]:
      name = payload["fileName"]
      pkg_dir  = OUTPUT_CACHE_MSVC_DIR /pkg
      pkg_dir.mkdir(exist_ok=True)
      pkg_path = pkg_dir / name
      if not os.path.exists(pkg_path):
        with tempfile.NamedTemporaryFile(delete=False) as f:
          data = download_progress(payload["url"], payload["sha256"], f"{pkg}/{name}", f)
          total_download += len(data)
        shutil.move(f.name, pkg_path)
    msi = OUTPUT_CACHE_MSVC_DIR / pkg / first(dbg["payloads"], lambda p: p["fileName"].endswith(".msi"))["fileName"]

    with tempfile.TemporaryDirectory() as d2:
      subprocess.check_call(["msiexec.exe", "/a", str(msi), "/quiet", "/qn", f"TARGETDIR={d2}"])
      for f in first(Path(d2).glob("System*"), lambda x: True).iterdir():
        f.replace(dst / f.name)

    # download DIA SDK and put msdia140.dll file into MSVC folder
    pkg = "microsoft.visualc.140.dia.sdk.msi"
    dia = packages[pkg][0]
    for payload in dia["payloads"]:
      name    = payload["fileName"]
      pkg_dir = OUTPUT_CACHE_MSVC_DIR /pkg
      pkg_dir.mkdir(exist_ok=True)
      pkg_path = OUTPUT_CACHE_MSVC_DIR / pkg / name
      if not os.path.exists(pkg_path):
        with tempfile.NamedTemporaryFile(delete=False) as f:
          data = download_progress(payload["url"], payload["sha256"], f"{pkg}/{name}", f)
          total_download += len(data)
        shutil.move(f.name, pkg_path)
    msi = OUTPUT_CACHE_MSVC_DIR / pkg / first(dia["payloads"], lambda p: p["fileName"].endswith(".msi"))["fileName"]

    with tempfile.TemporaryDirectory() as d2:
      subprocess.check_call(["msiexec.exe", "/a", str(msi), "/quiet", "/qn", f"TARGETDIR={d2}"])

      if HOST == "x86": msdia = "msdia140.dll"
      elif HOST == "x64": msdia = "amd64/msdia140.dll"
      else: exit("unknown")

      src = Path(d2) / "Program Files" / "Microsoft Visual Studio 14.0" / "DIA SDK" / "bin" / msdia
      src.replace(dst / "msdia140.dll")

# place the folders under the Redist folder in the SDK under a versioned folder to allow other versions to be installed

if install_sdk:
    redist_dir           = OUTPUT / "Windows Kits/10/Redist"
    redist_versioned_dir = redist_dir / f'{sdkv}'
    shutil.rmtree(redist_versioned_dir, ignore_errors=True)
    redist_versioned_dir.mkdir(exist_ok=True)

    for file_name in os.listdir(redist_dir):
        if not file_name.startswith('10.0.'): # Simple heuristic
          shutil.move((redist_dir / file_name), redist_versioned_dir)

### cleanup

shutil.rmtree(OUTPUT / "Common7", ignore_errors=True)
for f in ["Auxiliary", f"lib/{TARGET}/store", f"lib/{TARGET}/uwp"]:
  shutil.rmtree(OUTPUT / "VC/Tools/MSVC" / msvcv / f, ignore_errors=True)
for f in OUTPUT.glob("*.msi"):
  f.unlink()
for f in ["Catalogs", "DesignTime", f"bin/{sdkv}/chpe", f"Lib/{sdkv}/ucrt_enclave"]:
  shutil.rmtree(OUTPUT / "Windows Kits/10" / f, ignore_errors=True)
for arch in ["x86", "x64", "arm", "arm64"]:
  if arch != TARGET:
    shutil.rmtree(OUTPUT / "Windows Kits/10/Lib" / sdkv / "ucrt" / arch, ignore_errors=True)
    shutil.rmtree(OUTPUT / "Windows Kits/10/Lib" / sdkv / "um" / arch, ignore_errors=True)
  if arch != HOST:
    shutil.rmtree(OUTPUT / "Windows Kits/10/bin" / sdkv / arch, ignore_errors=True)
# executable that is collecting & sending telemetry every time cl/link runs
(OUTPUT / "VC/Tools/MSVC" / msvcv / f"bin/Host{HOST}/{TARGET}/vctip.exe").unlink(missing_ok=True)


### setup.bat

if install_msvc and install_sdk:
    SETUP = f"""@echo off

set MSVC_VERSION={msvcv}
set MSVC_HOST=Host{HOST}
set MSVC_ARCH={TARGET}
set MSVC_ROOT=%~dp0VC\\Tools\\MSVC\\%MSVC_VERSION%

set SDK_VERSION={sdkv}
set SDK_ARCH={TARGET}
set SDK_INCLUDE=%~dp0Windows Kits\\10\\Include\\%SDK_VERSION%
set SDK_LIBS=%~dp0Windows Kits\\10\\Lib\\%SDK_VERSION%

set WindowsSDKVersion=%SDK_VERSION%
set WindowsSDKDir=%~dp0Windows Kits\\10

set VCToolsInstallDir=%MSVC_ROOT%
set VSCMD_ARG_TGT_ARCH=%MSVC_ARCH%

set PATH=%MSVC_ROOT%\\bin\\%MSVC_HOST%\\%MSVC_ARCH%;%~dp0Windows Kits\\10\\bin\\%SDK_VERSION%\\%SDK_ARCH%;%~dp0Windows Kits\\10\\bin\\%SDK_VERSION%\\%SDK_ARCH%\\ucrt;%PATH%
set INCLUDE=%MSVC_ROOT%\\include;%SDK_INCLUDE%\\ucrt;%SDK_INCLUDE%\\shared;%SDK_INCLUDE%\\um;%SDK_INCLUDE%\\winrt;%SDK_INCLUDE%\\cppwinrt
set LIB=%MSVC_ROOT%\\lib\\%MSVC_ARCH%;%SDK_LIBS%\\ucrt\\%SDK_ARCH%;%SDK_LIBS%\\um\\%SDK_ARCH%
    """
    (OUTPUT / "setup.bat").write_text(SETUP)

if install_msvc:
    MSVC_SCRIPT = f"""@echo off

set MSVC_VERSION={msvcv}
set MSVC_HOST=Host{HOST}
set MSVC_ARCH={TARGET}
set MSVC_ROOT=%~dp0VC\\Tools\\MSVC\\%MSVC_VERSION%

set VCToolsInstallDir=%MSVC_ROOT%
set VSCMD_ARG_TGT_ARCH=%MSVC_ARCH%

set PATH=%MSVC_ROOT%\\bin\\%MSVC_HOST%\\%MSVC_ARCH%;%PATH%
set INCLUDE=%MSVC_ROOT%\\include;%INCLUDE%
set LIB=%MSVC_ROOT%\\lib\\%MSVC_ARCH%;%LIB%
"""
    (OUTPUT / f"msvc-{msvcv}.bat").write_text(MSVC_SCRIPT)
    (OUTPUT / f"msvc-{msvc_ver_key}.bat").write_text(MSVC_SCRIPT)

if install_sdk:
    WIN10_SDK_SCRIPT = f"""@echo off

set SDK_VERSION={sdkv}
set SDK_ARCH={TARGET}
set SDK_INCLUDE=%~dp0Windows Kits\\10\\Include\\%SDK_VERSION%
set SDK_LIBS=%~dp0Windows Kits\\10\\Lib\\%SDK_VERSION%

set WindowsSDKVersion=%SDK_VERSION%
set WindowsSDKDir=%~dp0Windows Kits\\10

set PATH=%~dp0Windows Kits\\10\\bin\\%SDK_VERSION%\\%SDK_ARCH%;%~dp0Windows Kits\\10\\bin\\%SDK_VERSION%\\%SDK_ARCH%\\ucrt;%PATH%
set INCLUDE=%SDK_INCLUDE%\\ucrt;%SDK_INCLUDE%\\shared;%SDK_INCLUDE%\\um;%SDK_INCLUDE%\\winrt;%SDK_INCLUDE%\\cppwinrt;%INCLUDE%
set LIB=%SDK_LIBS%\\ucrt\\%SDK_ARCH%;%SDK_LIBS%\\um\\%SDK_ARCH%;%LIB%
"""
    (OUTPUT / f"win-sdk-{sdkv}.bat").write_text(WIN10_SDK_SCRIPT)
    (OUTPUT / f"win-sdk-{sdk_ver_key}.bat").write_text(WIN10_SDK_SCRIPT)

print(f"Total downloaded: {total_download>>20} MB")
print("Done!")