"""Build the pinned native SPAC extensions on the Space CPU at startup."""
from __future__ import annotations
import hashlib
import importlib.util
import os
from pathlib import Path
import shutil
import site
import subprocess
import sys
import tarfile
import requests

ROOT=Path(__file__).resolve().parent
CACHE=ROOT/'.cache/zerogpu-native'
CUDA=CACHE/'cuda-12.8'
REDIST='https://developer.download.nvidia.com/compute/cuda/redist/'
ARCHIVES=[
 ('cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64-12.8.93-archive.tar.xz','9961b3484b6b71314063709a4f9529654f96782ad39e72bf1e00f070db8210d3'),
 ('cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-12.8.90-archive.tar.xz','8d566b5fe745c46842dc16945cf36686227536decd2302c372be86da37faca68'),
 ('cuda_cccl/linux-x86_64/cuda_cccl-linux-x86_64-12.8.90-archive.tar.xz','0740e9e01e4f15e17c5ab8d68bba4f8ec0eb6b84edccba4ac45112d2d2174e4b'),
]


def download(url,path,sha=None):
    path.parent.mkdir(parents=True,exist_ok=True)
    if not path.exists():
        session=requests.Session();session.trust_env=False
        with session.get(url,stream=True,timeout=(15,120)) as response:
            response.raise_for_status()
            temporary=path.with_suffix(path.suffix+'.partial')
            with temporary.open('wb') as stream:
                for chunk in response.iter_content(4*1024*1024):stream.write(chunk)
            temporary.replace(path)
    if sha and hashlib.sha256(path.read_bytes()).hexdigest()!=sha:
        raise RuntimeError('CUDA toolkit archive checksum mismatch')


def extract(path,folder):
    folder.mkdir(parents=True,exist_ok=True)
    with tarfile.open(path) as archive:
        for member in archive.getmembers():
            if not (folder/member.name).resolve().is_relative_to(folder.resolve()):
                raise RuntimeError('Unsafe archive member')
        archive.extractall(folder)


def ensure_native():
    for key in list(os.environ):
        if 'proxy' in key.lower():os.environ.pop(key)
    CUDA.mkdir(parents=True,exist_ok=True)
    if not (CUDA/'bin/nvcc').exists():
        for relative,sha in ARCHIVES:
            path=CACHE/Path(relative).name
            print('CUDA_TOOLCHAIN_COMPONENT',path.name,flush=True)
            download(REDIST+relative,path,sha)
            folder=CACHE/path.name.replace('.tar.xz','')
            extract(path,folder)
            source=next(folder.iterdir())
            shutil.copytree(source,CUDA,dirs_exist_ok=True,symlinks=True)
    os.environ['CUDA_HOME']=str(CUDA)
    os.environ['PATH']=str(CUDA/'bin')+os.pathsep+os.environ.get('PATH','')
    os.environ['TORCH_CUDA_ARCH_LIST']='12.0+PTX'
    os.environ['FORCE_CUDA']='1'
    os.environ['MAX_JOBS']='2'
    includes=[str(CUDA/'include')];libraries=[str(CUDA/'lib'),str(CUDA/'lib64')]
    for base in site.getsitepackages():
        nvidia=Path(base)/'nvidia'
        includes.extend(str(p) for p in nvidia.glob('*/include'))
        libraries.extend(str(p) for p in nvidia.glob('*/lib'))
    for key,values in [('CPATH',includes),('CPLUS_INCLUDE_PATH',includes),('LIBRARY_PATH',libraries),('LD_LIBRARY_PATH',libraries)]:
        os.environ[key]=os.pathsep.join(values+[os.environ.get(key,'')])
    wheels=CACHE/'wheels'
    wheels.mkdir(parents=True,exist_ok=True)
    if importlib.util.find_spec('torchsparse') is None:
        revision='385f5ce8718fcae93540511b7f5832f4e71fd835'
        archive=CACHE/'torchsparse.tar.gz'
        download('https://codeload.github.com/mit-han-lab/torchsparse/tar.gz/'+revision,archive)
        folder=CACHE/'torchsparse-source';extract(archive,folder)
        source=next(folder.iterdir())
        print('BUILDING_TORCHSPARSE_FOR_BLACKWELL',flush=True)
        setup=source/'setup.py'
        setup.write_text(setup.read_text().replace('["-g", "-O3"','["-g0", "-O3"'))
        subprocess.run([sys.executable,'-m','pip','wheel','--no-deps','--no-build-isolation','-w',str(wheels),str(source)],check=True)
        wheel=next(wheels.glob('torchsparse-*.whl'))
        subprocess.run([sys.executable,'-m','pip','install','--no-deps',str(wheel)],check=True)
    if importlib.util.find_spec('spac_sdf_cuda') is None:
        print('BUILDING_SDF_KERNEL_FOR_BLACKWELL',flush=True)
        setup=CACHE/'setup_sdf.py'
        sources=[str(ROOT/'spac/_kernels/src/udf_kernel.cu'),str(ROOT/'spac/_kernels/src/udf_cuda.cpp')]
        setup.write_text('from setuptools import setup\nfrom torch.utils.cpp_extension import BuildExtension,CUDAExtension\nsetup(name="spac-sdf-cuda",version="0.1.0",ext_modules=[CUDAExtension("spac_sdf_cuda",'+repr(sources)+',extra_compile_args={"cxx":["-g0"],"nvcc":["-O3"]})],cmdclass={"build_ext":BuildExtension})\n')
        subprocess.run([sys.executable,str(setup),'bdist_wheel','--dist-dir',str(wheels)],cwd=ROOT,check=True)
        wheel=next(wheels.glob('spac_sdf_cuda-*.whl'))
        subprocess.run([sys.executable,'-m','pip','install','--no-deps',str(wheel)],check=True)
    print('NATIVE_EXTENSIONS_READY',flush=True)