# Drug Repurposing Advisor: A Multi-Agent Workflow Example
# This example uses dummy data for demonstration.
# In a production system, replace the dummy data with real pharmaceutical databases.

import os
import json
import requests
import streamlit as st
from typing import List, Union, Tuple
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langgraph.graph import END, StateGraph, START
from langgraph.prebuilt import ToolNode
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict, Annotated
from typing import Sequence

# Dummy data for drug mechanism research and clinical trial outcomes
drug_mechanism_texts = [
    "Drug A: Inhibits enzyme X and modulates receptor Y; potential anti-inflammatory effects.",
    "Drug B: Blocks ion channel Z; has been shown to reduce oxidative stress in preclinical models.",
    "Drug C: Activates nuclear receptor W; exhibits neuroprotective properties."
]

clinical_trials_texts = [
    "Trial 1: Drug A repurposed for rheumatoid arthritis showed a 30% improvement in joint function.",
    "Trial 2: Drug B evaluated in a pilot study for neurodegenerative disorders demonstrated a reduction in symptom severity.",
    "Trial 3: Drug C tested in a phase II trial for multiple sclerosis reported significant reduction in relapse rates."
]

# Text splitting settings
splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
mechanism_docs = splitter.create_documents(drug_mechanism_texts)
clinical_docs = splitter.create_documents(clinical_trials_texts)

# Here you would typically create vector embeddings and vectorstores (e.g., using ChromaDB)
# For demonstration, we define simple retriever functions that return dummy results.
def mechanism_retriever(query: str) -> str:
    # Dummy search: return first document that mentions a keyword from the query
    for doc in drug_mechanism_texts:
        if any(word.lower() in doc.lower() for word in query.split()):
            return f"[Mechanism Doc]: {doc}"
    return "No relevant mechanism data found."

def clinical_retriever(query: str) -> str:
    for doc in clinical_trials_texts:
        if any(word.lower() in doc.lower() for word in query.split()):
            return f"[Clinical Trial Doc]: {doc}"
    return "No relevant clinical trial data found."

# Define tools using a simple wrapper function
def create_retriever_tool(retriever_func, tool_name: str, description: str):
    def tool(query: str):
        return retriever_func(query)
    # Mimic a tool message (in a real system, you would wrap this in a ToolMessage object)
    tool.__name__ = tool_name
    tool.description = description
    return tool

mechanism_tool = create_retriever_tool(
    mechanism_retriever,
    "mechanism_db_tool",
    "Search drug mechanism data for repurposing insights."
)
clinical_tool = create_retriever_tool(
    clinical_retriever,
    "clinical_db_tool",
    "Search clinical trial outcomes for repurposing evidence."
)
tools = [mechanism_tool, clinical_tool]

# Define the AgentState type for our workflow
class AgentState(TypedDict):
    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]

# Agent function: Classifies queries as targeting drug mechanisms or clinical outcomes
def agent(state: AgentState):
    print("---CALL AGENT---")
    messages = state["messages"]
    user_message = messages[0].content if not isinstance(messages[0], tuple) else messages[0][1]

    # Build a prompt to classify the query
    prompt = f"""Given the user question: "{user_message}"
If the question is about the molecular mechanism or pharmacodynamics, respond EXACTLY in this format:
SEARCH_MECHANISM: <search terms>

If it's about clinical trial outcomes, efficacy, or safety evidence, respond EXACTLY in this format:
SEARCH_CLINICAL: <search terms>

Otherwise, answer directly with general repurposing insights.
    """
    # (Here we simulate a call to DeepSeek-R1 using a dummy response)
    # In a real implementation, make an API call to DeepSeek-R1.
    if "mechanism" in user_message.lower() or "how it works" in user_message.lower():
        response_text = f"SEARCH_MECHANISM: {user_message}"
    elif "trial" in user_message.lower() or "efficacy" in user_message.lower() or "safety" in user_message.lower():
        response_text = f"SEARCH_CLINICAL: {user_message}"
    else:
        response_text = "The system did not classify your query. Please rephrase to focus on drug mechanism or clinical data."
    
    print("Agent response:", response_text)
    # Format response into expected tool call format
    if "SEARCH_MECHANISM:" in response_text:
        query = response_text.split("SEARCH_MECHANISM:")[1].strip()
        result = mechanism_tool(query)
        return {"messages": [AIMessage(content=f'Action: mechanism_db_tool\n{{"query": "{query}"}}\n\nResults: {result}')]}

    elif "SEARCH_CLINICAL:" in response_text:
        query = response_text.split("SEARCH_CLINICAL:")[1].strip()
        result = clinical_tool(query)
        return {"messages": [AIMessage(content=f'Action: clinical_db_tool\n{{"query": "{query}"}}\n\nResults: {result}')]}

    else:
        return {"messages": [AIMessage(content=response_text)]}

