backend.py

import os
import time
import tempfile
import traceback
from pathlib import Path

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

from interface import (
    save_patient_registry,
    load_patient_registry,
    run_end_to_end,
    monitoring_table_rows,
    search_reasoning_records,
    load_history_records,
    RESULTS_DIR,
)
from live_analyze_route import router as live_analyze_router

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# NOTE: an eager @app.on_event("startup") model preload was tried here and
# reverted — loading SepFormer + Wav2Vec2 + the GNN checkpoint all at once
# during boot risks exceeding this Space's memory tier, which can crash
# the whole container (symptom: even /health, which touches no models,
# starts failing too). Lazy-loading on first use is slower for that first
# request but far safer. If cold-start latency needs solving, do it by
# upgrading the Space's hardware tier, or by loading models one at a time
# with memory checks in between — not all at once on every boot.

UPLOAD_DIR = RESULTS_DIR / "uploads"
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)


def save_upload(upload: UploadFile, subfolder: str = "general") -> str:
    folder = UPLOAD_DIR / subfolder
    folder.mkdir(parents=True, exist_ok=True)
    dest = folder / f"{int(time.time())}_{upload.filename}"
    with open(dest, "wb") as f:
        f.write(upload.file.read())
    return str(dest)


@app.get("/health")
async def health():
    return {"status": "ok"}


@app.post("/api/patients/")
async def register_patient(
    patient_code: str = Form(...),
    full_name: str = Form(default=""),
    age: int = Form(...),
    gender: str = Form(...),
    ward: str = Form(...),
    bed_id: str = Form(...),
    baseline_audio: UploadFile = File(...),
):
    try:
        audio_path = save_upload(baseline_audio, "baselines")
        display_name = full_name.strip() if full_name.strip() else patient_code
        save_patient_registry([{
            "patient_id": patient_code,
            "name": display_name,
            "reference_audio": audio_path,
            "age": age,
            "gender": gender,
            "ward": ward,
            "bed_id": bed_id,
        }])
        return {
            "patient_id": patient_code,
            "full_name": display_name,
            "status": "registered",
            "bed_id": bed_id,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/patients/")
async def get_patients():
    try:
        registry = load_patient_registry()
        # Make sure bed_id is included in the response
        patients = []
        for p in registry:
            patients.append({
                "patient_id": p.get("patient_id", ""),
                "name": p.get("name", ""),
                "age": p.get("age"),
                "gender": p.get("gender", ""),
                "ward": p.get("ward", ""),
                "bed_id": p.get("bed_id", ""),
                "reference_audio": p.get("reference_audio"),
            })
        return {"patients": patients}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/analyze/")
async def analyze_audio(
    bed_id: str = Form(...),
    file: UploadFile = File(...),
):
    try:
        audio_path = save_upload(file, "mic_uploads")
        from interface import load_audio, predict_sources
        waveform, sr = load_audio(audio_path)
        prediction, _ = predict_sources(waveform, sr)
        source = prediction[0].numpy()
        duration = len(source) / 16000
        from gnn import estimate_breathing_rate_bpm
        bpm = estimate_breathing_rate_bpm(source, 16000, duration)
        if bpm < 12 or bpm > 25:
            severity = 0.75
            condition = "abnormal"
        elif bpm < 16 or bpm > 20:
            severity = 0.45
            condition = "mild"
        else:
            severity = 0.15
            condition = "normal"
        return {
            "bed_id": bed_id,
            "condition": condition,
            "severity_score": severity,
            "respiratory_rate": round(bpm, 1),
            "confidence": 87,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ── Live monitoring: one request per mic chunk, covering all beds that
#    mic serves (see live_analyze_route.py / interface.analyze_live_chunk_for_beds).
#    Registered here, before the catch-all static mount below.
app.include_router(live_analyze_router)


@app.post("/api/pipeline/")
async def run_pipeline(
    mix_audio: UploadFile = File(...),
    patient_1: str = Form(default=""),
    patient_2: str = Form(default=""),
    patient_3: str = Form(default=""),
):
    try:
        audio_path = save_upload(mix_audio, "mix_uploads")
        patient_names = [patient_1 or None, patient_2 or None, patient_3 or None]
        registry = load_patient_registry()
        reference_audio_paths = [None, None, None]
        for i, name in enumerate(patient_names):
            if name:
                for entry in registry:
                    if entry.get("name") == name:
                        reference_audio_paths[i] = (
                            entry.get("local_reference_audio")
                            or entry.get("reference_audio")
                        )
                        break
        outputs, reasoning_summaries, history_record = run_end_to_end(
            audio_path,
            patient_names=patient_names,
            reference_audio_paths=reference_audio_paths,
        )
        return {
            "status": "complete",
            "separated_sources": outputs[:3],
            "reasoning_count": len(reasoning_summaries),
            "timestamp": history_record.get("timestamp"),
            "summaries": reasoning_summaries,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=traceback.format_exc())


@app.get("/api/monitoring/")
async def get_monitoring():
    try:
        rows = monitoring_table_rows()
        return {"rows": rows}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/history/")
async def get_history():
    try:
        records = load_history_records()
        return {"records": records}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/history/search/")
async def search_history(query: str = ""):
    try:
        results = search_reasoning_records(query)
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ── Serve React (MUST be last) ────────────────────────────────
app.mount("/", StaticFiles(directory="dist", html=True), name="static")