"""VoiceCheck DAM Inference Server

Wraps KintsugiHealth/dam model in a Gradio app with a REST-friendly API.
"""
import io
import json
import tempfile
import os
import gradio as gr
import torch

# ---- Load model at startup ------------------------------------------------
print("Loading KintsugiHealth/dam pipeline...")
from huggingface_hub import snapshot_download

model_dir = snapshot_download("KintsugiHealth/dam")

# Add model dir to path so we can import its modules
import sys
sys.path.insert(0, model_dir)

from pipeline import Pipeline

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
pipe = Pipeline(
    checkpoint=os.path.join(model_dir, "dam3.1.ckpt"),
    device=device,
)
print("Model loaded successfully!")


# ---- Inference function ----------------------------------------------------
def predict(audio_filepath):
    """Run DAM inference on an audio file."""
    if audio_filepath is None:
        return json.dumps({"error": "No audio provided"})

    try:
        result = pipe.run_on_file(audio_filepath, quantize=True)
        return json.dumps({
            "depression": result["depression"],
            "anxiety": result["anxiety"],
        })
    except Exception as e:
        return json.dumps({"error": str(e)})


# ---- Gradio UI ------------------------------------------------------------
with gr.Blocks(title="VoiceCheck DAM Inference") as demo:
    gr.Markdown("## VoiceCheck DAM Inference")
    gr.Markdown("Upload or record audio (30+ seconds recommended).")

    audio_input = gr.Audio(
        label="Upload audio or record (30+ seconds recommended)",
        type="filepath",
        sources=["upload", "microphone"],
    )
    output = gr.Textbox(label="Analysis Result (JSON)", lines=4)
    btn = gr.Button("Analyze", variant="primary")

    btn.click(fn=predict, inputs=audio_input, outputs=output)

demo.launch()