app.py

# app_bilingual_parts.py
"""
Streamlit First Aid Assistant — Bilingual (EN/VN)
Single-file app divided into 3 labeled parts below.
Drop this into your Hugging Face Space as `app.py` or merge into your existing file.

PART 1: imports, session flags, i18n, small helpers
PART 2: Firebase/data rendering and simple config checks
PART 3: OpenAI integration, forms, saving victim info, chat UI
"""

# ------------------
# PART 1: Imports, session flags, translations, helpers
# ------------------
import os
import json
import traceback
from datetime import datetime
from pathlib import Path

import streamlit as st

# Optional imports that may fail at runtime; we handle missing packages gracefully
try:
    import pyrebase
except Exception:
    pyrebase = None

# Session defaults
if "firebase_record_id" not in st.session_state:
    st.session_state.firebase_record_id = None
if "victim_saved" not in st.session_state:
    st.session_state.victim_saved = False
if "basic_done" not in st.session_state:
    st.session_state.basic_done = False

# ------------------
# Simple i18n (EN / VN)
# ------------------
LANGS = {"en": "English", "vn": "Tiếng Việt"}
if "lang" not in st.session_state:
    st.session_state.lang = os.getenv("DEFAULT_LANG", "en")

# Short translations table. Add keys as you expand the app.
TRANSLATIONS = {
    "sidebar_title": {"en": "🎓 Quick Guide", "vn": "🎓 Hướng dẫn nhanh"},
    "capture_line": {
        "en": "Capture: Turn on edge kit → Capture. Take clear close-up (snake photo only if safe).",
        "vn": "Chụp: Bật edge kit → Chụp. Chụp cận (ảnh rắn chỉ khi an toàn).",
    },
    "first_aid_line": {
        "en": "First aid: Start top 3 actions immediately (call, immobilize, stay calm).",
        "vn": "Sơ cứu: Bắt đầu 3 hành động chính ngay (gọi, cố định, giữ bình tĩnh).",
    },
    "symptoms_line": {
        "en": "Symptoms: Optional — tick checklist, add age/weight/time, then submit.",
        "vn": "Triệu chứng: Tùy chọn — chọn danh sách, thêm tuổi/cân nặng/thời gian rồi gửi.",
    },
    "monitor_line": {
        "en": "Monitor: Watch status page, answer prompts, follow guidance.",
        "vn": "Giám sát: Xem trang trạng thái, trả lời câu hỏi, theo hướng dẫn.",
    },
    "escalation_line": {
        "en": "Escalation: High-risk cases notify hospitals automatically.",
        "vn": "Tăng cường: Các ca nguy cơ cao sẽ tự động thông báo bệnh viện.",
    },
    "emergency_call": {
        "en": "**Emergency:** If the victim collapses or cannot breathe, **CALL EMERGENCY SERVICES NOW**.",
        "vn": "**Khẩn cấp:** Nếu nạn nhân ngã hoặc không thở, **GỌI CẤP CỨU NGAY**.",
    },
    # UI labels used later
    "save_info": {"en": "💾 Save info & get advanced guidance", "vn": "💾 Lưu & nhận hướng dẫn nâng cao"},
    "emergency_btn": {"en": "🚨 Emergency — Call local services", "vn": "🚨 Khẩn cấp — Gọi dịch vụ địa phương"},
    "ask_btn": {"en": "Ask", "vn": "Hỏi"},
    "assistant_guidance": {"en": "Assistant guidance:", "vn": "Hướng dẫn trợ lý:"},
    "victim_name": {"en": "Victim name (optional)", "vn": "Tên nạn nhân (không bắt buộc)"},
    "phone": {"en": "Phone number", "vn": "Số điện thoại"},
    "age": {"en": "Age (years)", "vn": "Tuổi (năm)"},
    "no_firebase": {"en": "Firebase configuration not found in secrets.toml", "vn": "Không tìm thấy cấu hình Firebase trong secrets.toml"},
    "no_data": {"en": "No valid data found in Firebase.", "vn": "Không tìm thấy dữ liệu hợp lệ trong Firebase."},
    "weight": {"en": "Weight (kg)", "vn": "Cân Nặng (kg)"},
    "height": {"en": "Height (cm)", "vn": "Chiều Cao (cm)"},
    "weight_help": {"en": "Enter a number or leave blank", "vn": "Nhập số hoặc để trống"},
    "height_help": {"en": "Enter a number or leave blank", "vn": "Nhập số hoặc để trống"},
}


