# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import os
import requests
import shutil

# This script updates the current directory with parts of the ONNX Runtime headers
# needed in Chromium.

# The ONNX Runtime headers needed in Chromium
onnxruntime_headers = [
    "onnxruntime_c_api.h",
    "onnxruntime_session_options_config_keys.h",
]

# The target revision of ONNX Runtime
revision = "f217402897f40ebba457e2421bc0a4702771968e"
# These headers are extracted from a git repository and versioned by revision
version = "v1.22.0"

base_url = "https://raw.githubusercontent.com/microsoft/onnxruntime"
onnxruntime_headers_path = ["include", "onnxruntime", "core", "session"]

# Source URL for the current revision of ONNX Runtime headers
src_url = "/".join([base_url, revision, *onnxruntime_headers_path])
# License URL for the ONNX Runtime repository
license_url = "/".join([base_url, revision, "LICENSE"])

abs_path = os.path.dirname(os.path.abspath(__file__))

# Delete the src directory
abs_src_dir = os.path.join(abs_path, "src")
if os.path.exists(abs_src_dir):
    shutil.rmtree(abs_src_dir)

# Create the directory for ONNX Runtime headers
include_dir = os.path.join("src", *onnxruntime_headers_path)
abs_include_dir = os.path.join(abs_path, include_dir)
os.makedirs(abs_include_dir)

# Update the ONNX Runtime headers
for header in onnxruntime_headers:
    with open(os.path.join(abs_include_dir, header), "wb") as f:
        file_url = "/".join([src_url, header])
        response = requests.get(file_url)
        if response.status_code != 200:
            raise Exception("Failed to download", file_url)
        f.write(response.content)

# Update the license file
with open(os.path.join(abs_path, "src/LICENSE"), "wb") as f:
    response = requests.get(license_url)
    if response.status_code != 200:
        raise Exception("Failed to download", file_url)
    f.write(response.content)

# Generate the BUILD.gn file

def format_headers(headers):
    formatted_string = '[\n'
    for i in range(len(headers)):
        header = os.path.join(include_dir, headers[i]).replace("\\", "/")
        formatted_string += f'    "{header}",\n'
    formatted_string += '  ]'
    return formatted_string

build_file = '''# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# This file is generated by Update.py and should not be edited.

# onnxruntime_headers is used by WebNN on Windows.

source_set("onnxruntime_headers") {{
  sources = {sources}

  visibility = [ "//services/webnn/*" ]
}}
'''.format(sources=format_headers(onnxruntime_headers))

with open(os.path.join(abs_path, "BUILD.gn"), "w", newline="\n") as f:
    f.write(build_file)

# Generate the README.chromium file

readme_file = '''Name: Headers for ONNX Runtime C/C++ APIs
Short Name: onnxruntime_headers
URL: https://github.com/microsoft/onnxruntime
Version: {version}
Revision: {revision}
License: MIT
License File: src/LICENSE
Shipped: Yes
Security Critical: Yes

Description:
This folder is a snapshot of part of the ONNX Runtime headers for C/C++
APIs. It is used by WebNN for inference on Windows.

To avoid pulling in unnecessary files this snapshot only contains the
headers that are required by WebNN.

Local Modifications: None

# This file is generated by Update.py and should not be edited.
'''.format(version=version, revision=revision)

with open(os.path.join(abs_path, "README.chromium"), "w", newline="\n") as f:
    f.write(readme_file)