import streamlit as st
from graph import EssayWriter, RouteQuery, GraphState
import os
import base64

st.set_page_config(page_title="Multi-Agent Essay Writer", page_icon="🤖🤖🤖✍️")

# Button HTML
button_html = f'''
    <div style="display: flex; justify-content: center;">
        <a href=" " target="_blank">
            <button style="
                background-color: #FFDD00;
                border: none;
                color: black;
                padding: 10px 20px;
                text-align: center;
                text-decoration: none;
                display: inline-block;
                font-size: 16px;
                margin: 4px 2px;
                cursor: pointer;
                border-radius: 10px;
                box-shadow: 2px 2px 10px rgba(0, 0, 0, 0.2);
            ">
            </button>
        </a>
    </div>
'''

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
    st.session_state.app = None
    st.session_state.chat_active = True

# Sidebar with instructions and API key input
with st.sidebar:
    st.subheader("About:")
    st.info(
        "\n\n* This app uses the 'gpt-4o-mini-2024-07-18' model."
        "\n\n* Writing essays may take some time, approximately 1-2 minutes."
    )
    
    openai_key = st.secrets.get("OPENAI_API_KEY", "")

    st.divider()

    # Reference section
    st.subheader("References:")
    st.markdown(
        '1. [Multi-Agent System with CrewAI and LangChain](https://discuss.streamlit.io/t/new-project-i-have-build-a-multi-agent-system-with-crewai-and-langchain/84002)',
        unsafe_allow_html=True
    )

# Initialize agents function
def initialize_agents():
    if not openai_key:
        st.error("OpenAI API key is missing! Please provide a valid key through Hugging Face Secrets.")
        st.session_state.chat_active = True
        return None

    os.environ["OPENAI_API_KEY"] = openai_key
    try:
        if "app" in st.session_state and st.session_state.app is not None:
            return st.session_state.app  # Prevent re-initialization
        
        essay_writer = EssayWriter().graph
        st.success("Agents successfully initialized!")
        st.session_state.chat_active = False
        return essay_writer
    except Exception as e:
        st.error(f"Error initializing agents: {e}")
        st.session_state.chat_active = True
        return None

# Automatically initialize agents on app load
if "app" not in st.session_state or st.session_state.app is None:
    st.session_state.app = initialize_agents()

app = st.session_state.app  

# Function to invoke the agent and generate a response
def generate_response(topic):
    return app.invoke(input={"topic": topic}) if app else {"response": "Error: Agents not initialized."}

# Display chat messages from the session
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"], unsafe_allow_html=True)

# Handle user input
if topic := st.chat_input(placeholder="Ask a question or provide an essay topic...", disabled=st.session_state.chat_active):
    st.chat_message("user").markdown(topic)
    st.session_state.messages.append({"role": "user", "content": topic})

    with st.spinner("Thinking..."):
        response = generate_response(topic)

    # Handle the assistant's response
    with st.chat_message("assistant"):
        if "essay" in response:  # Display essay preview and download link
            st.markdown("### Essay Preview:")
            st.markdown(f"#### {response['essay']['header']}")
            st.markdown(response["essay"]["entry"])
            for para in response["essay"]["paragraphs"]:
                st.markdown(f"**{para['sub_header']}**")
                st.markdown(para["paragraph"])
            st.markdown("**Conclusion:**")
            st.markdown(response["essay"]["conclusion"])

            # Provide download link for the PDF
            with open(response["pdf_name"], "rb") as pdf_file:
                b64 = base64.b64encode(pdf_file.read()).decode()
                href = f"<a href='data:application/octet-stream;base64,{b64}' download='{response['pdf_name']}'>Click here to download the PDF</a>"
                st.markdown(href, unsafe_allow_html=True)

            # Save the assistant's message to session state
            st.session_state.messages.append(
                {"role": "assistant", "content": "Here is your essay preview and the download link."}
            )
        else:  # For other responses (e.g., general answers)
            st.markdown(response["response"])
            st.session_state.messages.append({"role": "assistant", "content": response["response"]})

# Add acknowledgment at the bottom
st.markdown("---")
st.markdown(
    """
    <div style="text-align: center; font-size: 14px; color: #555;">
        <strong>Acknowledgment:</strong> This project is based on Mesut Duman's work: 
        <a href="https://github.com/mesutdmn/Autonomous-Multi-Agent-Systems-with-CrewAI-Essay-Writer/tree/main" 
           target="_blank" style="color: #007BFF; text-decoration: none;">
           CrewAI Essay Writer
        </a>
    </div>
    """,
    unsafe_allow_html=True,
)