def t(key: str, **fmt):
    """Translate key for current language; fallback to English or the key itself."""
    lang = st.session_state.get("lang", "en")
    entry = TRANSLATIONS.get(key, {})
    text = entry.get(lang) or entry.get("en") or key
    if fmt:
        try:
            return text.format(**fmt)
        except Exception:
            return text
    return text

# Sidebar language chooser + compact guide
with st.sidebar:
    st.selectbox(
        "Language / Ngôn ngữ",
        options=list(LANGS.keys()),
        format_func=lambda k: LANGS[k],
        key="lang",
        help="Choose UI language"
    )

st.sidebar.title(t("sidebar_title"))
st.sidebar.markdown(
    f"## Quick steps\n"
    f"- {t('capture_line')}\n"
    f"- {t('first_aid_line')}\n"
    f"- {t('symptoms_line')}\n"
    f"- {t('monitor_line')}\n"
    f"- {t('escalation_line')}\n\n"
    f"---\n{t('emergency_call')}"
)

# ===============================
# ⭐ Evaluation button (NEW)
# ===============================
st.sidebar.markdown("---")

eval_label = (
    "⭐ Give Feedback"
    if st.session_state.lang == "en"
    else "⭐ Gửi phản hồi"
)

st.sidebar.link_button(
    eval_label,
    "https://docs.google.com/forms/d/e/1FAIpQLSfGD5MlrP9VMCXh_QxVtss68is4k_aRBjegjQsLc1NVEwckqg/viewform?usp=header",
    use_container_width=True,
)


# ------------------
# PART 2: Paths, simple first-aid content (localized), Firebase setup, read/render functions
# ------------------
APP_DIR = Path.cwd()
DATA_DIR = APP_DIR / "data"
DATA_DIR.mkdir(parents=True, exist_ok=True)
DATA_FILE = DATA_DIR / "victims.json"  # local fallback / cache

# Localized step lists
SNAKEBITE_STEPS = {
    "en": [
        "Keep the victim calm and still — limit movement.",
        "Call emergency services immediately.",
        "Position the affected limb at or slightly below heart level.",
        "Remove tight clothing or jewelry near the bite area.",
        "Clean the wound gently with soap and water; do NOT cut or suck the bite.",
        "Cover with a clean, loose dressing. Monitor breathing and consciousness.",
    ],
    "vn": [
        "Giữ nạn nhân bình tĩnh và cố định — hạn chế di chuyển.",
        "Gọi cấp cứu ngay lập tức.",
        "Đặt chi bị thương ngang hoặc thấp hơn tim một chút.",
        "Tháo quần áo chật hoặc trang sức gần vết cắn.",
        "Rửa nhẹ vết thương bằng xà phòng và nước; KHÔNG cắt hoặc hút vết cắn.",
        "Băng lỏng vết thương sạch. Giám sát hô hấp và ý thức.",
    ],
}

NORMAL_WOUND_STEPS = {
    "en": [
        "Wash your hands if possible.",
        "Apply gentle pressure with a clean cloth to stop bleeding.",
        "Rinse the wound with clean water (do not use strong antiseptics directly).",
        "Apply an antiseptic wipe or ointment if available.",
        "Cover with a sterile dressing or bandage.",
        "If bleeding continues or wound is deep, seek medical help.",
    ],
    "vn": [
        "Rửa tay nếu có thể.",
        "Dùng khăn sạch ép nhẹ để cầm máu.",
        "Rửa vết thương bằng nước sạch (không dùng chất tẩy mạnh trực tiếp).",
        "Dùng khăn sát trùng hoặc thuốc mỡ nếu có.",
        "Băng vết thương bằng băng sạch.",
        "Nếu chảy máu không dứt hoặc vết thương sâu, đi khám y tế.",
    ],
}


