import gradio as gr
import logging

from rag import generate_chat_answer

# ==============================
# Logging
# ==============================
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ==============================
# Clear Chat
# ==============================
def clear_chat():

    return "", [], []

# ==============================
# Chat Function
# ==============================
def respond(query, chat_history, session_history):

    if not query.strip():
        return "", chat_history, session_history

    if chat_history is None:
        chat_history = []

    if session_history is None:
        session_history = []

    answer = generate_chat_answer(
        query=query,
        history=session_history
    )

    # UI chat
    chat_history.append({
        "role": "user",
        "content": query
    })

    chat_history.append({
        "role": "assistant",
        "content": answer
    })

    # Session memory
    session_history.append({
        "role": "user",
        "content": query
    })

    session_history.append({
        "role": "assistant",
        "content": answer
    })

    logger.info(f"Session Updated:\n{session_history}")

    return "", chat_history, session_history

# ==============================
# CSS
# ==============================
custom_css = """
#container {
    background: white;
    border-radius: 15px;
    box-shadow: 0 4px 10px rgba(0,0,0,0.1);
    padding: 20px;
    margin: 20px auto;
    max-width: 900px;
}

#header {
    background: #4CAF50;
    padding: 20px;
    border-radius: 12px;
    margin-bottom: 20px;
    text-align: center;
    color: white;
}

#buttons-section {
    display: flex;
    gap: 10px;
    justify-content: center;
    margin-top: 20px;
}

#submit-btn {
    background: #4CAF50 !important;
    color: white !important;
    border: none !important;
}

#clear-btn {
    background: white !important;
    color: #4CAF50 !important;
    border: 1px solid #4CAF50 !important;
}
"""

# ==============================
# UI
# ==============================
with gr.Blocks(css=custom_css) as interface:

    with gr.Column(elem_id="container"):

        with gr.Column(elem_id="header"):

            gr.Markdown("""
# 🌿 Health & Nutrition Assistant
### Your AI Nutrition Companion
""")

        chatbot = gr.Chatbot(
            label="Chat History",
            elem_id="chatbot"
        )

        session_history = gr.State([])

        query = gr.Textbox(
            label="Ask your question",
            placeholder="Example: Best high protein Egyptian foods",
            lines=3
        )

        with gr.Row(elem_id="buttons-section"):

            submit_btn = gr.Button(
                "Send",
                elem_id="submit-btn"
            )

            clear_btn = gr.Button(
                "Clear Chat",
                elem_id="clear-btn"
            )

        submit_btn.click(
            fn=respond,
            inputs=[
                query,
                chatbot,
                session_history
            ],
            outputs=[
                query,
                chatbot,
                session_history
            ]
        )

        clear_btn.click(
            fn=clear_chat,
            outputs=[
                query,
                chatbot,
                session_history
            ]
        )

# ==============================
# Launch
# ==============================
if __name__ == "__main__":

    interface.launch()