#!python

import argparse
import contextlib
import inspect
import os
import os.path
import pathlib
import sys
import subprocess
import typing

whl_platform = "linux_roborio"

setup_py_tmpl = inspect.cleandoc(
    """
    #!/usr/bin/env python3
    #
    # Automatically generated by roborio-gen-whl
    #

    from setuptools import setup
    
    setup(
        name=%(name)r,
        version=%(version)r,
        license=%(license)r,
        install_requires=%(install_requires)r,
        url=%(url)r,
        maintainer="RobotPy",
        maintainer_email="robotpy@googlegroups.com",
        long_description=%(description)r,
        long_description_content_type="text/markdown",
        data_files=%(data_files)r,
    )

"""
)


@contextlib.contextmanager
def switch_cwd(path: pathlib.Path):
    oldcwd = os.getcwd()
    os.chdir(path)
    try:
        yield path
    finally:
        os.chdir(oldcwd)


def get_soname(path: str) -> str:
    output = subprocess.check_output(["objdump", "-p", path], encoding="utf-8")
    for line in output.splitlines():
        if line.startswith("  SONAME"):
            _, libname = line.split()
            return libname

    raise ValueError(f"SONAME not found for `{path}`")


def enforce_structure(
    prefix: str, strip: typing.Optional[pathlib.Path], soname_must_match: bool
) -> typing.List[typing.Tuple[str, typing.List[str]]]:
    # Can only install files to the python prefix, which is /usr/local
    data_files = []
    for root, _, files in os.walk(prefix):
        root_files = []
        for fname in files:
            full_fname = os.path.join(root, fname)
            if os.path.islink(full_fname):
                raise ValueError(f"{full_fname}: is a symlink")

            if (fname.endswith("so") or ".so." in fname) and not fname.endswith(
                ".debug"
            ):
                if strip:
                    print("+", strip, full_fname)
                    subprocess.check_call([strip, full_fname])

                if not fname.endswith("-arm-linux-gnueabi.so"):
                    soname = get_soname(full_fname)
                    if soname_must_match != (soname == fname):
                        raise ValueError(f"{full_fname}: SONAME is {soname}")

            root_files.append(full_fname)

        if root_files:
            root = os.path.relpath(root, prefix)
            data_files.append((root, root_files))

    return data_files


def create_wheel(
    data_py_fname: str,
    data_root: pathlib.Path,
    dev: bool,
    dbg: bool,
    strip: typing.Optional[pathlib.Path],
    output_path: pathlib.Path,
):
    """
    Generates a setup.py to create a python wheel for non-python artifacts.
    It's a bit of a hack, but it frees us to only have users deal with pip
    when installing packages.

    This can only install artifacts to the python 'prefix', which on RoboRIO
    we have set to /usr/local

    Ideally, we would create platform-specific wheels. Unfortunately, wheels
    do not currently support symlinks
    """
    data = {}

    with open(data_py_fname) as fp:
        exec(fp.read(), data)

    if dbg or dev:
        data["install_requires"] = [data["name"] + "==" + data["version"]]
    if dbg:
        data["name"] = data["name"] + "-dbg"
    if dev:
        data["name"] = data["name"] + "-dev"
        data["install_requires"].extend(data.get("install_dev_requires", []))

    # hack for implicit dev packages
    if data["name"].endswith("-dev"):
        dev = True

    # Add description to make 'twine check' happy
    data["description"] = f"{data['name']}\n{'='*len(data['name'])}"

    # enforce library restrictions
    # -> normal package must have .so that match SONAME
    # -> dev package must not have .so that match SONAME

    with switch_cwd(data_root):
        data_files = enforce_structure(
            os.path.join("usr", "local"), strip, soname_must_match=not dev
        )
    data["data_files"] = data_files

    setup_py = setup_py_tmpl % data

    with open(data_root / "setup.py", "w") as fp:
        fp.write(setup_py)

    subprocess.check_call(
        [
            sys.executable,
            "setup.py",
            "bdist_wheel",
            "-p",
            whl_platform,
            "-d",
            output_path.absolute(),
        ],
        cwd=data_root,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("data_py", type=pathlib.Path)
    parser.add_argument("data_root", type=pathlib.Path)
    parser.add_argument("--dev", action="store_true", default=False)
    parser.add_argument("--dbg", action="store_true", default=False)
    parser.add_argument(
        "-o", "--output-path", type=pathlib.Path, default=pathlib.Path(".")
    )
    parser.add_argument(
        "--strip",
        help="Executable to call to strip any shared libraries present",
        default=None,
        type=pathlib.Path,
    )

    args = parser.parse_args()

    create_wheel(
        args.data_py, args.data_root, args.dev, args.dbg, args.strip, args.output_path
    )