def get_steps_localized(wound_type: str):
    lang = st.session_state.get("lang", "en")
    wound_type_norm = (wound_type or "").strip().lower()
    if wound_type_norm in ("snakebite", "poisoned"):
        return SNAKEBITE_STEPS.get(lang, SNAKEBITE_STEPS["en"])
    return NORMAL_WOUND_STEPS.get(lang, NORMAL_WOUND_STEPS["en"])

# Firebase init (uses st.secrets if present)
firebase = None
db = None
if pyrebase and "firebase" in st.secrets:
    try:
        fb_cfg = st.secrets["firebase"]
        firebase_config = {
            "apiKey": fb_cfg.get("apiKey"),
            "authDomain": fb_cfg.get("authDomain"),
            "projectId": fb_cfg.get("projectId"),
            "storageBucket": fb_cfg.get("storageBucket"),
            "messagingSenderId": fb_cfg.get("messagingSenderId"),
            "appId": fb_cfg.get("appId"),
            "measurementId": fb_cfg.get("measurementId"),
            "databaseURL": fb_cfg.get("databaseURL"),
        }
        firebase = pyrebase.initialize_app(firebase_config)
        db = firebase.database()
    except Exception as e:
        st.warning(f"Firebase init error: {e}")
else:
    if not pyrebase:
        st.info("pyrebase not installed; Firebase disabled.")
    else:
        st.info(t("no_firebase"))


# Read latest wound data from Firebase (safe fallback to local file)
def get_latest_wound_data():
    # Try Firebase first
    if db:
        try:
            data = db.child("data").get().val()
            if data and isinstance(data, dict):
                latest_key = list(data.keys())[-1]
                latest_data = data[latest_key]
                if isinstance(latest_data, dict):
                    return latest_key, latest_data
        except Exception as e:
            st.warning(f"Firebase read error: {e}")
    # Fallback: read local cache file (if exists)
    if DATA_FILE.exists():
        try:
            all_v = json.loads(DATA_FILE.read_text(encoding="utf-8") or "{}")
            if isinstance(all_v, dict):
                # return last inserted
                k = list(all_v.keys())[-1] if all_v else None
                return (k, all_v[k]) if k else (None, None)
        except Exception as e:
            st.warning(f"Local cache read error: {e}")
    return None, None


def render_latest_firebase_data(show_toast=False):
    if show_toast:
        try:
            st.toast("Refreshing data...", icon="🔁")
        except Exception:
            pass
    record_id, data = get_latest_wound_data()
    if not record_id or not data:
        st.warning(t("no_data"))
        return
    st.session_state.firebase_record_id = record_id
    st.session_state.victim_saved = False

    st.subheader("📡 Latest record")
    st.json(data)
    st.markdown("---")

    wound_type = str(data.get("Wound", data.get("wound", ""))).strip().lower()
    st.session_state.wound_type = wound_type
    st.session_state.basic_done = True

    steps = get_steps_localized(wound_type)
    if wound_type in ("snakebite", "poisoned"):
        st.subheader("🐍 " + ("Snakebite" if st.session_state.lang == "en" else "Rắn cắn"))
    else:
        st.subheader("🩹 " + ("Normal wound" if st.session_state.lang == "en" else "Vết thương thông thường"))

    st.markdown("\n".join([f"{i+1}. {s}" for i, s in enumerate(steps)]))

# ------------------
# PART 3: OpenAI integration, save, forms, chat UI
# ------------------

