import streamlit as st
import requests
import os 

API_URL = os.getenv("API_URL", "http://127.0.0.1:8000")

st.set_page_config(page_title="GeneSeek V2", page_icon="🧬", layout="wide")
st.title("GeneSeek V2: Clinical Trial Assistant")

with st.sidebar:
    st.header("Data Ingestion")
    uploaded_file = st.file_uploader("Upload Protocol (PDF/TXT)", type=["pdf", "txt"])
    
    if uploaded_file and st.button("Ingest Document"):
        with st.spinner("Indexing..."):
            files = {"file": (uploaded_file.name, uploaded_file, uploaded_file.type)}
            try:
                res = requests.post(f"{API_URL}/upload", files=files)
                if res.status_code == 200:
                    st.success("File Indexed!")
                else:
                    st.error(f"Error: {res.text}")
            except Exception as e:
                st.error(f"Connection Failed: {e}")

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

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

if prompt := st.chat_input("Ask a clinical question..."):
    st.chat_message("user").markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            try:
                res = requests.post(f"{API_URL}/chat", json={"question": prompt})
                if res.status_code == 200:
                    data = res.json()
                    st.markdown(data["answer"])
                    
                    with st.expander("View Citations"):
                        for i, ctx in enumerate(data["contexts"]):
                            st.markdown(f"**Ref {i+1}:** {ctx['metadata']}")
                            st.caption(ctx['content'])
                            st.divider()

                    st.session_state.messages.append({"role": "assistant", "content": data["answer"]})
                else:
                    st.error(f"Backend Error: {res.text}")
            except Exception as e:
                st.error(f"Connection Error: {e}")