import gradio as gr
from answer import answer_question


# ---------- Helper: format retrieved chunks ----------
def format_chunks(chunks):
    if not chunks:
        return "No knowledge-base retrieval was required for this question."

    formatted = []
    for i, doc in enumerate(chunks, 1):
        source = doc.metadata.get("source", "unknown")
        text = doc.page_content.strip()

        formatted.append(
            f"### Chunk {i}\n"
            f"**Source:** {source}\n\n"
            f"{text}"
        )

    return "\n\n---\n\n".join(formatted)


# ---------- Chat handler (tuple-based, HF compatible) ----------
def chat_handler(message, chat_history):
    chat_history = chat_history or []

    # Build history text for RAG backend
    history_text = ""
    for msg in chat_history:
        history_text += f"{msg['role'].capitalize()}: {msg['content']}\n"

    # Call RAG backend
    answer, chunks = answer_question(message, history_text)

    # Append messages as DICTS (REQUIRED)
    chat_history.append({"role": "user", "content": message})
    chat_history.append({"role": "assistant", "content": answer})

    chunk_display = format_chunks(chunks)

    return chat_history, chunk_display, ""


# ---------- Gradio UI ----------
with gr.Blocks(title="BlinkNow – DSA RAG Assistant") as demo:

    gr.Markdown(
        """
        # 📘 BlinkNow – Data Structures & Algorithms Assistant
        Ask a DSA question and see the **retrieved knowledge chunks**
        used to generate the answer.
        """
    )

    with gr.Row():
        # Left: Chat
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="Chat",
                height=500
            )

            with gr.Row():
                user_input = gr.Textbox(
                    placeholder="Ask a DSA question...",
                    label="Your Question",
                    scale=4
                )
                send_btn = gr.Button("Send", variant="primary", scale=1)

        # Right: Retrieved Context
        with gr.Column(scale=3):
            gr.Markdown("### 🔍 Retrieved Context")
            retrieved_chunks = gr.Markdown(
                value="*Retrieved knowledge chunks will appear here after your query.*"
            )

    # ---------- Events ----------
    send_btn.click(
        fn=chat_handler,
        inputs=[user_input, chatbot],
        outputs=[chatbot, retrieved_chunks, user_input]
    )

    user_input.submit(
        fn=chat_handler,
        inputs=[user_input, chatbot],
        outputs=[chatbot, retrieved_chunks, user_input]
    )


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