# Grading function: Checks if retrieved documents were found
def simple_grade_documents(state: AgentState):
    messages = state["messages"]
    last_message = messages[-1]
    print("Evaluating message:", last_message.content)
    if "Results:" in last_message.content and "No relevant" not in last_message.content:
        print("---DATA FOUND, PROCEED TO GENERATE INSIGHTS---")
        return "generate"
    else:
        print("---NO DATA FOUND, TRY REWRITE---")
        return "rewrite"

# Generate function: Synthesizes repurposing insights from retrieved data
def generate(state: AgentState):
    print("---GENERATE FINAL INSIGHTS---")
    messages = state["messages"]
    question = messages[0].content
    last_message = messages[-1]
    # Extract data from results
    data_start = last_message.content.find("Results:")
    retrieved_data = last_message.content[data_start:] if data_start != -1 else "No data available"
    # Build a prompt to synthesize insights
    prompt = f"""Based on the following retrieved data:
    {retrieved_data}
and considering the question:
    {question}
Summarize potential drug repurposing opportunities and any recommended next steps for further investigation.
    """
    # Dummy generation using a simple echo for demonstration.
    final_answer = f"Summary Insight: Considering the data, a promising repurposing opportunity is to explore Drug A for anti-inflammatory applications beyond its original use, and Drug B might be repurposed for neurodegenerative conditions. Further research should validate these hypotheses."
    print("Final Answer:", final_answer)
    return {"messages": [AIMessage(content=final_answer)]}

# Rewrite function: If no data is found, help rephrase the query for clarity
def rewrite(state: AgentState):
    print("---REWRITE QUESTION---")
    messages = state["messages"]
    original_question = messages[0].content if messages else "N/A"
    # Dummy rewrite that just appends "Please specify mechanism or trial data" for demonstration.
    rewritten = f"{original_question} (Please specify if you are asking about drug mechanism or clinical trial outcomes.)"
    print("Rewritten question:", rewritten)
    return {"messages": [AIMessage(content=rewritten)]}

# Decision function: Determines next step based on last message content
def custom_tools_condition(state: AgentState):
    messages = state["messages"]
    last_message = messages[-1]
    content = last_message.content
    if content.startswith("Action:"):
        print("Tool action detected. Proceed to retrieval.")
        return "tools"
    return END

# Create the workflow graph
workflow = StateGraph(AgentState)
workflow.add_node("agent", agent)
retrieve_node = ToolNode(tools)
workflow.add_node("retrieve", retrieve_node)
workflow.add_node("rewrite", rewrite)
workflow.add_node("generate", generate)

# Define workflow edges
workflow.add_edge(START, "agent")
workflow.add_conditional_edges("agent", custom_tools_condition, {"tools": "retrieve", END: END})
workflow.add_conditional_edges("retrieve", simple_grade_documents)
workflow.add_edge("generate", END)
workflow.add_edge("rewrite", "agent")
app = workflow.compile()

# Function to process a query through the workflow
def process_question(user_question: str, app, config: dict):
    events = []
    for event in app.stream({"messages": [("user", user_question)]}, config):
        events.append(event)
    return events

# Streamlit UI for the Drug Repurposing Advisor
def main():
    st.set_page_config(
        page_title="Drug Repurposing Advisor",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    st.title("💊 Drug Repurposing Advisor")
    st.markdown("### Explore potential drug repurposing opportunities with AI-driven insights.")
    query = st.text_area("Enter your research question:",
                          placeholder="e.g., Can Drug A be repurposed for neurodegenerative diseases?")
    col1, col2 = st.columns([1, 2])
    with col1:
        if st.button("🔍 Get Insights", use_container_width=True):
            if query:
                with st.spinner("Processing your query..."):
                    events = process_question(query, app, {"configurable": {"thread_id": "1"}})
                    for event in events:
                        if 'agent' in event:
                            with st.expander("Agent Processing Step", expanded=True):
                                content = event['agent']['messages'][0].content
                                st.markdown(f"**Agent Step Output:**\n\n{content}")
                        elif 'generate' in event:
                            st.markdown("### Final Insights:")
                            st.success(event['generate']['messages'][0].content)
                        elif 'rewrite' in event:
                            st.markdown("### Suggestion:")
                            st.warning(event['rewrite']['messages'][0].content)
            else:
                st.warning("⚠️ Please enter a query.")
    with col2:
        st.markdown("""
        **How to Use the Drug Repurposing Advisor:**
        1. **Input Query:** Describe your research question. Specify whether you are interested in drug mechanisms or clinical outcomes.
        2. **Get Insights:** Click "Get Insights" and let the system process your query.
        3. **Review Output:** Explore the retrieved data and the final synthesized insights.
        **Example Questions:**
        - "How does Drug A work and could its mechanism be useful in treating inflammatory diseases?"
        - "What are the clinical trial outcomes of Drug B and can it be repurposed for neurodegenerative conditions?"
        """)
if __name__ == "__main__":
    main()