app.py

"""
Health Voice Assistant
- Whisper ASR (incremental)
- GPT health guidance (NON-diagnostic)
- Auto Text-to-Speech (autoplay, no button, no file saved)

⚠️ GENERAL guidance only – NOT medical advice
"""

# ================= IMPORTS =================
import time
import traceback
import re
import textwrap
import base64
from io import BytesIO

import numpy as np
import torch
import requests
import streamlit as st

from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read
from openai import OpenAI
from gtts import gTTS

# ================= PAGE CONFIG =================
st.set_page_config(page_title="Health Voice Assistant", layout="wide")

# ================= SESSION STATE =================
if "transcript" not in st.session_state:
    st.session_state.transcript = None
if "gpt_text" not in st.session_state:
    st.session_state.gpt_text = None
if "audio_html" not in st.session_state:
    st.session_state.audio_html = None
if "done" not in st.session_state:
    st.session_state.done = False

# ================= CONFIG =================
MODEL_NAME = "openai/whisper-small.en"
CHUNK_SEC = 6.0
OVERLAP_SEC = 0.05
MAX_CLIP_SEC = 59.0

CHATGPT_API = st.secrets["gpt_apikey"]
CHATGPT_MODEL = st.secrets.get("gpt_model", "gpt-4.0-mini")

FIREBASE_DB_URL = st.secrets["firebase_database_url"]

device = 0 if torch.cuda.is_available() else -1

# ================= UI HEADER =================
st.title("🩺 Health Voice-Based Assistant")

MEDICAL_DISCLAIMER = """
⚠️ **IMPORTANT MEDICAL DISCLAIMER**

This assistant is **NOT** a medical professional.  
It does **NOT** diagnose diseases or prescribe medication.  
Information provided is for **GENERAL guidance only**.
"""

# ================= SIDEBAR =================
st.sidebar.title("🎓 User Manual")
st.sidebar.markdown(
    """
1. Choose input method  
2. Upload audio **or** record via web  
3. App will automatically:
   - Transcribe (or read transcript)
   - Analyze with GPT
   - Speak the guidance
"""
)
st.sidebar.markdown("---")
st.sidebar.markdown(MEDICAL_DISCLAIMER)

# ================= LOAD ASR =================
@st.cache_resource
def load_pipeline():
    pipe = pipeline(
        task="automatic-speech-recognition",
        model=MODEL_NAME,
        chunk_length_s=CHUNK_SEC,
        device=device,
    )
    try:
        forced_ids = pipe.tokenizer.get_decoder_prompt_ids(
            language="en", task="transcribe"
        )
        pipe.model.config.forced_decoder_ids = forced_ids
        pipe._forced_decoder_ids = forced_ids
    except Exception:
        pass
    return pipe

pipe = load_pipeline()

# ================= GPT CLIENT =================
client = OpenAI(api_key=CHATGPT_API)

# ================= HELPERS =================
def format_text(text: str, width: int = 90) -> str:
    text = re.sub(r"\s+", " ", text).strip()
    sentences = re.split(r"(?<=[.!?])\s+", text)
    paragraphs, buf = [], []
    for s in sentences:
        buf.append(s)
        if len(buf) >= 3:
            paragraphs.append(" ".join(buf))
            buf = []
    if buf:
        paragraphs.append(" ".join(buf))
    return "\n\n".join(textwrap.fill(p, width) for p in paragraphs)

def merge_overlap(old: str, new: str, max_words: int = 10) -> str:
    if not old:
        return new
    a, b = old.split(), new.split()
    for k in range(min(max_words, len(a), len(b)), 0, -1):
        if a[-k:] == b[:k]:
            return " ".join(a + b[k:])
    return old + " " + new

def tts_bytes(text: str) -> BytesIO:
    mp3 = BytesIO()
    gTTS(text=text, lang="en").write_to_fp(mp3)
    mp3.seek(0)
    return mp3

# ================= FIREBASE =================
def fetch_firebase_transcript():
    """
    Read transcript text from Firebase Realtime Database
    Path: /transcribe
    """
    try:
        url = f"{FIREBASE_DB_URL}/transcribe.json"
        res = requests.get(url, timeout=5)
        res.raise_for_status()
        data = res.json()

        if not data:
            return None

        text = data.get("text", "").strip()
        return text if text else None

    except Exception as e:
        st.error(f"Firebase read error: {e}")
        return None

