# app.py

import os
import streamlit as st
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
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_knowledge_graph
from pathlib import Path
import pandas as pd
from fpdf import FPDF
import asyncio
import plotly.express as px
import streamlit.components.v1 as components

ROOT_DIR = Path(__file__).resolve().parent
LOGO_PATH = ROOT_DIR / "assets" / "logo.png"

# --- FASTAPI BACKEND ---

api = FastAPI(
    title="MedGenesis MCP Server",
    version="2.0.0",
    description="MedGenesis AI unifies PubMed, ArXiv, OpenFDA, UMLS, and GPT-4o into a single biomedical intelligence platform."
)

api.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@api.post("/unified_search", response_model=UnifiedSearchResult)
async def unified_search_endpoint(data: UnifiedSearchInput):
    return await orchestrate_search(data.query)

@api.post("/ask_ai")
async def ask_ai_endpoint(question: str, context: str = ""):
    return await answer_ai_question(question, context)

# --- PDF Export Utility ---

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

# --- STREAMLIT UI ---

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

    # --- SIDEBAR WORKSPACE ---
    with st.sidebar:
        st.header("🗂️ Your Workspace")
        saved_queries = get_workspace()
        if saved_queries:
            for i, item in enumerate(saved_queries, 1):
                with st.expander(f"{i}. {item['query']}"):
                    st.write("**AI Summary:**", item["result"]["ai_summary"])
                    st.write("**First Paper:**", item["result"]["papers"][0]["title"] if item["result"]["papers"] else "None")
                    df = pd.DataFrame(item["result"]["papers"])
                    st.download_button(
                        label="📥 Download as CSV",
                        data=df.to_csv(index=False),
                        file_name=f"workspace_query_{i}.csv",
                        mime="text/csv",
                    )
        else:
            st.info("Run a search and save it here!")

    # --- MAIN APP HEADER ---
    col1, col2 = st.columns([0.15, 0.85])
    with col1:
        if LOGO_PATH.exists():
            st.image(str(LOGO_PATH), width=100)
        else:
            st.markdown("🧬")
    with col2:
        st.markdown("""
            ## 🧬 **MedGenesis AI** – Biomedical Research Reimagined  
            *Unified Intelligence from PubMed, ArXiv, OpenFDA, UMLS, and GPT-4o*
        """)
        st.caption("Created by Oluwafemi Idiakhoa | Hugging Face Spaces")

    st.markdown("---")

    # Unified Semantic Search
    st.subheader("🔍 Unified Semantic Search")
    query = st.text_input("Enter your biomedical research question:", placeholder="e.g. New treatments for glioblastoma using CRISPR")

    results = None
    if st.button("Run Search 🚀"):
        with st.spinner("Thinking... Gathering and analyzing data across 5 systems..."):
            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]:
            for i, paper in enumerate(results["papers"], 1):
                st.markdown(f"**{i}. [{paper['title']}]({paper['link']})**  \n*{paper['authors']}* ({paper['source']})")
                st.markdown(f"<div style='font-size: 0.9em; color: gray'>{paper['summary']}</div>", unsafe_allow_html=True)

            # Save to workspace (user-initiated, for clarity)
            if st.button("Save this search to Workspace"):
                save_query(query, results)
                st.success("Saved to your workspace!")

            # Export as CSV
            if results["papers"]:
                df = pd.DataFrame(results["papers"])
                csv = df.to_csv(index=False)
                st.download_button(
                    label="📥 Download results as CSV",
                    data=csv,
                    file_name="medgenesis_results.csv",
                    mime="text/csv",
                )

            # Export as PDF
            if results["papers"]:
                pdf_bytes = generate_pdf(results["papers"])
                st.download_button(
                    label="📄 Download results as PDF",
                    data=pdf_bytes,
                    file_name="medgenesis_results.pdf",
                    mime="application/pdf",
                )

            # UMLS Concepts
            st.markdown("### 🧠 Biomedical Concept Enrichment (UMLS)")
            for concept in results["umls"]:
                if concept["cui"]:
                    st.markdown(f"🔹 **{concept['name']}** (CUI: `{concept['cui']}`): {concept['definition'] or 'No definition available'}")

            # Drug Safety
            st.markdown("### 💊 Drug Safety Insights (OpenFDA)")
            for drug_report in results["drug_safety"]:
                if drug_report:
                    st.json(drug_report)

            # AI Summary
            st.markdown("### 🤖 AI-Powered Summary")
            st.info(results["ai_summary"])

            # Suggested Reading
            st.markdown("### 📖 Suggested Links")
            for link in results["suggested_reading"]:
                st.write(f"- {link}")

        # --- TAB 2: Knowledge Graph ---
        with tabs[1]:
            st.markdown("#### Explore Connections")
            kg_html_path = build_knowledge_graph(results["papers"], results["umls"], results["drug_safety"])
            with open(kg_html_path, 'r', encoding='utf-8') as f:
                kg_html = f.read()
            components.html(kg_html, height=600)

        # --- 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)
            # Placeholder for more charts

    # Follow-up AI Q&A
    st.markdown("---")
    st.subheader("💬 Ask AI a Follow-up Question")
    follow_up = st.text_input("What do you want to ask based on the above?", placeholder="e.g. What's the most promising therapy?")
    if st.button("Ask AI"):
        with st.spinner("Analyzing and responding..."):
            ai_answer = asyncio.run(answer_ai_question(follow_up, context=query))
            st.success("AI's Response:")
            st.write(ai_answer["answer"])

    # Footer
    st.markdown("---")
    st.markdown(
        "<div style='text-align: center; font-size: 0.9em;'>"
        "✨ Built with ❤️ by <strong>Oluwafemi Idiakhoa</strong> • Powered by FastAPI, Streamlit, Hugging Face, OpenAI, UMLS, OpenFDA, and NCBI</div>",
        unsafe_allow_html=True
    )

# --- MAIN ENTRY ---

if __name__ == "__main__":
    import sys
    if "runserver" in sys.argv:
        import uvicorn
        uvicorn.run(api, host="0.0.0.0", port=7860)
    else:
        render_ui()