import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline

# Load a lightweight, CPU-friendly model
model_id = "google/flan-t5-base"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id)

# Pipeline setup
chatbot = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

# Function to format prompt for chat-like interaction
def format_prompt(user_input):
    base_prompt = (
        "You are Cinco, a helpful assistant that answers customer questions ONLY about product returns, refunds, and exchanges.\n"
        "Respond concisely, clearly, and don't repeat the question. If the question is not about returns, politely say so.\n\n"
        f"Customer: {user_input}\n"
        f"Cinco Assistant:"
    )
    
# Chatbot logic
def chat_fn(user_input, history):
    history = history or []
    prompt = format_prompt(history, user_input)
    response = chatbot(prompt, max_length=256, do_sample=False, clean_up_tokenization_spaces=True)[0]["generated_text"]

    # Extract only the latest assistant response
    if "Cinco Assistant:" in response:
        assistant_reply = response.split("Cinco Assistant:")[-1].strip()
    else:
        assistant_reply = response.strip()

    history.append((user_input, assistant_reply))
    return "", history

# Build Gradio UI
with gr.Blocks(title="Cinco Returns Chatbot") as demo:
    gr.Markdown("## 🧾 Cinco Returns Chatbot\nAsk anything about returns, refunds, or exchanges.")
    chatbot_ui = gr.Chatbot(label="Cinco Assistant", show_label=True)
    with gr.Row():
        user_input = gr.Textbox(placeholder="Example: Can I return a used item without a receipt?", scale=6)
        submit_btn = gr.Button("Send", scale=1)

    state = gr.State([])

    submit_btn.click(fn=chat_fn, inputs=[user_input, state], outputs=[user_input, chatbot_ui])
    user_input.submit(fn=chat_fn, inputs=[user_input, state], outputs=[user_input, chatbot_ui])

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