#!/usr/bin/env fuchsia-vendored-python
# Copyright 2026 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import collections
import configparser
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

from depfile import DepFile


def main() -> int:
    """Builds a Python wheel from GN-provided python_library and C extension metadata.

    This script reads JSON metadata generated by GN describing Python library sources,
    shared libraries, and data files. It stages these files into a temporary directory,
    generates declarative setup.cfg and minimal setup.py files, invokes `pip wheel` to
    produce the wheel package, and writes a Ninja depfile tracking all input dependencies.
    """
    parser = argparse.ArgumentParser("Builds a Python wheel")
    parser.add_argument("--target_name", required=True)
    parser.add_argument("--gen_dir", type=Path, required=True)
    parser.add_argument(
        "--library_infos", type=argparse.FileType("r"), required=True
    )
    parser.add_argument("--pyproject_toml", type=Path, required=True)
    parser.add_argument("--depfile", type=Path, required=True)
    parser.add_argument("--stamp_file", type=Path, required=True)
    parser.add_argument("--wheel_file", type=Path, required=True)
    parser.add_argument("--python_interpreter", required=True)
    parser.add_argument("--version", default="0.0.1")
    parser.add_argument("--package_name", required=True)
    parser.add_argument("--packages", nargs="*", default=[])
    args = parser.parse_args()

    with args.library_infos as f:
        infos = json.load(f)

    staged_dir = args.gen_dir / f"{args.target_name}_wheel_staged"
    if staged_dir.is_dir():
        shutil.rmtree(staged_dir)
    staged_dir.mkdir(parents=True, exist_ok=True)

    # Copy pyproject.toml to the staging directory.
    # Note: The wheel build generates a setuptools-based setup.cfg / setup.py, assuming
    # setuptools (setuptools.build_meta) as the build backend in pyproject.toml.
    # Other backends (e.g. hatchling, flit) will ignore or conflict with setup.py.
    shutil.copy2(args.pyproject_toml, staged_dir / "pyproject.toml")

    inputs: list[str] = [
        str(args.pyproject_toml),
        args.library_infos.name,
    ]
    outputs: list[str] = [
        str(args.stamp_file),
        str(args.wheel_file),
        str(staged_dir / "pyproject.toml"),
        str(staged_dir / "setup.cfg"),
        str(staged_dir / "setup.py"),
    ]
    packages: list[str] = list(args.packages)
    package_data: collections.defaultdict[
        str, list[str]
    ] = collections.defaultdict(list)
    for lib_info in infos:
        lib_name = lib_info["library_name"]
        dest_lib_dir = staged_dir / lib_name.replace(".", os.sep)
        dest_lib_dir.mkdir(parents=True, exist_ok=True)

        src_root = Path(lib_info["source_root"])
        for src in lib_info["sources"]:
            src_path = src_root / src
            dest_path = dest_lib_dir / src
            dest_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(src_path, dest_path)
            inputs.append(str(src_path))
            outputs.append(str(dest_path))

        if "shared_libs" in lib_info:
            for shlib in lib_info["shared_libs"]:
                shlib_path = Path(shlib)
                dest_path = dest_lib_dir / shlib_path.name
                shutil.copy2(shlib_path, dest_path)
                inputs.append(str(shlib_path))
                outputs.append(str(dest_path))
                package_data[lib_name].append(shlib_path.name)

        if lib_info.get("data_sources") and lib_info.get("data_package_name"):
            data_pkg_name = lib_info["data_package_name"]
            data_pkg_rel_dir = data_pkg_name.replace(".", os.sep)
            dest_data_dir = dest_lib_dir / data_pkg_rel_dir
            dest_data_dir.mkdir(parents=True, exist_ok=True)
            init_file = dest_data_dir / "__init__.py"
            init_file.touch()
            outputs.append(str(init_file))
            full_data_pkg = f"{lib_name}.{data_pkg_name}"
            if full_data_pkg not in packages:
                packages.append(full_data_pkg)
            for data_src in lib_info["data_sources"]:
                data_src_path = Path(data_src)
                dest_path = dest_data_dir / data_src_path.name
                shutil.copy2(data_src_path, dest_path)
                inputs.append(str(data_src_path))
                outputs.append(str(dest_path))
                rel_data_path = (
                    f"{data_pkg_name.replace('.', '/')}/{data_src_path.name}"
                )
                package_data[lib_name].append(rel_data_path)

    has_shared_libs = any(
        f.endswith(".so") for files in package_data.values() for f in files
    )

    # Generate declarative setup.cfg configuration via configparser.
    config = configparser.ConfigParser()
    config["metadata"] = {
        "name": args.package_name,
        "version": args.version,
    }
    config["options"] = {}
    if packages:
        config["options"]["packages"] = "\n" + "\n".join(packages)
    if package_data:
        config["options.package_data"] = {
            pkg: "\n" + "\n".join(files) for pkg, files in package_data.items()
        }
    with open(staged_dir / "setup.cfg", "w", encoding="utf-8") as f:
        config.write(f)

    # Generate minimal setup.py providing the custom Distribution class for binary extensions.
    setup_py_content = f"""from setuptools import setup
from setuptools.dist import Distribution


class BinaryDistribution(Distribution):
    def is_pure(self) -> bool:
        return {not has_shared_libs}

    def has_ext_modules(self) -> bool:
        return {has_shared_libs}


setup(distclass=BinaryDistribution)
"""
    with open(staged_dir / "setup.py", "w", encoding="utf-8") as f:
        f.write(setup_py_content)

    wheel_out_dir = args.gen_dir / f"{args.target_name}_wheel_out"
    if wheel_out_dir.is_dir():
        shutil.rmtree(wheel_out_dir)
    wheel_out_dir.mkdir(parents=True, exist_ok=True)

    cmd_env = os.environ.copy()
    cmd_env["SETUPTOOLS_SCM_PRETEND_VERSION"] = args.version
    cmd_env["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
    # Ensure pip does not implicitly read host/user-level configuration (~/.config/pip/pip.conf).
    cmd_env["PIP_CONFIG_FILE"] = os.devnull
    # Ensure bit-for-bit reproducible wheel zip archive timestamps (1980-01-01 00:00:00 UTC).
    cmd_env["SOURCE_DATE_EPOCH"] = "315532800"
    try:
        subprocess.run(
            [
                args.python_interpreter,
                "-m",
                "pip",
                "wheel",
                "--no-cache-dir",
                "--no-deps",
                "--no-index",
                "--no-build-isolation",
                "-w",
                str(wheel_out_dir),
                str(staged_dir),
            ],
            env=cmd_env,
            capture_output=True,
            text=True,
            check=True,
        )

        # Locate the generated wheel and move it to the declared wheel output path.
        # Note: The wheel filename conforms to the PEP 427 naming convention
        # ({distribution}-{version}(-{build tag})?-{python tag}-{abi tag}-{platform tag}.whl).
        wheels = list(wheel_out_dir.glob("*.whl"))
        if not wheels:
            print(
                f"No wheel generated in '{wheel_out_dir}'.",
                file=sys.stderr,
            )
            return 1
        if len(wheels) > 1:
            print(
                f"Expected exactly one wheel in '{wheel_out_dir}', found: {[w.name for w in wheels]}",
                file=sys.stderr,
            )
            return 1
        generated_wheel = wheels[0]
        args.wheel_file.parent.mkdir(parents=True, exist_ok=True)
        if args.wheel_file.exists() or args.wheel_file.is_symlink():
            args.wheel_file.unlink()
        shutil.move(generated_wheel, args.wheel_file)

        dep_file = DepFile.from_deps(str(args.stamp_file), sorted(set(inputs)))
        for out in sorted(set(outputs)):
            dep_file.add_output(out)
        args.depfile.parent.mkdir(parents=True, exist_ok=True)
        with open(args.depfile, "w", encoding="utf-8") as f:
            dep_file.write_to(f)

        args.stamp_file.touch()
        return 0
    except subprocess.CalledProcessError as e:
        print(
            f"Failed to build wheel (exit code {e.returncode}). Staged files retained at '{staged_dir}' for debugging.",
            file=sys.stderr,
        )
        if e.stdout:
            print(e.stdout, file=sys.stderr)
        if e.stderr:
            print(e.stderr, file=sys.stderr)
        return e.returncode


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