import os

import requests
import streamlit as st

BACKEND_URL = os.environ["BACKEND_URL"].rstrip("/")
HF_TOKEN = os.environ.get("HF_TOKEN")

st.set_page_config(page_title="Medical Assistant", page_icon="🩺")

st.title("Medical Assistant")
st.caption(
    "Ask a medical question and get an answer grounded in the MedQuAD dataset, "
    "with sources cited below each response. Not a substitute for professional "
    "medical advice."
)

with st.expander("ℹ️ How this works"):
    st.markdown(
        """
This chat UI is the **frontend** half of a two-Space RAG system — it carries
no model logic of its own, it just calls a backend API.

**Pipeline:** your question → embedded with
[`sentence-transformers/all-MiniLM-L6-v2`](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
→ top-k passages retrieved from a FAISS index built over the
[MedQuAD](https://www.kaggle.com/datasets/pythonafroz/medquad-medical-question-answer-for-ai-research)
medical Q&A dataset → assembled into a prompt that restricts the model to
answering **only** from that retrieved context (no diagnosis, no
prescriptions) → answer generated by a hosted LLM and returned with its
sources.

**Models in use:** a stock, unmodified sentence-embedding model for
retrieval, plus a remote hosted LLM for generation. Nothing here is
fine-tuned or trained — there's no custom model artifact behind this
project, so there's no separate model card.

**Limitations:** answers are only as good as MedQuAD's coverage; the
backend has no control over the hosted generation model beyond the
grounding prompt; no published latency/accuracy numbers yet.

Full details: see this Space's
[README](https://huggingface.co/spaces/whosouravsharma/Medical-Assistant-ChatBot/blob/main/README.md)
and the [backend Space](https://huggingface.co/spaces/whosouravsharma/medical-assistant-ai-backend).
        """
    )

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])


def ask_backend(question: str) -> str:
    headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
    response = requests.post(
        f"{BACKEND_URL}/ask",
        json={"question": question},
        headers=headers,
        timeout=60,
    )
    response.raise_for_status()
    result = response.json()
    sources = "\n".join(f"- {s['source']} — {s['focus_area']}" for s in result["sources"])
    return f"{result['answer']}\n\n**Sources**\n{sources}"


if question := st.chat_input("What are the symptoms of glaucoma?"):
    st.session_state.messages.append({"role": "user", "content": question})
    with st.chat_message("user"):
        st.markdown(question)

    with st.chat_message("assistant"):
        try:
            answer = ask_backend(question)
        except requests.RequestException as exc:
            answer = f"Sorry, the backend request failed: {exc}"
        st.markdown(answer)

    st.session_state.messages.append({"role": "assistant", "content": answer})