#!/usr/bin/env python3
# Copyright 2023 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Update in-tree checkout of Rust toolchain

Minimal Skia adapation of [1] to fetch a Rust revision manually
specified in MANUAL_REVISION below. Frequently roll this to latest
version available in [2].

[1] https://source.chromium.org/chromium/chromium/src/+/main:tools/rust/update_rust.py
[2] https://commondatastorage.googleapis.com/chromium-browser-clang/index.html?path=Linux_x64/rust-toolchain-

"""

import argparse
import os
import re
import shutil
import sys
import tempfile
import urllib

from pathlib import Path

# Chromium's Rust builds (for Linux) that worked are found at:
# https://commondatastorage.googleapis.com/chromium-browser-clang/index.html?path=Linux_x64/rust-toolchain-
# The latest builds are prefixed with a date, such as `20230101-1`.
# To update, roll this to versions that Chromium's tools/rust/build-rust.py has produced and which are found from the CDS url.
MANUAL_REVISION = (
    "ac4379fea9e83465d814bb05005689f49bd2141e-1-llvmorg-17-init-3874-g93a2fecc-1"
)

THIS_DIR = os.path.abspath(os.path.dirname(__file__))
SKIA_DIR = os.path.abspath(os.path.join(THIS_DIR, ".."))
THIRD_PARTY_DIR = os.path.join(SKIA_DIR, "third_party")
RUST_TOOLCHAIN_OUT_DIR = os.path.join(THIRD_PARTY_DIR, "rust-toolchain")
VERSION_STAMP_PATH = os.path.join(RUST_TOOLCHAIN_OUT_DIR, "VERSION")


def GetDownloadPackageVersion():
    # TODO(https://crbug.com/14191): This is hardcoded in Skia right
    # now, as Skia does not rebuild rust on its own and can't directly
    # access the most recent version that's been built for
    # Chromium. Could the Chromium side-build publish something like a
    # symbolic link to last-known-good?
    return MANUAL_REVISION


# Get the version of the toolchain package we already have.
def GetStampVersion():
    if os.path.exists(RUST_TOOLCHAIN_OUT_DIR):
        with open(VERSION_STAMP_PATH) as version_file:
            existing_stamp = version_file.readline().rstrip()
        version_re = re.compile(r"rustc [0-9.]+ [0-9a-f]+ \((.+?) chromium\)")
        match = version_re.fullmatch(existing_stamp)
        if match is None:
            return None
        return match.group(1)

    return None


def main():
    parser = argparse.ArgumentParser(description="Update Rust package")
    parser.add_argument(
        "--print-package-version",
        action="store_true",
        help="Print Rust package version (including both the "
        "Rust and Clang revisions) and quit.",
    )
    args = parser.parse_args()

    if args.print_package_version:
        print(GetDownloadPackageVersion())
        return 0

    from clang_update import DownloadAndUnpack, GetDefaultHostOs, GetPlatformUrlPrefix

    # Exit early if the existing package is up-to-date. Note that we cannot
    # simply call DownloadAndUnpack() every time: aside from unnecessarily
    # downloading the toolchain if it hasn't changed, it also leads to multiple
    # versions of the same rustlibs. build/rust/std/find_std_rlibs.py chokes in
    # this case.
    if os.path.exists(RUST_TOOLCHAIN_OUT_DIR):
        if GetDownloadPackageVersion() == GetStampVersion():
            return 0

    if os.path.exists(RUST_TOOLCHAIN_OUT_DIR):
        shutil.rmtree(RUST_TOOLCHAIN_OUT_DIR)

    try:
        host_os = GetDefaultHostOs()
        # TODO(https://crbug.com/skia/14190): Enable this on additional
        # platforms.
        if not "linux" in host_os:
            print(
                "Unsupported platform, Rust support only available on Linux "
                "at the moment, see https://crbug.com/skia/14190"
            )
            return 1
        platform_prefix = GetPlatformUrlPrefix(host_os)
        version = GetDownloadPackageVersion()
        url = f"{platform_prefix}rust-toolchain-{version}.tgz"
        DownloadAndUnpack(url, THIRD_PARTY_DIR)
    except urllib.error.HTTPError as e:
        # Fail softly for now. This can happen if a Rust package was not
        # produced, e.g. if the Rust build failed upon a Clang update, or if a
        # Rust roll and a Clang roll raced against each other.
        #
        # TODO(https://crbug.com/1245714): Reconsider how to handle this.
        print(f"warning: could not download Rust package")

    # Ensure the newly extracted package has the correct version.
    assert GetDownloadPackageVersion() == GetStampVersion()


if __name__ == "__main__":
    sys.exit(main())