from pathlib import Path
import shutil
import sys
import sysconfig

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension


ROOT = Path("/workspace/MapTRv2/mmdetection3d")


def build_extensions():
    setup(
        name="maptr_bev_pool_ext",
        ext_modules=[
            CUDAExtension(
                "mmdet3d.ops.bev_pool.bev_pool_ext",
                [
                    str(ROOT / "mmdet3d/ops/bev_pool/src/bev_pool.cpp"),
                    str(ROOT / "mmdet3d/ops/bev_pool/src/bev_pool_cuda.cu"),
                ],
            ),
            CUDAExtension(
                "mmdet3d.ops.bev_pool_v2.bev_pool_v2_ext",
                [
                    str(ROOT / "mmdet3d/ops/bev_pool_v2/src/bev_pool.cpp"),
                    str(ROOT / "mmdet3d/ops/bev_pool_v2/src/bev_pool_cuda.cu"),
                ],
            ),
        ],
        cmdclass={"build_ext": BuildExtension},
        script_args=["build_ext", "--inplace"],
    )


def install_extensions():
    installed_root = Path(sysconfig.get_paths()["purelib"]) / "mmdet3d"
    if not installed_root.exists():
        for item in sys.path:
            candidate = Path(item) / "mmdet3d"
            if candidate.exists():
                installed_root = candidate
                break
    installed_ops = installed_root / "ops"
    for name in ("bev_pool", "bev_pool_v2"):
        shutil.copytree(
            ROOT / "mmdet3d" / "ops" / name,
            installed_ops / name,
            dirs_exist_ok=True,
        )

    init_file = installed_ops / "__init__.py"
    init_text = init_file.read_text()
    imports = "from .bev_pool import *\nfrom .bev_pool_v2 import *\n"
    if "from .bev_pool import *" not in init_text:
        marker = "from .paconv import PAConv, PAConvCUDA, assign_score_withk\n"
        if marker in init_text:
            init_text = init_text.replace(marker, marker + imports)
        else:
            init_text += "\n" + imports
        init_file.write_text(init_text)


if __name__ == "__main__":
    build_extensions()
    install_extensions()
