from __future__ import annotations

import os
import shutil
import subprocess
from pathlib import Path

from setuptools import setup
from setuptools.command.build_py import build_py as _build_py


def _binary_name() -> str:
    return "sglangmuxd.exe" if os.name == "nt" else "sglangmuxd"


class build_py(_build_py):
    def run(self) -> None:
        root = Path(__file__).parent.resolve()
        binary = _binary_name()

        cargo_cmd = ["cargo", "build", "--release", "--bin", "sglangmuxd"]
        subprocess.run(cargo_cmd, cwd=root, check=True)

        cargo_target = os.environ.get("CARGO_BUILD_TARGET")
        if cargo_target:
            built_binary = root / "target" / cargo_target / "release" / binary
        else:
            built_binary = root / "target" / "release" / binary

        if not built_binary.exists():
            raise RuntimeError(f"Expected Rust binary not found: {built_binary}")

        out_dir = Path(self.build_lib) / "sglangmux" / "bin"
        out_dir.mkdir(parents=True, exist_ok=True)
        out_binary = out_dir / binary
        shutil.copy2(built_binary, out_binary)
        out_binary.chmod(0o755)

        super().run()


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