"""Temporary ZeroGPU isolation entrypoint; select a stage with DVLT_DIAGNOSTIC_MODE."""
import os
os.environ.setdefault('OMP_NUM_THREADS', '1')
os.environ.setdefault('MKL_NUM_THREADS', '1')
os.environ.setdefault('OPENBLAS_NUM_THREADS', '1')
import faulthandler
import sys
from pathlib import Path
import platform
import time

faulthandler.enable()
def child_trace():
    print('DIAGNOSTIC: child forked', flush=True)
    faulthandler.dump_traceback_later(30, repeat=True)
os.register_at_fork(after_in_child=child_trace)

import spaces
import torch
import gradio as gr

torch.set_num_threads(1)
sys.path.insert(0, str(Path(__file__).resolve().parent / 'src'))
mode = os.getenv('DVLT_DIAGNOSTIC_MODE', 'minimal')
print(f'DIAGNOSTIC: mode={mode}, python={platform.python_version()}, torch={torch.__version__}, cuda={torch.version.cuda}, gradio={gr.__version__}', flush=True)
if mode in ('open3d', 'onnxruntime', 'rerun', 'cv2', 'accelerate', 'torchvision'):
    import importlib
    importlib.import_module(mode)
    print(f'DIAGNOSTIC: imported {mode}', flush=True)
if mode in ('imports', 'model'):
    import dvlt.scripts.gradio_demo
    print('DIAGNOSTIC: DVLT imports complete', flush=True)
if mode == 'model':
    from accelerate import Accelerator
    from dvlt.util.model_registry import ModelEntry
    accelerator = Accelerator(mixed_precision='bf16')
    entry = ModelEntry(label='DVLT', config_name='dvlt')
    entry.ensure_loaded(accelerator)
    print('DIAGNOSTIC: DVLT model loaded', flush=True)

@spaces.GPU(duration=30)
def probe():
    print('DIAGNOSTIC: GPU function entered', flush=True)
    start = time.monotonic()
    x = torch.ones((32, 32), device='cuda')
    result = float((x @ x).sum().item())
    torch.cuda.synchronize()
    print('DIAGNOSTIC: CUDA calculation complete', flush=True)
    return {'mode': mode, 'result': result, 'expected': 32768.0, 'gpu': torch.cuda.get_device_name(0), 'seconds': time.monotonic()-start}

with gr.Blocks() as demo:
    gr.Markdown(f'# ZeroGPU diagnostic — {mode}')
    run = gr.Button('Run GPU probe')
    out = gr.JSON()
    run.click(probe, outputs=out, api_name='probe')
if __name__ == '__main__':
    demo.queue().launch(ssr_mode=False)