app.py

import streamlit as st
import base64, tempfile, requests, os, re
import whisper  # local Whisper model
from piper import PiperVoice

import wave


# ─────────────────────────────────────────────────────────────
st.set_page_config(page_title="Med11 – AI Medical Emergency guide", page_icon="🩺")

# Groq key (unchanged)
GROQ_API_KEY = st.secrets["GROQ_API_KEY"]

# Session flags
st.session_state.setdefault("show_numbers", False)
st.session_state.setdefault("upload_id", 0)
st.session_state.setdefault("chat_history", [])

# ─────────────────────────────────────────────────────────────
@st.cache_resource(show_spinner="Loading Whisper model…")
def _load_whisper():
    # Loads the *base* model once and re-uses it across reruns
    return whisper.load_model("base")

@st.cache_resource(show_spinner="Loading Piper voice model…")
def load_voice():
    # Adjust the path as needed, model files must be in 'models' directory
    return PiperVoice.load("models/en_US-amy-low.onnx")

def groq_reply(prompt: str, model="llama3-8b-8192"):
    url = "https://api.groq.com/openai/v1/chat/completions"
    payload = {
        "model": model,
        "messages": [
            {
                "role": "system",
                "content": (
                    "You are an emergency-response medic. "
                    "Classify the emergency, rate severity, and give clear, step-by-step first-aid. "
                    "Limit to 120 words."
                ),
            },
            *st.session_state.chat_history,
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.2,
        "max_tokens": 256,
    }
    r = requests.post(
        url, json=payload,
        headers={"Authorization": f"Bearer {GROQ_API_KEY}"},
        timeout=45,
    )
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"]

def transcribe_audio(file_path: str) -> str:
    model = _load_whisper()
    result = model.transcribe(file_path)
    return result["text"]

# ────────────────────────────────
# Gibberish / Markdown sanitization
# ────────────────────────────────
def clean_text_for_tts(text: str) -> str:
    text = re.sub(r"\*\*|\*|__|_", "", text)              # bold/italic
    text = re.sub(r"[|~>#]", "", text)                   # misc. MD
    text = re.sub(r"[\u2600-\u26FF\u2700-\u27BF]+", "", text)  # emojis
    text = re.sub(r"\s+", " ", text)                      # whitespace
    return text.strip()


def tts_player(text: str):
    voice = load_voice()
    clean_text = clean_text_for_tts(text)

    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_wav:
        with wave.open(tmp_wav.name, "wb") as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)  # 16-bit samples
            wav_file.setframerate(voice.config.sample_rate)
            voice.synthesize(clean_text, wav_file)

        audio_bytes = open(tmp_wav.name, "rb").read()

    st.audio(audio_bytes, format="audio/wav")
    os.remove(tmp_wav.name)



# ─────────────────────────────────────────────────────────────
st.title("🩺 Med11 – AI Medical Emergency guide")

if st.button("📞 Emergency Numbers"):
    st.session_state.show_numbers = not st.session_state.show_numbers

if st.session_state.show_numbers:
    with st.expander("Key Emergency Numbers (Pakistan)", expanded=True):
        st.markdown(
            """
            | Service | Number |
            |---------|--------|
            | 🚑 **Edhi Ambulance** | **115** |
            | 🚒 **Rescue / Fire** | **1122** |
            | 🚓 Police | **15** |
            | 💊 Poison Control (NIH) | **051-9255075** |
            """
        )

# Upload audio
uploader_key = f"audio_{st.session_state.upload_id}"
audio_file = st.file_uploader(
    "Upload a WAV/MP3 describing the emergency",
    type=["wav", "mp3", "m4a"],
    key=uploader_key,
)

if audio_file:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
        tmp.write(audio_file.read())
        tmp_path = tmp.name

    with st.spinner("Transcribing…"):
        try:
            transcript = transcribe_audio(tmp_path).strip()
        except Exception as e:
            st.error(f"Transcription error: {e}")
            transcript = ""
    os.remove(tmp_path)

    with st.expander("📜 Transcript", expanded=False):
        st.write(transcript or "_(no speech detected)_")

    if transcript:
        with st.spinner("Analyzing…"):
            try:
                result = groq_reply(
                    f'Analyze this emergency: "{transcript}". '
                    "Classify type, rate severity, and give brief first-aid steps."
                )
                st.session_state.chat_history.append({"role": "user", "content": transcript})
                st.session_state.chat_history.append({"role": "assistant", "content": result})
            except Exception as e:
                st.error(f"Groq error: {e}")
                result = ""

        if result:
            tts_player(result)
            st.markdown("### 🧠 Guidance")
            st.markdown(result)

            st.session_state.upload_id += 1
            st.markdown("---")
            st.markdown("**Need to tell me more? Upload another voice clip below.**")

# ─────────────────────────────────────────────────────────────
col1, col2 = st.columns(2)
with col1:
    if st.button("🏥 Nearest Hospital (Google Maps)"):
        st.markdown(
            "[Open Maps →](https://www.google.com/maps/search/hospital/)",
            unsafe_allow_html=True,
        )
with col2:
    if st.button("🗺️ Nearby Pharmacies"):
        st.markdown(
            "[Open Maps →](https://www.google.com/maps/search/pharmacy/)",
            unsafe_allow_html=True,
        )