# OpenAI config from secrets (Spaces: set in Secrets panel)
OPENAI_api_key = st.secrets.get("OPENAI_API_KEY")
OPENAI_model = st.secrets.get("OPENAI_MODEL", "gpt-4o-mini")

# Minimal robust client creation supporting modern + legacy SDKs

def make_openai_client(api_key):
    if not api_key:
        return None, "OpenAI API key not set. Set OPENAI_API_KEY in secrets."
    try:
        from openai import OpenAI
        client = OpenAI(api_key=api_key)
        return client, None
    except Exception:
        try:
            import openai
            openai.api_key = api_key
            return openai, None
        except Exception as e:
            return None, f"Failed to create OpenAI client: {e}"


# Simple extractor (keeps many shapes)

def extract_text_from_response(resp):
    try:
        if resp is None:
            return ""
        if isinstance(resp, str):
            return resp.strip()
        # dict-like shapes
        if isinstance(resp, dict):
            # responses API often has 'output' or 'choices'
            if "output_text" in resp and isinstance(resp["output_text"], str):
                return resp["output_text"].strip()
            if "choices" in resp and isinstance(resp["choices"], list) and resp["choices"]:
                c = resp["choices"][0]
                if isinstance(c, dict):
                    m = c.get("message") or c
                    if isinstance(m, dict) and isinstance(m.get("content"), str):
                        return m["content"].strip()
                    if isinstance(m.get("text"), str):
                        return m.get("text").strip()
            # responses API: 'output' might be list
            if "output" in resp:
                out = resp["output"]
                if isinstance(out, list) and out:
                    first = out[0]
                    if isinstance(first, dict):
                        c = first.get("content") or first.get("text")
                        if isinstance(c, list) and c:
                            f = c[0]
                            if isinstance(f, dict) and isinstance(f.get("text"), str):
                                return f.get("text").strip()
                        if isinstance(c, str):
                            return c.strip()
        # objects with attributes
        for attr in ("output_text", "text", "content"):
            if hasattr(resp, attr):
                val = getattr(resp, attr)
                if isinstance(val, str):
                    return val.strip()
                try:
                    return str(val).strip()
                except Exception:
                    pass
        return str(resp).strip()
    except Exception:
        try:
            return str(resp)
        except Exception:
            return ""


def normalize_assistant_text(x):
    if x is None:
        return ""
    if isinstance(x, str):
        return x.strip()
    try:
        return extract_text_from_response(x)
    except Exception:
        return str(x)


# Query wrapper: tries modern chat, responses, legacy shapes

def query_openai(system_prompt: str, user_message: str, max_tokens=800, temperature=0.2):
    """
    Robust OpenAI query that supports:
    - modern client.chat.create(...)
    - legacy openai.ChatCompletion.create(...)
    - client.responses.create(...) (tries several kwarg combinations)
    - older Completion.create(...)
    Returns {"text": "..."} on success or {"error": "..."} on failure.
    """
    client, err = make_openai_client(OPENAI_api_key)
    if err:
        return {"error": err}

    model = OPENAI_model or "gpt-4o-mini"

    try:
        # 1) Modern OpenAI SDK: client.chat.create(...)
        if hasattr(client, "chat") and hasattr(client.chat, "create"):
            resp = client.chat.create(
                model=model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_message},
                ],
                temperature=temperature,
                max_tokens=max_tokens,
            )
            return {"text": extract_text_from_response(resp)}

        # 2) Legacy openai package: ChatCompletion.create(...)
        if hasattr(client, "ChatCompletion") and hasattr(client.ChatCompletion, "create"):
            resp = client.ChatCompletion.create(
                model=model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_message},
                ],
                temperature=temperature,
                max_tokens=max_tokens,
            )
            return {"text": extract_text_from_response(resp)}

        # 3) Responses API: client.responses.create(...)
        if hasattr(client, "responses") and hasattr(client.responses, "create"):
            prompt = system_prompt + "\n\n" + user_message
            response = None
            last_exc = None

            # Try several kwarg shapes because different SDK versions accept different names
            candidate_kw_sets = [
                {"model": model, "input": prompt, "max_output_tokens": max_tokens, "temperature": temperature},
                {"model": model, "input": prompt, "max_tokens": max_tokens, "temperature": temperature},
                {"model": model, "input": prompt, "temperature": temperature},
                {"model": model, "input": prompt},
            ]

            for kw in candidate_kw_sets:
                try:
                    response = client.responses.create(**kw)
                    break
                except TypeError as te:
                    # argument mismatch for this kw set — try next
                    last_exc = te
                    continue
                except Exception as e:
                    # other runtime errors (auth/quota) -> return as error
                    return {"error": f"OpenAI responses request failed: {e}"}

            if response is None:
                return {"error": f"OpenAI responses client didn't accept tried args. Last error: {last_exc}"}

            return {"text": extract_text_from_response(response)}

        # 4) Very old Completion API fallback
        if hasattr(client, "Completion") and hasattr(client.Completion, "create"):
            prompt = system_prompt + "\n\n" + user_message
            resp = client.Completion.create(model=model, prompt=prompt, max_tokens=max_tokens, temperature=temperature)
            return {"text": extract_text_from_response(resp)}

        return {"error": "OpenAI client present but no compatible chat method found. Consider upgrading the openai package."}

    except Exception as e:
        return {"error": f"OpenAI request failed: {e}"}



