# 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

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")

    # Workspace in sidebar
    with st.sidebar:
        st.header("🗂️ Workspace")
        ws = get_workspace()
        if ws:
            for i, item in enumerate(ws, 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"query_{i}.csv", "text/csv")
        else:
            st.info("Run & save searches here.")

    # Header and logo
    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 biomedical research question:", placeholder="e.g. CRISPR treatments for glioblastoma")

    results = None
    if st.button("Run Search 🚀"):
        with st.spinner("Gathering and analyzing data..."):
            results = asyncio.run(orchestrate_search(query))
            st.success("✅ Search complete")

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

        # Tab 1: Results
        with tabs[0]:
            st.header("📚 Most Relevant Papers")
            for i, p in enumerate(results["papers"], 1):
                st.markdown(f"**{i}. [{p['title']}]({p['link']})**  \n*{p['authors']}* ({p['source']})")
                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 to workspace!")
            df = pd.DataFrame(results["papers"])
            st.download_button("📥 Download CSV", df.to_csv(index=False), "results.csv", "text/csv")
            pdf = generate_pdf(results["papers"])
            st.download_button("📄 Download 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','No definition')}")
            st.subheader("💊 Drug Safety (OpenFDA)")
            for d in results["drug_safety"]:
                st.json(d)
            st.subheader("🤖 AI-Powered Summary")
            st.info(results["ai_summary"])

        # Tab 2: Knowledge Graph (with Node Search Highlight)
        with tabs[1]:
            st.header("🗺️ Knowledge Graph Explorer")
            search_term = st.text_input("🔎 Highlight nodes containing:", value="")
            try:
                nodes, edges, config = build_agraph(results["papers"], results["umls"], results["drug_safety"])
                if search_term.strip():
                    pattern = re.compile(re.escape(search_term), re.IGNORECASE)
                    for node in nodes:
                        if pattern.search(node.label) or (hasattr(node, "tooltip") and pattern.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("Knowledge graph unavailable: " + str(e))

        # Tab 3: Visualizations
        with tabs[2]:
            pub_years = [p["published"] for p in results["papers"] if p.get("published")]
            if pub_years:
                fig = px.histogram(pub_years, nbins=10, title="Publication Year Distribution")
                st.plotly_chart(fig)

    # Follow-up AI Q&A
    st.markdown("---")
    st.subheader("💬 Ask a Follow-up Question")
    follow_up = st.text_input("Based on above:", placeholder="e.g. What's the most promising therapy?")
    if st.button("Ask AI"):
        with st.spinner("Analyzing..."):
            ai_ans = asyncio.run(answer_ai_question(follow_up, context=query))
            st.write(ai_ans.get("answer", ai_ans))

    st.markdown("---")
    st.caption("✨ Built by Oluwafemi Idiakhoa • Powered by FastAPI, Streamlit, OpenAI, UMLS, OpenFDA, NCBI")

if __name__ == "__main__":
    render_ui()