import gradio as gr
import os
from langchain_core.messages import HumanMessage    
import src.srf_bot as sb
import prompts.system_prompts as sp



# Initialize chatbot
chatbot = sb.SRFChatbot()

# Define the respond function
def respond(query, history):
    formatted_query = [HumanMessage(content=query)]
    # Invoke the graph with properly formatted input
    result = chatbot.graph.invoke({"messages": formatted_query}, chatbot.config)
    # Get the passages from the graph and append to history if documents exist
    state = chatbot.graph.get_state(config=chatbot.config).values
    documents = state.get("documents")
    passages = ''
    if documents and len(documents) > 0:
        for d in documents:
            passages += f'<b>{d.metadata["publication_name"]} - {d.metadata["chapter_name"]}</b>\n{d.page_content}\n\n'
        history.append((f'Passages: {query}', passages))
    # Extract the assistant's response and append to history
    response = result["messages"][-1].content
    system_message_dropdown = state.get("system_message_dropdown")
    history.append((query, f"<i>[{system_message_dropdown}]</i>\n" + response))
    return history


# Gradio interface with black and grey color scheme
with gr.Blocks(css="""
    .gradio-container {
        background-color: #F0F0F0; 
        font-family: 'Arial', sans-serif;
    }
    h1, h2, p {
        color: black;
    }
    h1 {
        font-size: 32px;
        text-align: left;
    }
    h2 {
        font-size: 24px;
    }
    p {
        font-size: 18px;
        margin-bottom: 15px;
    }
    .gr-button {
        background-color: #333333; 
        color: white;
        font-size: 18px;
        padding: 10px;
    }
    .gr-textbox textarea {
        font-size: 18px;
        color: black;
    }
    .gr-dropdown {
        font-size: 18px;
        color: black;
    }
    .source-box {
        background-color: white; 
        padding: 10px; 
        border-radius: 8px; 
        margin-top: 20px;
        color: black;
        border: 1px solid #D0D0D0;
    }
    @media (max-width: 600px) {
        .gr-row { flex-direction: column !important; }
        .gr-column { width: 100% !important; }
    }
""") as demo:
    
    # Title
    gr.Markdown("# SRF Chatbot")
    
    with gr.Row(elem_classes="gr-row"):
        with gr.Column(scale=4, elem_classes="gr-column"):
            # Chatbot interface
            chatbot_output = gr.Chatbot(height=600)  # Increased height for longer chat interface
            user_input = gr.Textbox(placeholder="Type your question here...", label="Your Question", value="What is the meaning of life?")
            submit_button = gr.Button("Submit")
        
        with gr.Column(scale=1, elem_classes="gr-column"):
            # Dropdown to select system prompts
            system_prompt_dropdown = gr.Dropdown(
                choices=list(sp.system_prompt_templates.keys()),
                label="Select Chatbot Instructions",
                value=list(sp.system_prompt_templates.keys())[0],
                elem_classes="gr-dropdown"
            )
            # Display the selected system prompt
            system_prompt_display = gr.Textbox(
                value=sp.system_prompt_templates[list(sp.system_prompt_templates.keys())[0]],
                label="Current Chatbot Instructions",
                lines=5,
                interactive=False
            )
            
            # Sources box (Now white, matching the other boxes)
            gr.Markdown("""
            <div class="source-box">
                <strong>Available sources:</strong>
                <ul>
                    <li>Journey to Self-Realization</li>
                    <li>The Second Coming of Christ</li>
                    <li>Autobiography of a Yogi</li>
                </ul>
            </div>
            """)
    
    # Update system prompt display when a new prompt is selected
    system_prompt_dropdown.change(
        fn=chatbot.reset_system_prompt,
        inputs=[system_prompt_dropdown],
        outputs=[system_prompt_display]
    )
    
    # Submit button logic to handle chatbot conversation
    submit_button.click(
        fn=respond,
        inputs=[user_input, chatbot_output],
        outputs=[chatbot_output]
    )

# Access the secrets
username = os.getenv("USERNAME")
password = os.getenv("PASSWORD")

# Launch the interface
demo.launch(share=True, auth=(username, password))