# System prompt builder that respects UI language

def make_system_prompt(wound_type: str, victim_info: dict) -> str:
    lang = st.session_state.get("lang", "en")
    lang_instruction = {"en": "Please reply in English.", "vn": "Vui lòng trả lời bằng tiếng Việt."}
    base = (
        f"You are an expert first aid assistant. The victim has wound_type={wound_type}.\n"
        "- Give clear, actionable, short steps that a non-professional can follow.\n"
        "- Mention when to call emergency services or go to hospital.\n"
        f"- Use the victim info below to adjust guidance:\n{json.dumps(victim_info, indent=2)}\n"
        "Keep language concise. Use numbered short steps when appropriate.\n"
    )
    return base + lang_instruction.get(lang, "Please reply in English.")


# Save victim info to Firebase (requires db initialized)
def save_victim_info(victim: dict):
    if not db:
        # Fallback: append to local DATA_FILE
        try:
            all_v = json.loads(DATA_FILE.read_text(encoding="utf-8") or "{}")
        except Exception:
            all_v = {}
        record_id = datetime.utcnow().strftime("%Y%m%d%H%M%S")
        all_v[record_id] = victim
        DATA_FILE.write_text(json.dumps(all_v, ensure_ascii=False, indent=2), encoding="utf-8")
        return victim
    record_id = st.session_state.get("firebase_record_id")
    if not record_id:
        # If we don't have an id, generate one under victims
        record_id = datetime.utcnow().strftime("%Y%m%d%H%M%S")
    victim_data = victim.copy()
    victim_data["timestamp"] = datetime.utcnow().isoformat() + "Z"
    try:
        db.child("victims").child(record_id).update(victim_data)
        return victim_data
    except Exception as e:
        raise RuntimeError(f"Failed to save victim info: {e}")


# ------------------
# Streamlit UI & Flow
# ------------------
st.title("🩺 First Aid Assistant — Firebase Connected (EN/VN)")

# Auto display latest on load
render_latest_firebase_data()

st.markdown("---")
if st.button("🔄 " + ("Refresh Firebase data" if st.session_state.lang == "en" else "Làm mới dữ liệu Firebase")):
    render_latest_firebase_data(show_toast=True)