# ================= GPT =================
def run_gpt(transcript: str) -> str:
    prompt = f"""
You are a NON-diagnostic health assistant.

{MEDICAL_DISCLAIMER}

Transcript:
{transcript}

TASK:
- Summarize possible concern categories (NOT diagnosis)
- Explain in simple language
- Give safe general advice
- Clearly state when to seek professional medical help
"""
    res = client.chat.completions.create(
        model=CHATGPT_MODEL,
        messages=[
            {"role": "system", "content": "Health assistant"},
            {"role": "user", "content": prompt},
        ],
        temperature=0.3,
    )
    return res.choices[0].message.content.strip()

# ================= ASR =================
def transcribe(audio_bytes: bytes):
    sr = pipe.feature_extractor.sampling_rate
    audio = ffmpeg_read(audio_bytes, sr)

    if audio.ndim > 1:
        audio = audio.mean(axis=1)

    if audio.shape[0] / sr > MAX_CLIP_SEC:
        raise ValueError("Audio too long (max 59s)")

    chunk_samples = int(CHUNK_SEC * sr)
    overlap_samples = int(OVERLAP_SEC * sr)

    total_chunks = max(1, (len(audio) + chunk_samples - 1) // chunk_samples)
    acc = ""

    for i in range(total_chunks):
        start = max(0, i * chunk_samples - overlap_samples)
        end = min(len(audio), (i + 1) * chunk_samples + overlap_samples)

        result = pipe(
            {"array": audio[start:end], "sampling_rate": sr},
            generate_kwargs={
                "forced_decoder_ids": getattr(pipe, "_forced_decoder_ids", None)
            },
        )

        acc = merge_overlap(acc, result["text"].strip())
        yield acc, i + 1, total_chunks

# ================= INPUT MODE =================
st.markdown("### 🎤 Input Mode")

mode = st.radio(
    "Choose input method:",
    ["Upload audio file", "Live recording (via Firebase)"],
)

# ================= MODE 2: FIREBASE =================
if mode == "Live recording (via Firebase)":
    st.info("Reading transcript from Firebase...")

    firebase_text = fetch_firebase_transcript()

    if firebase_text:
        st.success("Transcript found. Processing with AI...")

        st.session_state.audio_html = None
        st.session_state.transcript = firebase_text

        with st.spinner("Analyzing with GPT..."):
            st.session_state.gpt_text = run_gpt(firebase_text)

        voice = tts_bytes(st.session_state.gpt_text)
        b64 = base64.b64encode(voice.read()).decode()
        st.session_state.audio_html = f"""
        <audio autoplay>
            <source src="data:audio/mp3;base64,{b64}" type="audio/mp3">
        </audio>
        """

        st.session_state.done = True

    else:
        st.warning("No transcript found. Redirecting to recorder...")
        st.markdown(
            """
            <meta http-equiv="refresh" content="0; url=https://nolist-ssp-recordoption.hf.space">
            """,
            unsafe_allow_html=True,
        )

# ================= MODE 1: UPLOAD AUDIO =================
if mode == "Upload audio file":
    uploaded = st.file_uploader(
        "Upload audio (wav, mp3, m4a, ogg, flac)",
        type=["wav", "mp3", "m4a", "ogg", "flac"],
    )

    if uploaded:
        audio_bytes = uploaded.read()
        st.audio(audio_bytes)

        if st.button("Transcribe"):
            try:
                st.session_state.audio_html = None

                placeholder = st.empty()
                progress = st.progress(0.0)

                final_text = ""

                for txt, i, total in transcribe(audio_bytes):
                    final_text = txt
                    placeholder.markdown(
                        f"```text\n{format_text(txt)}\n```"
                    )
                    progress.progress(i / total)

                st.session_state.transcript = final_text

                with st.spinner("Analyzing with GPT..."):
                    st.session_state.gpt_text = run_gpt(final_text)

                voice = tts_bytes(st.session_state.gpt_text)
                b64 = base64.b64encode(voice.read()).decode()
                st.session_state.audio_html = f"""
                <audio autoplay>
                    <source src="data:audio/mp3;base64,{b64}" type="audio/mp3">
                </audio>
                """

                st.session_state.done = True

            except Exception as e:
                st.error(f"Failed:\n{e}\n{traceback.format_exc()}")

# ================= RENDER RESULT =================
if st.session_state.transcript:
    st.markdown("### 📝 Transcript")
    st.markdown(f"```text\n{format_text(st.session_state.transcript)}\n```")

if st.session_state.gpt_text:
    st.markdown("### 🧠 Health Guidance")
    st.markdown(st.session_state.gpt_text)

if st.session_state.audio_html:
    st.markdown(st.session_state.audio_html, unsafe_allow_html=True)

st.markdown("---")
st.markdown("⚠️ **NOT medical advice. Seek professional help when needed.**")