#!/usr/bin/env python3
"""SplatAnything Gradio demo: python app.py (local GPU or CPU smoke run)."""
import argparse
import os
from pathlib import Path

ROOT = Path(__file__).resolve().parent


def parse_args(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--host', default=os.environ.get('HOST', '127.0.0.1'))
    parser.add_argument('--port', type=int, default=int(os.environ.get('PORT', '7860')))
    parser.add_argument('--device', default=os.environ.get('SPLATANYTHING_DEVICE', 'cuda'))
    parser.add_argument('--config_path', default=os.environ.get('SPLATANYTHING_CONFIG', str(ROOT / 'configs/inference.yaml')))
    parser.add_argument('--ckpt_path', default=os.environ.get('SPLATANYTHING_CKPT'))
    parser.add_argument('--geometry_ckpt_path', default=os.environ.get('WORLDMIRROR2_CKPT'))
    parser.add_argument('--sky_onnx_path', default=os.environ.get('SPLATANYTHING_SKY_CKPT'))
    parser.add_argument('--output_dir', default=os.environ.get('SPLATANYTHING_OUTPUT_DIR'))
    parser.add_argument('--examples_dir', default=os.environ.get('SPLATANYTHING_EXAMPLES_DIR', str(ROOT / 'examples')))
    parser.add_argument('--install-viewer', action='store_true', help='Install the pinned npm viewer dependencies before launch.')
    args = parser.parse_args(argv)
    return args


def build_app(args):
    # ZeroGPU must initialize CUDA emulation before Gradio or inference imports torch.
    zero = os.environ.get('SPLATANYTHING_ZEROGPU', os.environ.get('SPACES_ZERO_GPU', '0')).lower() in ('1', 'true')
    if zero:
        import spaces  # noqa: F401
    from demo.runtime import Pipeline
    from demo.ui import create_app
    pipeline = Pipeline(args.config_path, args.device, args.ckpt_path, args.geometry_ckpt_path, args.sky_onnx_path)
    predict_fn = None
    if zero:
        from demo.zero_gpu import initialize, predict
        initialize(pipeline)
        predict_fn = predict
    app = create_app(pipeline, args.output_dir, args.examples_dir, predict_fn=predict_fn)
    app.state.zero_gpu = zero
    return app


def main(argv=None):
    args = parse_args(argv)
    if args.install_viewer:
        from demo.viewer import install_converter
        install_converter()
    app = build_app(args)
    import uvicorn
    if app.state.zero_gpu:
        # mount_gradio_app does not call Blocks.launch(), whose spaces SDK hook
        # packs preloaded CUDA tensors and reports the GPU functions to the host.
        # Expose the server as well so ZeroGPU workers close inherited listener FDs.
        server = uvicorn.Server(uvicorn.Config(app, host=args.host, port=args.port, access_log=False))
        app.state.demo.server = server
        from demo.zero_gpu import startup
        startup()
        server.run()
    else:
        uvicorn.run(app, host=args.host, port=args.port, access_log=False)


if __name__ == '__main__':
    main()