# -----------------------
# Advanced assistance block (fixed: adds weight/height, VN/EN, saves to victim profile)
# -----------------------
if st.session_state.get("basic_done") and st.session_state.get("wound_type"):
    st.header("Advanced assistance")
    st.markdown(f"**Wound type:** `{st.session_state.wound_type}`")

    with st.form("victim_form", clear_on_submit=True):
        st.subheader(t("victim_name"))
        name = st.text_input(t("victim_name"))
        phone = st.text_input(t("phone"),
                              help=("Emergency contact number" if st.session_state.lang == "en" else "Số liên hệ khẩn cấp"))
        age = st.text_input(t("age"),
                            help=("Enter an integer or leave blank" if st.session_state.lang == "en" else "Nhập số nguyên hoặc để trống"))
        # new localized weight / height fields
        weight = st.text_input(t("weight"), help=t("weight_help"))
        height = st.text_input(t("height"), help=t("height_help"))
        # localized notes label (falls back if t("notes") missing)
        notes_label = t("notes") if "notes" in TRANSLATIONS else (
            "Notes (allergies, meds, conscious status)" if st.session_state.lang == "en"
            else "Ghi chú thêm (dị ứng, thuốc đang dùng, tình trạng tỉnh táo)"
        )
        notes = st.text_area(notes_label, max_chars=500)

        col1, col2 = st.columns(2)
        with col1:
            submit = st.form_submit_button(t("save_info"))
        with col2:
            emergency_submit = st.form_submit_button(t("emergency_btn"))

        # common validation (now parses weight & height as floats)
        def validate_fields(age_str, phone_str, weight_str, height_str):
            age_v = None
            weight_v = None
            height_v = None

            # age
            if age_str and age_str.strip():
                try:
                    age_v = int(age_str.strip())
                except Exception:
                    st.error(("Age must be an integer (or leave blank)." if st.session_state.lang == "en"
                              else "Tuổi phải là số nguyên (hoặc để trống)."))
                    st.stop()

            # weight
            if weight_str and weight_str.strip():
                try:
                    weight_v = float(weight_str.strip())
                except Exception:
                    st.error(("Weight must be a number (kg) or leave blank." if st.session_state.lang == "en"
                              else "Cân nặng phải là số (kg) hoặc để trống."))
                    st.stop()

            # height
            if height_str and height_str.strip():
                try:
                    height_v = float(height_str.strip())
                except Exception:
                    st.error(("Height must be a number (cm) or leave blank." if st.session_state.lang == "en"
                              else "Chiều cao phải là số (cm) hoặc để trống."))
                    st.stop()

            # phone
            phone_clean = phone_str.strip()
            if phone_clean and not phone_clean.replace("+", "").isdigit():
                st.error(("Phone number should contain digits only." if st.session_state.lang == "en"
                          else "Số điện thoại chỉ chứa chữ số."))
                st.stop()

            return age_v, weight_v, height_v, phone_clean

        if (submit or emergency_submit) and not st.session_state.victim_saved:
            # validate including weight/height
            age_v, weight_v, height_v, phone_clean = validate_fields(age, phone, weight, height)

            victim = {
                "name": name or "unknown",
                "phone": phone_clean or None,
                "age": age_v,
                "weight_kg": weight_v,
                "height_cm": height_v,
                "notes": notes or "",
                "wound_type": st.session_state.wound_type,
                "saved_at": datetime.utcnow().isoformat() + "Z",
                "emergency": bool(emergency_submit),
            }

            try:
                record = save_victim_info(victim)
                st.session_state.last_saved_record = record
                st.session_state.victim_saved = True
                if emergency_submit:
                    st.error(("EMERGENCY recorded. Call local emergency services now." if st.session_state.lang == "en"
                              else "Đã ghi nhận KHẨN CẤP. Gọi dịch vụ khẩn cấp ngay."))
                else:
                    st.success(("Victim info saved." if st.session_state.lang == "en" else "Đã lưu thông tin nạn nhân."))
            except Exception as e:
                st.error(str(e))
                st.write(traceback.format_exc())

            # For non-emergency: ask OpenAI for guidance
            if submit and st.session_state.victim_saved:
                system_prompt = make_system_prompt(st.session_state.wound_type, victim)
                initial_user_msg = (
                    "The patient has completed the basic steps. Provide the next, in-depth first aid guidance and any monitoring steps. Keep guidance concise and list immediate priorities first."
                    if st.session_state.lang == "en"
                    else "Bệnh nhân đã hoàn thành các bước cơ bản. Cung cấp hướng dẫn sơ cứu chi tiết tiếp theo và các bước giám sát. Giữ gọn và nêu ưu tiên trước."
                )
                with st.spinner(("Contacting assistant..." if st.session_state.lang == "en" else "Đang liên hệ trợ lý...")):
                    resp = query_openai(system_prompt, initial_user_msg)
                if "error" in resp:
                    st.error(f"Assistant error: {resp['error']}")
                else:
                    assistant_text = normalize_assistant_text(resp.get("text"))
                    st.session_state.assistant_text = assistant_text
                    st.session_state.chat_history = [{"role": "assistant", "text": assistant_text}]
                    st.markdown("**" + t("assistant_guidance") + "**")
                    st.text_area(t("assistant_guidance"), value=assistant_text, height=250, key="assistant_output")


    # Chat follow-ups
    st.subheader(("Ask follow-up questions" if st.session_state.lang == "en" else "Hỏi thêm"))
    st.markdown(("Ask the assistant further questions related to the victim or care. Keep questions concise." if st.session_state.lang == "en" else "Hỏi trợ lý các câu liên quan tới nạn nhân hoặc chăm sóc. Giữ câu hỏi ngắn."))
    col_in, col_btn = st.columns([4, 1])
    with col_in:
        user_q = st.text_input(("Your question" if st.session_state.lang == "en" else "Câu hỏi của bạn"), key="chat_input")
    with col_btn:
        if st.button(t("ask_btn")):
            if not user_q or user_q.strip() == "":
                st.warning(("Type a question first." if st.session_state.lang == "en" else "Nhập câu hỏi trước."))
            else:
                victim_info = st.session_state.get("last_saved_record") or {}
                system_prompt = make_system_prompt(st.session_state.wound_type, victim_info)
                with st.spinner(("Assistant is replying..." if st.session_state.lang == "en" else "Trợ lý đang trả lời...")):
                    resp = query_openai(system_prompt, user_q)
                if "error" in resp:
                    st.error(f"Assistant error: {resp['error']}")
                else:
                    assistant_reply = normalize_assistant_text(resp.get("text"))
                    history = st.session_state.get("chat_history", [])
                    history.append({"role": "user", "text": user_q})
                    history.append({"role": "assistant", "text": assistant_reply})
                    st.session_state.chat_history = history
                    st.success(("Assistant responded. See conversation below." if st.session_state.lang == "en" else "Trợ lý đã trả lời. Xem cuộc trò chuyện bên dưới."))

    # Conversation display
    if st.session_state.get("chat_history"):
        st.subheader(("Conversation log" if st.session_state.lang == "en" else "Nhật ký cuộc trò chuyện"))
        for turn in st.session_state.chat_history:
            role = turn.get("role", "assistant")
            text = turn.get("text", "")
            if role == "assistant":
                st.markdown(f"**Assistant:** {text}")
            else:
                st.markdown(f"**User:** {text}")

else:
    st.info(("Start by reviewing the quick steps for Snakebite or Normal wound and confirm completion using the checkbox, then click Next → Advanced." if st.session_state.lang == "en" else "Bắt đầu bằng cách xem các bước nhanh cho rắn cắn hoặc vết thương thông thường và xác nhận, sau đó nhấp Tiếp → Nâng cao."))

st.markdown("---")
st.caption(
    ("🔒 Security: Prefer setting OPENAI_API_KEY as a secret. If you paste a key into the UI it is stored only in session memory and not written to disk." if st.session_state.lang == "en" else "🔒 Bảo mật: Ưu tiên đặt OPENAI_API_KEY trong Secrets. Nếu dán vào UI thì chỉ lưu trong session và không viết ra đĩa.")
)

# END of file