"""Shim setup.py: extracts pre-built wheels vendored under wheels/ at install time.

IQ.Lvbs fork of commaai/dependencies. Upstream downloads wheels from GitHub
release storage at install time; we vendor the wheel bytes inside the git
branch so the install works against a private gitea repo with no extra auth.
"""
import os
import platform
import zipfile
from io import BytesIO

try:
  import tomllib
except ImportError:
  import tomli as tomllib

from setuptools import setup
from setuptools.command.build_py import build_py

_HERE = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(_HERE, "pyproject.toml"), "rb") as _f:
  _cfg = tomllib.load(_f)

DATADIR = _cfg["tool"]["shim"]["datadir"]
VERSION = _cfg["project"]["version"]
MODULE = _cfg["project"]["name"].replace("-", "_")

# all top-level packages bundled in this wheel.
# `packages` may be a flat list or a dict with find.include patterns.
_pkgs = _cfg.get("tool", {}).get("setuptools", {}).get("packages", [f"{MODULE}*"])
if isinstance(_pkgs, list):
  _INCLUDE = _pkgs
elif isinstance(_pkgs, dict):
  _INCLUDE = _pkgs.get("find", {}).get("include", [f"{MODULE}*"])
else:
  _INCLUDE = [f"{MODULE}*"]
TOP_PACKAGES = sorted({p.rstrip("*").rstrip("/") for p in _INCLUDE if p.rstrip("*").rstrip("/")})

PLATFORM_MAP = {
  ("Linux", "x86_64"): "linux_x86_64",
  ("Linux", "aarch64"): "linux_aarch64",
  ("Darwin", "arm64"): "macosx_11_0_arm64",
}


class InstallPrebuilt(build_py):
  def run(self):
    module_dir = os.path.join(_HERE, MODULE)
    data_dir = os.path.join(module_dir, DATADIR)

    if not os.path.exists(os.path.join(data_dir, "lib")) and not os.path.exists(os.path.join(data_dir, "bin")):
      key = (platform.system(), platform.machine())
      plat = PLATFORM_MAP.get(key)
      if plat is None:
        raise RuntimeError(f"unsupported platform: {key}")

      candidates = [
        f"{MODULE}-{VERSION}-py3-none-{plat}.whl",
        f"{MODULE}-{VERSION}-py3-none-any.whl",
      ]
      whl_path = None
      wheels_dir = os.path.join(_HERE, "wheels")
      for name in candidates:
        candidate = os.path.join(wheels_dir, name)
        if os.path.exists(candidate):
          whl_path = candidate
          break
      if whl_path is None:
        raise RuntimeError(
          f"vendored wheel not found for platform {key} under {wheels_dir} "
          f"(looked for: {', '.join(candidates)})"
        )

      print(f"Extracting vendored wheel {whl_path} ...")
      with open(whl_path, "rb") as f:
        raw = f.read()
      with zipfile.ZipFile(BytesIO(raw)) as zf:
        purelib_data = f"{MODULE}-{VERSION}.data/purelib/"
        for info in zf.infolist():
          if info.is_dir():
            continue
          name = info.filename
          if name.startswith(purelib_data):
            rel = name[len(purelib_data):]
          else:
            rel = name
          if "/" not in rel:
            continue
          top = rel.split("/", 1)[0]
          if top not in TOP_PACKAGES:
            continue
          dest = os.path.join(_HERE, rel)
          os.makedirs(os.path.dirname(dest), exist_ok=True)
          with open(dest, "wb") as f:
            f.write(zf.read(info))
          if info.external_attr >> 16 & 0o111:
            os.chmod(dest, 0o755)

    super().run()


setup(cmdclass={"build_py": InstallPrebuilt})