#!/usr/bin/env python3
# Copyright (C) 2026 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Packages the Perfetto UI, as deployed on ui.perfetto.dev, for self-hosting.

For a given release tag, this downloads the exact UI build that the
stable-channel Cloud Build deploy publishes to the public
gs://ui.perfetto.dev bucket, and zips it up into a tree that can be dropped
onto any static file server to self-host that release's UI.

Usage: ./tools/release/package-ui-github-release-artifact v57.1

This will generate perfetto-ui.zip in /tmp/perfetto-v57.1-github-release/
(the same staging dir package-github-release-artifacts uses, so a single
`gh release upload $STAGING/*.zip` picks up both).
"""

import argparse
import base64
import concurrent.futures
import hashlib
import json
import os
import re
import subprocess
import sys
import tempfile
import urllib.error
import urllib.request
import zipfile

_GCS_ROOT = 'https://commondatastorage.googleapis.com/ui.perfetto.dev'
_REPO_ROOT = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..', '..'))
_NUM_WORKERS = 16


def _fetch(url):
  with urllib.request.urlopen(url) as resp:
    return resp.read()


def _resolve_ui_version(version):
  try:
    sha = subprocess.check_output(['git', 'rev-parse', f'{version}^{{commit}}'],
                                  cwd=_REPO_ROOT,
                                  stderr=subprocess.PIPE,
                                  text=True).strip()
  except subprocess.CalledProcessError as e:
    raise SystemExit(
        f'ERROR: tag {version} not found in this checkout. Fetch tags '
        f'(git fetch --tags) and re-run.\n{e.stderr}')
  return f'{version}-{sha[:9]}'


def _download_manifest(ui_version):
  url = f'{_GCS_ROOT}/{ui_version}/manifest.json'
  try:
    return json.loads(_fetch(url))
  except urllib.error.HTTPError as e:
    if e.code == 404:
      raise SystemExit(
          f'ERROR: {url} not found (404). The UI for {ui_version} has not '
          'been deployed to ui.perfetto.dev yet: the stable-channel Cloud '
          'Build deploy runs on push to the stable branch. Wait for it to '
          'complete and re-run this script.')
    raise


def _download_resource(ui_dir, ui_version, relpath, digest):
  url = f'{_GCS_ROOT}/{ui_version}/{relpath}'
  data = _fetch(url)
  actual = 'sha256-' + base64.b64encode(hashlib.sha256(data).digest()).decode()
  if actual != digest:
    raise SystemExit(
        f'ERROR: checksum mismatch for {relpath}: manifest says {digest}, '
        f'downloaded content hashes to {actual}')
  dst = os.path.join(ui_dir, ui_version, relpath)
  os.makedirs(os.path.dirname(dst), exist_ok=True)
  with open(dst, 'wb') as f:
    f.write(data)


def _download_ui_tree(ui_dir, ui_version, manifest):
  resources = manifest['resources']
  print(f'--- Downloading {len(resources)} resources for {ui_version} ---')
  with concurrent.futures.ThreadPoolExecutor(max_workers=_NUM_WORKERS) as pool:
    futures = [
        pool.submit(_download_resource, ui_dir, ui_version, relpath, digest)
        for relpath, digest in resources.items()
    ]
    for future in concurrent.futures.as_completed(futures):
      future.result()  # Re-raises on failure.

  version_dir = os.path.join(ui_dir, ui_version)
  os.makedirs(version_dir, exist_ok=True)
  for extra in ('index.html', 'manifest.json'):
    data = _fetch(f'{_GCS_ROOT}/{ui_version}/{extra}')
    with open(os.path.join(version_dir, extra), 'wb') as f:
      f.write(data)


def _download_root_files(ui_dir):
  sw_path = os.path.join(ui_dir, 'service_worker.js')
  with open(sw_path, 'wb') as f:
    f.write(_fetch(f'{_GCS_ROOT}/service_worker.js'))

  try:
    data = _fetch(f'{_GCS_ROOT}/service_worker.js.map')
  except urllib.error.HTTPError as e:
    if e.code != 404:
      raise
    print('NOTE: service_worker.js.map not found on the bucket; skipping.')
    return
  with open(os.path.join(ui_dir, 'service_worker.js.map'), 'wb') as f:
    f.write(data)


def _write_root_index_html(ui_dir, ui_version):
  version_index = os.path.join(ui_dir, ui_version, 'index.html')
  with open(version_index) as f:
    html = f.read()

  version_map = json.dumps({'stable': ui_version})
  patched, count = re.subn(
      r"data-perfetto_version='[^']*'",
      f"data-perfetto_version='{version_map}'",
      html,
      count=1)
  if count == 0:
    raise SystemExit(
        "ERROR: data-perfetto_version attribute not found in the downloaded "
        f'{ui_version}/index.html; the archival copy format may have '
        'changed.')

  with open(os.path.join(ui_dir, 'index.html'), 'w') as f:
    f.write(patched)


def _zip_ui_dir(ui_dir, zip_path):
  file_count = 0
  with zipfile.ZipFile(
      zip_path, 'w', zipfile.ZIP_DEFLATED, compresslevel=9) as zf:
    for root, _, files in os.walk(ui_dir):
      for name in sorted(files):
        abs_path = os.path.join(root, name)
        arcname = os.path.relpath(abs_path, ui_dir)
        zf.write(abs_path, arcname)
        file_count += 1
  return file_count


def main():
  parser = argparse.ArgumentParser(epilog='Example: %s v57.1' % __file__)
  parser.add_argument('version', help='Release version tag (e.g. v57.1)')
  parser.add_argument(
      '--out-dir',
      help='Staging dir (default: /tmp/perfetto-<version>-github-release, '
      'matching package-github-release-artifacts).')
  args = parser.parse_args()

  if not re.match(r'^v\d+\.\d+$', args.version):
    print(
        f'ERROR: version must match vX.Y (got: {args.version})',
        file=sys.stderr)
    return 1

  out_dir = args.out_dir or f'/tmp/perfetto-{args.version}-github-release'
  os.makedirs(out_dir, exist_ok=True)

  ui_version = _resolve_ui_version(args.version)
  print(f'Resolved {args.version} -> ui_version {ui_version}')

  manifest = _download_manifest(ui_version)

  # Stage the tree outside out_dir: package-github-release-artifacts'
  # verify_downloads() hard-fails on unexpected directories in there.
  with tempfile.TemporaryDirectory(prefix='perfetto-ui-') as ui_dir:
    _download_ui_tree(ui_dir, ui_version, manifest)
    _download_root_files(ui_dir)
    _write_root_index_html(ui_dir, ui_version)

    zip_path = os.path.join(out_dir, 'perfetto-ui.zip')
    file_count = _zip_ui_dir(ui_dir, zip_path)

  print('')
  print('=' * 70)
  print(f'Packaged {file_count} files into {zip_path}')
  print(f'Zip size: {os.path.getsize(zip_path)} bytes')
  print('=' * 70)
  return 0


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