import spaces
import torch
import gradio as gr
from native_bootstrap import ensure_native,CACHE

ensure_native()

@spaces.GPU(duration=20)
def probe_gpu():
    import tempfile
    from pathlib import Path
    import trimesh
    import torchsparse
    from spac.sparse import SparseTensor,SparseConv3d
    from spac.geometry import mesh_to_sdf
    torchsparse.backends.allow_tf32=False
    torch.backends.cuda.matmul.allow_tf32=False
    torch.backends.cudnn.allow_tf32=False
    torch.manual_seed(42)
    xyz=torch.cartesian_prod(*(torch.arange(3,device='cuda') for _ in range(3))).int()
    coords=torch.cat([torch.zeros((len(xyz),1),device='cuda',dtype=torch.int32),xyz],dim=1)
    feats=torch.ones((len(xyz),4),device='cuda',dtype=torch.float32)
    conv=SparseConv3d(4,4,3).to('cuda').eval()
    with torch.no_grad():output=conv(SparseTensor(feats,coords))
    assert output.feats.dtype==torch.float32 and torch.isfinite(output.feats).all()
    with tempfile.TemporaryDirectory() as directory:
        mesh=Path(directory)/'cube.obj';trimesh.creation.box().export(mesh)
        sdf=mesh_to_sdf(str(mesh),torch.device('cuda'),resolution=16,band=.5)
        assert len(sdf['coords'])>0 and torch.isfinite(sdf['feats']).all()
    return {'gpu':torch.cuda.get_device_name(),'torch':torch.__version__,'cuda':torch.version.cuda,
            'cxx11_abi':torch._C._GLIBCXX_USE_CXX11_ABI,'sparse_conv_fp32':True,'sdf_kernel':True,
            'external_inference_backend':False}

def native_artifacts():
    return [str(p) for p in sorted((CACHE/'wheels').glob('*.whl'))]

with gr.Blocks(title='SPAC-4B ZeroGPU') as demo:
    gr.Markdown('# SPAC-4B ZeroGPU native compatibility check\nTemporary CUDA-extension test. Full model inference is not enabled yet.')
    button=gr.Button('Run native GPU compatibility check')
    result=gr.JSON(label='Hugging Face runtime')
    button.click(probe_gpu,outputs=result,api_name='probe_gpu')
    artifacts=gr.File(label='Native extension artifacts',file_count='multiple',visible=False)
    gr.Button(visible=False).click(native_artifacts,outputs=artifacts,api_name='native_artifacts',queue=False)
demo.queue().launch(allowed_paths=[str(CACHE/'wheels')])