import streamlit as st
from paper import (
    literature_research_task, outline_task, draft_writing_task,
    citation_task, editing_task, chatbot_task,
    run_task
)

st.set_page_config(
    page_title="📘 Academic Research Assistant",
    layout="wide",
    initial_sidebar_state="expanded"
)

st.title("📘 Academic Research & Writing Assistant")
st.markdown(
    """
    Welcome to your AI-powered academic assistant! Choose a task from the sidebar,
    enter your topic and keywords, and let the agents do the rest.
    """
)

st.sidebar.header("🔧 Task Selector")
task_choice = st.sidebar.radio(
    "Select an academic task to perform:",
    [
        "Literature Review",
        "Generate Paper Outline",
        "Write Draft",
        "Generate Citations",
        "Edit Draft",
        "Ask Chatbot"
    ]
)

task_map = {
    "Literature Review": literature_research_task,
    "Generate Paper Outline": outline_task,
    "Write Draft": draft_writing_task,
    "Generate Citations": citation_task,
    "Edit Draft": editing_task,
    "Ask Chatbot": chatbot_task,
}

st.subheader("🔍 Input Details")

topic = st.text_input("Enter your research topic:", placeholder="e.g. Impact of AI in Healthcare")
keywords = st.text_input("Enter related keywords (comma-separated):", placeholder="e.g. AI, healthcare, diagnosis")

query = ""
if task_choice == "Ask Chatbot":
    query = st.text_area("Ask your question:", placeholder="e.g. How to format a research abstract?")

if st.button("🚀 Run Agent"):
    if not topic and task_choice != "Ask Chatbot":
        st.warning("Please enter a topic for the selected task.")
    elif task_choice == "Ask Chatbot" and not query:
        st.warning("Please enter a chatbot query.")
    else:
        with st.spinner("Agent is working on your task..."):
            # Format input text based on task type
            input_text = (
                query if task_choice == "Ask Chatbot"
                else f"Topic: {topic}\nKeywords: {keywords}"
            )
            result = run_task(task_map[task_choice], input_text)

        st.success("✅ Task Completed!")
        st.subheader("📄 Result")
        st.markdown(result, unsafe_allow_html=True)

st.markdown(
    """
    <hr>
    <center>
        <small>Built By Parthib</small>
    </center>
    """,
    unsafe_allow_html=True
)