# 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

# Initialize FastAPI app for API users
api = FastAPI(title="MCP Research Server", version="2.0")
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 for Hugging Face Space
def render_ui():
    st.set_page_config(page_title="Ultimate Research Assistant", page_icon=":microscope:", layout="wide")
    st.image("assets/logo.png", width=100)
    st.title("🔬 Next-Gen AI-Powered Biomedical Research Assistant")
    st.markdown(
        """
        *Combine the power of ArXiv, PubMed, UMLS, OpenFDA, and OpenAI.  
        Get instant, unified, semantically-ranked answers—plus drug safety, concept enrichment, and expert Q&A!*
        """
    )

    query = st.text_input("Enter your research question or topic:", value="What are the latest treatments for Alzheimer's disease?")
    if st.button("Run Unified Search 🚀"):
        with st.spinner("Retrieving and synthesizing knowledge..."):
            results = orchestrate_search(query)
            st.success("Here are the results!")
            for i, paper in enumerate(results['papers'], 1):
                st.markdown(f"**{i}. [{paper['title']}]({paper['link']})**  \n*{paper['authors']}*")
                st.write(paper['summary'])
            st.subheader("UMLS Concept Enrichment")
            for c in results['umls']:
                st.write(f"**{c['name']}** (CUI: {c['cui']}): {c['definition']}")
            st.subheader("Drug & Safety Insights")
            for d in results['drug_safety']:
                st.write(d)
            st.subheader("AI-Generated Synthesis")
            st.info(results['ai_summary'])
            st.markdown("#### 📚 Suggested Reading")
            for link in results['suggested_reading']:
                st.write(f"- {link}")

    st.markdown("---")
    st.subheader("🤖 Ask a follow-up (AI Q&A):")
    follow_up = st.text_input("Type your question here:")
    if st.button("Ask AI"):
        with st.spinner("AI is thinking..."):
            answer = answer_ai_question(follow_up, context=query)
            st.success("AI says:")
            st.write(answer['answer'])

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