# 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 pathlib import Path
import asyncio

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)

# --- STREAMLIT UI ---

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

    # Header with logo and branding
    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")

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

            # Papers
            st.markdown("### 📚 Most Relevant Papers")
            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)

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

    # 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()