import streamlit as st
from services.auth import authenticator, require_login
from services.logger import logger
from services.metrics import CHAT_COUNT, OPTIMIZE_COUNT
from agent.gemini_agent import chat_with_gemini
from clinical_nlp.umls_bioportal import lookup_umls, lookup_bioportal
from quantum.bf_dcqo import optimize_hubo
from services.pdf_report import generate_pdf
from repositories.chat_repo import ChatRepo

# Initialize DB
from models.db import init_db
init_db()

# UI
username = require_login()
st.set_page_config(page_title="Quantum Health AI", layout="wide")
st.image("assets/logo.png", width=64)
st.title(f"Welcome, {username}!")

tab1, tab2 = st.tabs(["🩺 Consult", "📊 Reports"])

with tab1:
    query = st.text_area("Enter clinical query or symptoms:")
    if st.button("Ask Gemini"):
        CHAT_COUNT.inc()
        with st.spinner("Consulting AI..."):
            response = chat_with_gemini(username, query)
        st.markdown(f"**AI**: {response}")
        # save chat
        ChatRepo().save(username, query, response)
        # clinical NLP
        with st.expander("UMLS Results"):
            st.write(lookup_umls(query))
        with st.expander("BioPortal Results"):
            st.write(lookup_bioportal(query))

    if st.button("Quantum Optimize"):
        OPTIMIZE_COUNT.inc()
        with st.spinner("Running quantum optimizer..."):
            result = optimize_hubo({"query":query})
        st.json(result)

with tab2:
    if st.button("Generate PDF Report"):
        # gather last 5 messages
        pdf_data = {"Last Chats": ChatRepo().get_recent(username, limit=5)}
        fname = generate_pdf(pdf_data)
        st.success("Report Generated")
        st.download_button("Download Report", data=open(fname,"rb"), file_name=fname)