import gradio as gr
from askstein import ask

def main():
    with gr.Blocks() as demo:
        # Store user info and chat history in state
        user_info = gr.State({"first_name": None, "last_name": None})
        chat_history = gr.State([])  # list of (user_msg, bot_response, feedback)

        # --- User info input ---
        with gr.Row():
            first_name = gr.Textbox(label="First Name", placeholder="Enter your first name", interactive=True)
            last_name = gr.Textbox(label="Last Name", placeholder="Enter your last name", interactive=True)
        submit_names = gr.Button("Submit")

        # --- Chatbot interface (hidden initially) ---
        chatbot = gr.Chatbot(visible=False)
        user_msg = gr.Textbox(placeholder="Ask me anything...", visible=False)
        send_btn = gr.Button("Send", visible=False)
        feedback = gr.Radio(choices=["👍", "👎"], label="Your feedback", visible=False)
        submit_feedback = gr.Button("Submit Feedback", visible=False)
        clear_btn = gr.Button("Clear Chat", visible=False)

        # --- Functions ---

        def submit_user_info(fn, ln):
            fn = fn.strip()
            ln = ln.strip()
            if not fn or not ln:
                return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), "Please enter both first and last names."
            user_info.value = {"first_name": fn, "last_name": ln}
            return (
                gr.update(visible=False),  # hide first_name
                gr.update(visible=False),  # hide last_name
                gr.update(visible=True),   # show chatbot
                gr.update(visible=True),   # show user_msg
                gr.update(visible=True),   # show send_btn
                gr.update(visible=False),  # hide feedback (initially)
                gr.update(visible=False),  # hide submit_feedback
                gr.update(visible=True),   # show clear_btn
                ""
            )

        def send_message(message, history):
            if not message.strip():
                return history, "", gr.update(visible=False), gr.update(visible=True), gr.update(value=None)
            response = ask(message)
            history = history + [(message, response, None)]
            return history, "", gr.update(visible=True), gr.update(visible=False), gr.update(value=None)

        def submit_user_feedback(feedback_value, history):
            if not history:
                return history, gr.update(visible=False), "No conversation yet."
            if feedback_value is None:
                return history, gr.update(visible=True), "Please provide feedback."
            # Update last message feedback
            history[-1] = (history[-1][0], history[-1][1], feedback_value)
            return history, gr.update(visible=False), "Thank you for your feedback!"

        def clear_chat():
            return [], ""

        # --- Event bindings ---
        submit_names.click(
            submit_user_info,
            inputs=[first_name, last_name],
            outputs=[first_name, last_name, chatbot, user_msg, send_btn, feedback, submit_feedback, clear_btn, gr.Textbox(visible=True, interactive=False, label="Error Message")]
        )

        send_btn.click(
            send_message,
            inputs=[user_msg, chat_history],
            outputs=[chatbot, user_msg, feedback, submit_feedback, feedback]
        )

        submit_feedback.click(
            submit_user_feedback,
            inputs=[feedback, chat_history],
            outputs=[chatbot, feedback, gr.Textbox(visible=True, interactive=False, label="Feedback Status")]
        )

        clear_btn.click(
            clear_chat,
            outputs=[chatbot, user_msg]
        )

    demo.launch(server_name="0.0.0.0", server_port=7860)

if __name__ == "__main__":
    main()