# app.py

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

from mcp.orchestrator import orchestrate_search, answer_ai_question
from mcp.schemas import UnifiedSearchInput, UnifiedSearchResult
from mcp.workspace import get_workspace, save_query
from mcp.knowledge_graph import build_agraph
from streamlit_agraph import agraph, Node, Edge, Config

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

def generate_pdf(papers):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    pdf.cell(200, 10, "MedGenesis AI - Search Results", ln=True, align="C")
    pdf.ln(10)
    for i, p in enumerate(papers, 1):
        pdf.set_font("Arial", "B", 12)
        pdf.multi_cell(0, 10, f"{i}. {p['title']}")
        pdf.set_font("Arial", "", 10)
        pdf.multi_cell(0, 8, f"Authors: {p['authors']}\nLink: {p['link']}\nSummary: {p['summary']}\n")
        pdf.ln(2)
    return pdf.output(dest="S").encode("latin-1")

def render_ui():
    st.set_page_config(page_title="MedGenesis AI", layout="wide")

    with st.sidebar:
        st.header("🗂️ Workspace")
        for i, item in enumerate(get_workspace(), 1):
            with st.expander(f"{i}. {item['query']}"):
                st.write("**AI Summary:**", item["result"]["ai_summary"])
                st.write("**Top Paper:**", item["result"]["papers"][0]["title"] if item["result"]["papers"] else "None")
                df = pd.DataFrame(item["result"]["papers"])
                st.download_button("📥 CSV", df.to_csv(index=False), f"workspace_{i}.csv", "text/csv")
        if not get_workspace():
            st.info("Run and save searches here.")

    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\n*Unified Intelligence from PubMed, ArXiv, OpenFDA, UMLS, GPT‑4o*")
        st.caption("Built by Oluwafemi Idiakhoa • Hugging Face Spaces")

    st.markdown("---")
    query = st.text_input("🔍 Enter research question:", placeholder="e.g. CRISPR treatments for glioblastoma")

    results = None
    if st.button("Run Search 🚀"):
        with st.spinner("Analyzing..."):
            results = asyncio.run(orchestrate_search(query))
            st.success("Search complete")

    if results:
        tabs = st.tabs(["📝 Results", "🗺️ Knowledge Graph", "📊 Visualizations"])

        with tabs[0]:
            st.header("📚 Papers")
            for i, p in enumerate(results["papers"], 1):
                st.markdown(f"**{i}. [{p['title']}]({p['link']})** *{p['authors']}*")
                st.markdown(f"<div style='color:gray'>{p['summary']}</div>", unsafe_allow_html=True)
            if st.button("Save to Workspace"):
                save_query(query, results); st.success("Saved!")
            df = pd.DataFrame(results["papers"])
            st.download_button("📥 CSV", df.to_csv(index=False), "results.csv", "text/csv")
            pdf = generate_pdf(results["papers"])
            st.download_button("📄 PDF", pdf, "results.pdf", "application/pdf")
            st.subheader("🧠 UMLS Concepts")
            for c in results["umls"]:
                if c.get("cui"):
                    st.markdown(f"- **{c['name']}** (CUI:{c['cui']}): {c.get('definition','')}")
            st.subheader("💊 Drug Safety")
            for d in results["drug_safety"]:
                st.json(d)
            st.subheader("🤖 AI Summary")
            st.info(results["ai_summary"])

        with tabs[1]:
            st.header("🔎 Knowledge Graph")
            search_term = st.text_input("Highlight node containing text:", value="")
            try:
                nodes, edges, config = build_agraph(results["papers"], results["umls"], results["drug_safety"])
                if search_term:
                    pat = re.compile(re.escape(search_term), re.IGNORECASE)
                    for node in nodes:
                        if pat.search(node.label) or (hasattr(node, "tooltip") and pat.search(getattr(node, "tooltip", ""))):
                            node.color = "#f1c40f"; node.size = max(node.size, 30)
                        else:
                            node.color = "#ddd"
                agraph(nodes=nodes, edges=edges, config=config)
            except Exception as e:
                st.warning("Graph unavailable: " + str(e))

        with tabs[2]:
            years = [p["published"] for p in results["papers"] if p.get("published")]
            if years:
                st.plotly_chart(px.histogram(years, title="Pub Year"))

    st.markdown("---")
    st.subheader("💬 Ask AI a Follow-up")
    fq = st.text_input("Ask based on above:", "")
    if st.button("Ask AI"):
        with st.spinner("Thinking..."):
            ans = asyncio.run(answer_ai_question(fq, context=query))
            st.write(ans.get("answer", ans))

    st.markdown("---")
    st.caption("✨ Built by Oluwafemi Idiakhoa • Powered by Streamlit & AI")

if __name__ == "__main__":
    render_ui()