# ── app.py ─────────────────────────────────────────────────────────────
import asyncio
import re

import streamlit as st
import pandas as pd
import plotly.express as px
from fpdf import FPDF
from streamlit_agraph import agraph

from mcp.orchestrator import orchestrate_search, answer_ai_question
from mcp.knowledge_graph import build_agraph
from mcp.graph_metrics import build_nx, get_top_hubs, get_density
from mcp.protocols import draft_protocol

# ── Streamlit setup ────────────────────────────────────────────────────
st.set_page_config(layout="wide", page_title="MedGenesis AI")
if "res" not in st.session_state:
    st.session_state.res = None

st.title("🧬 MedGenesis AI")
llm = st.radio("LLM engine", ["openai", "gemini"], horizontal=True)
query = st.text_input("Enter biomedical question")

# PDF helper

def _make_pdf(papers):
    pdf = FPDF()
    pdf.add_page(); pdf.set_font("Helvetica", size=12)
    pdf.cell(0, 10, "MedGenesis AI – Results", ln=True, align="C"); pdf.ln(5)
    for i, p in enumerate(papers, 1):
        pdf.set_font("Helvetica", "B", 11)
        pdf.multi_cell(0, 7, f"{i}. {p.get('title','')}")
        pdf.set_font("Helvetica", size=9)
        body = f"{p.get('authors','')}
{p.get('summary','')}
{p.get('link','')}"
        pdf.multi_cell(0, 6, body); pdf.ln(3)
    return pdf.output(dest="S").encode("latin-1", errors="replace")

# Run search
if st.button("Run Search 🚀") and query:
    with st.spinner("Gathering data…"):
        st.session_state.res = asyncio.run(orchestrate_search(query, llm))
res = st.session_state.res
if not res:
    st.info("Enter a query and press Run Search")
    st.stop()

# Tabs
tabs = st.tabs(["Results", "Graph", "Clusters", "Variants", "Trials", "Metrics", "Visuals", "Protocols"])

# ── Results tab
with tabs[0]:
    for i, p in enumerate(res["papers"], 1):
        st.markdown(f"**{i}. [{p['title']}]({p['link']})**")
        st.write(p["summary"])
    c1, c2 = st.columns(2)
    c1.download_button("CSV", pd.DataFrame(res["papers"]).to_csv(index=False),
                       "papers.csv", "text/csv")
    c2.download_button("PDF", _make_pdf(res["papers"]),
                       "papers.pdf", "application/pdf")
    st.subheader("AI summary"); st.info(res["ai_summary"])

# ── Graph tab
with tabs[1]:
    nodes, edges, cfg = build_agraph(
        res["papers"], res["umls"], res["drug_safety"], res["umls_relations"]
    )
    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 n.color
    agraph(nodes, edges, cfg)

# ── Clusters tab
with tabs[2]:
    clusters = res.get("clusters", [])
    if clusters:
        df = pd.DataFrame({
            "title": [p['title'] for p in res['papers']],
            "cluster": clusters
        })
        st.write("### Paper Clusters")
        for c in sorted(set(clusters)):
            st.write(f"**Cluster {c}**")
            for t in df[df['cluster'] == c]['title']:
                st.write(f"- {t}")
    else:
        st.info("No clusters to show.")

# ── Variants tab
with tabs[3]:
    if res.get("variants"):
        st.json(res["variants"])
    else:
        st.warning("No variants found. Try 'TP53' or 'BRCA1'.")

# ── Trials tab
with tabs[4]:
    if res.get("clinical_trials"):
        st.json(res["clinical_trials"])
    else:
        st.warning("No trials found. Try a disease or drug.")

# ── Metrics tab
with tabs[5]:
    nodes_dicts = [n.__dict__ for n in nodes]
    edges_dicts = [e.__dict__ for e in edges]
    G = build_nx(nodes_dicts, edges_dicts)
    st.metric("Density", f"{get_density(G):.3f}")
    st.markdown("**Top hubs**")
    for nid, sc in get_top_hubs(G):
        lbl = next((n.label for n in nodes if n.id == nid), nid)
        st.write(f"- {lbl}: {sc:.3f}")

# ── Visuals tab
with tabs[6]:
    years = [p.get("published","")[:4] for p in res["papers"] if p.get("published")]
    if years:
        st.plotly_chart(px.histogram(years, nbins=10, title="Publication Year"))

# ── Protocols tab
with tabs[7]:
    proto_q = st.text_input("Enter hypothesis for protocol:", key="proto_q")
    if st.button("Draft Protocol") and proto_q.strip():
        with st.spinner("Generating protocol…"):
            proto = asyncio.run(draft_protocol(
                proto_q, context=res['ai_summary'], llm=llm
            ))
        st.subheader("Experimental Protocol")
        st.write(proto)
 ───────────────────────────────────────────────────
# In import section:
from mcp.embeddings import embed_texts, cluster_embeddings
from mcp.protocols import draft_protocol

# After creating tabs = st.tabs([...,'Clusters','Protocols']):
with tabs[-2]:  # second last tab = Clusters
    if res.get('clusters'):
        df = pd.DataFrame({
            'title': [p['title'] for p in res['papers']],
            'cluster': res['clusters']
        })
        st.write("### Paper Clusters")
        for c in sorted(set(res['clusters'])):
            st.write(f"**Cluster {c}**")
            for t in df[df['cluster']==c]['title']:
                st.write(f"- {t}")
    else:
        st.info("No clusters to show.")

with tabs[-1]:  # last tab = Protocols
    proto_q = st.text_input("Enter hypothesis for protocol:", key="proto_q")
    if st.button("Draft Protocol") and proto_q.strip():
        with st.spinner("Generating protocol…"):
            proto = asyncio.run(draft_protocol(
                proto_q, context=res['ai_summary'], llm=llm))
        st.subheader("Experimental Protocol")
        st.write(proto)