# app.py  –  MedGenesis AI (CPU-only, dual-LLM)

import asyncio, re
from pathlib import Path
import streamlit as st, pandas as pd, plotly.express as px
from fpdf import FPDF
from streamlit_agraph import agraph

from mcp.orchestrator   import orchestrate_search, answer_ai_question
from mcp.workspace      import get_workspace, save_query
from mcp.knowledge_graph import build_agraph
from mcp.graph_metrics   import build_nx, get_top_hubs, get_density
from mcp.alerts          import check_alerts

ROOT  = Path(__file__).parent
LOGO  = ROOT / "assets" / "logo.png"

# ────────────────────────────────────────────────────────────────
def _pdf(papers):
    pdf=FPDF(); pdf.add_page(); pdf.set_font("Arial",size=11)
    for i,p in enumerate(papers,1):
        pdf.multi_cell(0,7,f"{i}. {p['title']} – {p['authors']}\n{p['summary']}\n{p['link']}\n")
        pdf.ln(1)
    return pdf.output(dest="S").encode("latin-1")

def _sidebar_workspace():
    with st.sidebar:
        st.header("🗂️ Workspace")
        for i,item in enumerate(get_workspace(),1):
            with st.expander(f"{i}. {item['query']}"):
                st.write(item["result"]["ai_summary"])

# ────────────────────────────────────────────────────────────────
def render_ui():
    st.set_page_config("MedGenesis AI", layout="wide")
    _sidebar_workspace()

    # ── Header
    col1,col2 = st.columns([0.15,0.85])
    with col1:
        if LOGO.exists(): st.image(str(LOGO), width=100)
    with col2:
        st.markdown("## 🧬 **MedGenesis AI**")
        st.caption("Multi-source biomedical assistant · OpenAI / Gemini LLMs")

    llm = st.radio("Choose LLM Engine", ["openai","gemini"], horizontal=True)
    query = st.text_input("Enter biomedical question…",
                          placeholder="e.g. CRISPR glioblastoma therapy")

    # 🔔 Alert check
    if get_workspace():
        try:
            news = asyncio.run(check_alerts([q["query"] for q in get_workspace()]))
            if news:
                st.sidebar.subheader("🔔 New Papers")
                for q, lnks in news.items(): st.sidebar.write(f"**{q}** – {len(lnks)} new")
        except Exception: pass

    # ── Run
    if st.button("Run Search 🚀") and query:
        with st.spinner("Gathering literature & biomedical data…"):
            res = asyncio.run(orchestrate_search(query, llm))
        st.success(f"Completed with **{res['llm_used'].title()}**")

        # Tabs
        tabs = st.tabs(["Results","Genes","Trials","Graph","Metrics","Visuals"])
        # 1) Results
        with tabs[0]:
            for i,p in enumerate(res["papers"],1):
                st.markdown(f"**{i}. [{p['title']}]({p['link']})**  – *{p['authors']}*")
                st.write(p["summary"])
            st.download_button("📥 CSV", pd.DataFrame(res["papers"]).to_csv(index=False),
                               "papers.csv","text/csv")
            st.download_button("📄 PDF", _pdf(res["papers"]), "papers.pdf","application/pdf")
            st.info(res["ai_summary"])
            if st.button("Save to Workspace"): save_query(query,res); st.success("Saved!")

        # 2) Genes
        with tabs[1]:
            st.write("### Gene hits")
            for g in res["genes"]: st.write("-", g.get("name",g.get("geneid")))

        # 3) Trials
        with tabs[2]:
            if res["clinical_trials"]:
                for t in res["clinical_trials"]:
                    st.write(f"**{t['NCTId'][0]}** – {t['BriefTitle'][0]}")
            else:
                st.info("No trials (rate-limited or none found).")

        # 4) Graph
        with tabs[3]:
            nodes,edges,cfg = build_agraph(res["papers"],res["umls"],res["drug_safety"])
            hl = st.text_input("Highlight node…", key="hl")
            if hl:
                pat=re.compile(re.escape(hl),re.I)
                for n in nodes: 
                    n.color="#f1c40f" if pat.search(n.label) else "#d3d3d3"
            agraph(nodes=nodes,edges=edges,config=cfg)

        # 5) Metrics
        with tabs[4]:
            G=build_nx([n.__dict__ for n in nodes],[e.__dict__ for e in edges])
            st.metric("Density",f"{get_density(G):.3f}")
            st.write("**Top hubs**")
            for nid,sc in get_top_hubs(G): st.write(nid, f"{sc:.3f}")

        # 6) Visuals
        with tabs[5]:
            yrs=[p["published"] for p in res["papers"] if p.get("published")]
            if yrs: st.plotly_chart(px.histogram(yrs,title="Publication Year"))

        # Follow-up Q-A
        q2=st.text_input("Ask follow-up:")
        if st.button("Ask AI"): 
            ans=asyncio.run(answer_ai_question(q2,context=query,llm=llm))
            st.write(ans["answer"])
    else:
        st.info("Enter a question and press **Run Search 🚀**")

if __name__=="__main__":
    render_ui()