"""
Bharat Voice Assistant on Modal — Gradio UI on A10G (~11GB VRAM)
================================================================
MiniCPM-o 4.5 (STT + LLM) + VoxCPM2 (TTS) in one container.

  modal deploy deploy_modal.py
  → https://aniketchopde03--bharat-voice-ui.modal.run

Dashboard: https://modal.com/apps/aniketchopde03/main
"""

import modal

app = modal.App("bharat-voice")

OMNI_REVISION = "4382fcae8a551b54d18f18462db974ff312aa7f3"
BAKE_MODELS = (
    'python -c "from huggingface_hub import snapshot_download; '
    f"snapshot_download('openbmb/MiniCPM-o-4_5', revision='{OMNI_REVISION}', "
    "local_dir='/models/MiniCPM-o-4_5'); "
    "snapshot_download('openbmb/VoxCPM2', local_dir='/models/VoxCPM2')\""
)

web_image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("ffmpeg", "libsndfile1", "build-essential")
    .pip_install("torch==2.8.0", "torchaudio==2.8.0")
    .pip_install(
        "transformers==4.51.0",
        "huggingface_hub>=0.33.5,<1.0",
        "gradio>=6.0.0,<6.18.0",
        "accelerate>=0.34.0",
        "bitsandbytes>=0.43.0",
        "librosa==0.9.0",  # minicpmo-utils pins this exactly
        "setuptools>=65.0.0,<81",  # librosa 0.9.0 needs pkg_resources
        "soundfile>=0.12.1",
        "numpy>=1.26.0",
        "minicpmo-utils[all]>=1.0.5",
        "voxcpm>=2.0.3",
        "sentencepiece>=0.2.0",
        "fastapi[standard]",
    )
    # Bake to fixed paths — more reliable than HF cache + HF_HUB_OFFLINE
    .run_commands(BAKE_MODELS)
    .env({
        "OMNI_MODEL_PATH": "/models/MiniCPM-o-4_5",
        "TTS_MODEL_PATH": "/models/VoxCPM2",
        "HF_HUB_DISABLE_TELEMETRY": "1",
    })
    .add_local_file("app.py", "/root/bharat_voice_app.py")
)

with web_image.imports():
    from fastapi import FastAPI
    from gradio.routes import mount_gradio_app


@app.function(
    image=web_image,
    gpu="A10G",
    memory=16384,           # 16GB — faster A10G scheduling; model ~8GB 4-bit
    max_containers=1,       # Gradio sticky sessions
    scaledown_window=600,   # idle 10 min → scale to zero
    timeout=1800,           # cold GPU queue + model load can exceed 15 min
)
@modal.concurrent(max_inputs=50)  # REQUIRED: Gradio heartbeat + dozens of asset requests
@modal.asgi_app()
def ui():
    import importlib.util
    import threading

    spec = importlib.util.spec_from_file_location(
        "bharat_voice_app", "/root/bharat_voice_app.py"
    )
    bharat = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(bharat)

    # Mount UI first (queue/heartbeat need ASGI up); warm GPU model in background
    web_app = mount_gradio_app(
        app=FastAPI(),
        blocks=bharat.demo,
        path="/",
        css=bharat._TTS_CSS,
        js=bharat._TTS_JS,
        theme="soft",
        show_error=True,
    )

    def _warmup():
        print("[INFO] Loading MiniCPM-o on GPU (background warmup)...")
        bharat.get_omni()
        print("[INFO] MiniCPM-o ready.")
        # VoxCPM2 loads on first TTS (warmup is slow ~2min on A10G)

    threading.Thread(target=_warmup, daemon=True).start()
    return web_app