#!/usr/bin/env python3
import argparse
import hashlib
import re
import sys
import urllib.request
from pathlib import Path

REDIST_URL = "https://aka.ms/vs/17/release/vc_redist.x64.exe"
BEGIN = "; >>> VC_REDIST_INSTALL >>>"
END = "; <<< VC_REDIST_INSTALL <<<"

REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_CARGO_TOML = REPO_ROOT / "packages" / "oneclient_app" / "Cargo.toml"
DEFAULT_REDIST = REPO_ROOT / "target" / "vcredist" / "vc_redist.x64.exe"

BLOCK = """  {begin}
  ; Generated by distribution/windows/bundle-vcredist.py -- do not edit by hand.
  DetailPrint "Installing Microsoft Visual C++ Redistributable (x64)..."
  InitPluginsDir
  File "/oname=$PLUGINSDIR\\\\vc_redist.x64.exe" "{redist}"
  ExecWait '"$PLUGINSDIR\\\\vc_redist.x64.exe" /quiet /norestart' $R3
  ${{If}} $R3 == 3010
    SetRebootFlag true
  ${{EndIf}}
  {end}"""

REGION = re.compile(
    r"^[ \t]*" + re.escape(BEGIN) + r".*?^[ \t]*" + re.escape(END) + r"[ \t]*$",
    re.DOTALL | re.MULTILINE,
)


def download(dest: Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    print(f"downloading {REDIST_URL} -> {dest}")
    with urllib.request.urlopen(REDIST_URL) as resp, dest.open("wb") as out:
        while chunk := resp.read(1 << 20):
            out.write(chunk)

    data = dest.read_bytes()
    if data[:2] != b"MZ":
        sys.exit(f"error: {dest} is not a PE image (no MZ header)")
    if len(data) < 10 * 1024 * 1024:
        sys.exit(f"error: {dest} is only {len(data)} bytes, expected >10MB")
    print(f"  {len(data)} bytes, sha256={hashlib.sha256(data).hexdigest()}")


def patch(cargo_toml: Path, redist: Path) -> None:
    text = cargo_toml.read_text(encoding="utf-8")
    if not REGION.search(text):
        sys.exit(f"error: {BEGIN} .. {END} region not found in {cargo_toml}")

    block = BLOCK.format(begin=BEGIN, end=END, redist=str(redist).replace("\\", "\\\\"))
    cargo_toml.write_text(REGION.sub(lambda _: block, text, count=1), encoding="utf-8")
    print(f"patched {cargo_toml} to embed {redist}")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--cargo-toml", type=Path, default=DEFAULT_CARGO_TOML)
    ap.add_argument("--redist", type=Path, default=DEFAULT_REDIST)
    ap.add_argument(
        "--skip-download",
        action="store_true",
        help="reuse an already-downloaded --redist instead of fetching it",
    )
    args = ap.parse_args()

    redist = args.redist.resolve()
    if args.skip_download:
        if not redist.is_file():
            sys.exit(f"error: --skip-download given but {redist} does not exist")
    else:
        download(redist)

    patch(args.cargo_toml.resolve(), redist)


if __name__ == "__main__":
    main()