import gradio as gr
from transformers import pipeline

# Initialize the conversational model pipeline
chatbot_pipeline = pipeline("text-generation", model="Aditya0619/Medbot")

# Function to manage history and generate responses
def respond(message, history, system_message, max_tokens, temperature, top_p):
    # Initialize history if it's None
    if history is None:
        history = []

    # Build input by concatenating past messages
    chat_input = ""
    for user_input, bot_response in history:
        chat_input += f"User: {user_input}\nBot: {bot_response}\n"
    chat_input += f"User: {message}\nBot:"

    # Generate a response using the pipeline
    response = chatbot_pipeline(
        chat_input,
        max_length=max_tokens,
        temperature=temperature,
        top_p=top_p,
        pad_token_id=50256  # Avoid padding issues for models like GPT-2 variants
    )[0]["generated_text"].split("Bot:")[-1].strip()

    # Update the conversation history
    history.append((message, response))

    # Return the updated chat history
    return history, history

# Define the Gradio app layout
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 AI Chatbot with Memory\nChat with me! I’ll remember your messages.")

    # Taskbar with configurable parameters
    with gr.Row():
        with gr.Accordion("⚙️ Configure Chatbot Settings", open=False):
            system_message = gr.Textbox(
                label="System Message (Optional)",
                placeholder="e.g., You are a helpful assistant."
            )
            max_tokens = gr.Slider(
                label="Max Tokens", minimum=50, maximum=500, value=250, step=10
            )
            temperature = gr.Slider(
                label="Temperature", minimum=0.0, maximum=1.0, value=0.7, step=0.1
            )
            top_p = gr.Slider(
                label="Top P", minimum=0.0, maximum=1.0, value=0.9, step=0.1
            )

    # Chatbot and user input section
    chatbot = gr.Chatbot(label="Chat with AI")
    user_input = gr.Textbox(label="Your Message", placeholder="Type a message...", lines=2)

    # Hidden state to store the conversation history
    state = gr.State([])

    # Submit button
    submit = gr.Button("Send")

    # Connect user input to the chatbot response function
    submit.click(
        respond,
        inputs=[user_input, state, system_message, max_tokens, temperature, top_p],
        outputs=[chatbot, state]
    )

    # Display an initial greeting message
    demo.load(lambda: [("Hi! How can I assist you today?", "")], outputs=chatbot)

# Launch the Gradio app
demo.launch()