# app.py
# GPT Chatbot with thinking mode
# Users enter their HF_API_TOKEN in the UI
# Ready for Hugging Face Spaces

import time
import gradio as gr
from huggingface_hub import InferenceClient

# -----------------------
# Configuration
# -----------------------
GPT_MODEL_ID = "HuggingFaceH4/zephyr-7b-beta"  # default model, can change

# -----------------------
# Chatbot Core
# -----------------------
def gpt_chat(prompt, history, thinking, thinking_time, hf_token):
    """
    Generate a chat response using Hugging Face Inference API.
    Users provide their own HF_API_TOKEN.
    """
    if not hf_token or hf_token.strip() == "":
        return history, "❌ Please enter your Hugging Face API Token."

    if not prompt.strip():
        return history, "Please write a message."

    # Initialize client
    try:
        client = InferenceClient(model=GPT_MODEL_ID, token=hf_token)
    except Exception as e:
        return history, f"Error creating HF client: {str(e)}"

    # Append user message to history
    history = history + [(prompt, None)]

    # Simulate thinking
    if thinking:
        time.sleep(max(0.0, thinking_time))

    try:
        # Build conversation prompt
        conversation_text = ""
        for user_msg, bot_msg in history:
            if bot_msg is not None:
                conversation_text += f"User: {user_msg}\nAssistant: {bot_msg}\n"
        conversation_text += f"User: {prompt}\nAssistant:"

        # Generate response
        response = client.text_generation(
            prompt=conversation_text,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True
        )

        answer = response.replace(conversation_text, "").strip()
        if "Assistant:" in answer:
            answer = answer.split("Assistant:")[-1].strip()

        # Update last entry in history
        history = history[:-1] + [(prompt, answer)]
        return history, "OK"

    except Exception as e:
        history = history[:-1] + [(prompt, "⚠ Error generating response.")]
        return history, str(e)

# -----------------------
# Gradio UI
# -----------------------
with gr.Blocks(title="ThinkPad - Chatbot") as demo:
    gr.Markdown("""
    # 🤖 ThinkPad Chatbot
    Enter your Hugging Face API Token below to start chatting.
    Built with ❤️ by [JumpLander](https://jumplander.org)
    """)

    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="Chat")
            user_input = gr.Textbox(show_label=False, placeholder="Write a message and press Send")
            hf_token_input = gr.Textbox(
                label="Hugging Face API Token",
                placeholder="Paste your HF_API_TOKEN here",
                type="password"
            )
            thinking_mode = gr.Checkbox(label="Thinking mode", value=True)
            thinking_time = gr.Slider(0, 5, value=2.0, step=0.5, label="Thinking time (seconds)")
            status_box = gr.Textbox(label="Status")
            send_btn = gr.Button("Send")

    # handle send
    def respond(history, msg, thinking, thinking_t, hf_token):
        return gpt_chat(msg, history or [], thinking, thinking_t, hf_token)

    send_btn.click(respond, inputs=[chatbot, user_input, thinking_mode, thinking_time, hf_token_input], outputs=[chatbot, status_box])
    user_input.submit(respond, inputs=[chatbot, user_input, thinking_mode, thinking_time, hf_token_input], outputs=[chatbot, status_box])

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