import gradio as gr
from transformers import pipeline
import functools

# Cache the pipeline so it loads only once
@functools.lru_cache(maxsize=1)
def load_pipeline():
    # Using google/flan-t5-base which is instruction-tuned for improved responses.
    return pipeline("text2text-generation", model="google/flan-t5-base", device=-1)

def dental_chatbot_response(message, history):
    """
    Generates responses with dental expertise using google/flan-t5-base.
    """
    prompt = (
        "You are a highly knowledgeable and friendly dental expert chatbot. "
        "Provide detailed and accurate explanations of dental terms, procedures, and treatments. "
        "If the query is not dental-related, respond helpfully and informatively. "
        f"User: {message}"
    )
    
    gen_pipe = load_pipeline()
    generated = gen_pipe(
        prompt,
        max_length=200,
        do_sample=True,
        top_p=0.9,
        top_k=50,
    )
    ai_response = generated[0]["generated_text"].strip()
    return ai_response

# Custom CSS for a colorful, modern look with icons and visual effects.
custom_css = """
/* Overall background with a pleasant gradient */
body {
    background: linear-gradient(135deg, #ff9a9e, #fad0c4);
    font-family: Arial, sans-serif;
    margin: 0;
    padding: 0;
}

/* Main container styling */
.container {
    background: #ffffff;
    border-radius: 20px;
    box-shadow: 0 8px 16px rgba(0,0,0,0.2);
    max-width: 800px;
    margin: 40px auto;
    padding: 20px;
}

/* Header styling */
.header {
    text-align: center;
    margin-bottom: 20px;
}
.header h1 {
    font-size: 2.5rem;
    margin: 0;
    color: #d81b60;
}
.header p {
    font-size: 1.2rem;
    color: #555555;
}

/* Chat message styling */
.chat-msg {
    display: flex;
    margin: 10px 0;
    align-items: flex-start;
}
.chat-msg.user {
    justify-content: flex-start;
}
.chat-msg.bot {
    justify-content: flex-end;
}
.msg-content {
    max-width: 70%;
    padding: 10px 15px;
    border-radius: 15px;
    position: relative;
    font-size: 1rem;
    line-height: 1.3;
}
/* User message bubble */
.chat-msg.user .msg-content {
    background: #ffe0b2;
    color: #5d4037;
    margin-left: 40px;
}
/* Bot message bubble */
.chat-msg.bot .msg-content {
    background: #c5e1a5;
    color: #33691e;
    margin-right: 40px;
}

/* Icon styling for messages (using emojis) */
.msg-icon {
    font-size: 1.5rem;
    width: 30px;
    text-align: center;
}
.chat-msg.user .msg-icon {
    color: #d81b60;
}
.chat-msg.bot .msg-icon {
    color: #2e7d32;
}

/* Input area styling */
#input-area {
    display: flex;
    margin-top: 20px;
}
#input-area textarea {
    flex: 1;
    padding: 10px;
    font-size: 1rem;
    border: 2px solid #ffe0b2;
    border-radius: 10px;
    outline: none;
}
#input-area button {
    background: #d81b60;
    color: #fff;
    border: none;
    padding: 0 20px;
    margin-left: 10px;
    border-radius: 10px;
    cursor: pointer;
    transition: background 0.3s ease;
    font-size: 1rem;
}
#input-area button:hover {
    background: #ff5252;
}
"""

# Build the custom Blocks interface using gr.Column as our container.
with gr.Blocks(css=custom_css) as demo:
    with gr.Column(elem_classes=["container"]):
        # Header with a dental icon in the title
        with gr.Column(elem_classes=["header"]):
            gr.Markdown("## 🦷 Advanced Dental Terminology Chatbot")
            gr.Markdown("Ask me anything about dental terms, procedures, and treatments! This chatbot is powered by an instruction-tuned LLM for accurate and detailed answers.")
        
        # Chat display area
        chatbot = gr.HTML("<div id='chat-box'></div>")
        chat_state = gr.State([])  # To store chat history as a list of (user, bot) pairs

        # Input area
        with gr.Row(elem_id="input-area"):
            user_input = gr.Textbox(show_label=False, placeholder="Type your message here...", lines=2)
            send_btn = gr.Button("Send")

    def update_chat(user_message, history):
        # Append the user's message to the chat history.
        history = history + [(user_message, None)]
        # Generate bot response.
        bot_reply = dental_chatbot_response(user_message, history)
        history[-1] = (user_message, bot_reply)
        # Build HTML for the updated chat box with icons.
        chat_html = ""
        for user_msg, bot_msg in history:
            # User message bubble with a smiley icon.
            chat_html += f"""
            <div class='chat-msg user'>
                <div class='msg-icon'>🙂</div>
                <div class='msg-content'>{user_msg}</div>
            </div>
            """
            if bot_msg:
                # Bot message bubble with a dental icon.
                chat_html += f"""
                <div class='chat-msg bot'>
                    <div class='msg-content'>{bot_msg}</div>
                    <div class='msg-icon'>🦷</div>
                </div>
                """
        return chat_html, history, ""

    # Wire up interactions: when the textbox is submitted or the send button is clicked.
    user_input.submit(
        update_chat,
        inputs=[user_input, chat_state],
        outputs=[chatbot, chat_state, user_input]
    )
    send_btn.click(
        update_chat,
        inputs=[user_input, chat_state],
        outputs=[chatbot, chat_state, user_input]
    )

if __name__ == "__main__":
    demo.launch()