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

# Load a lightweight, CPU-friendly model
model_id = "google/flan-t5-small"
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(history, user_input):
 #   base_prompt = (
  #      "You are Cinco, a helpful and friendly customer support assistant. "
   #     "Answer customer questions politely and clearly about product returns, refunds, and exchanges.\n\n"
    #)
    #for user, bot in history:
    #    base_prompt += f"Customer: {user}\nCinco Assistant: {bot}\n"
    #base_prompt += f"Customer: {user_input}\nCinco Assistant:"
    #return base_prompt

def format_prompt(history, user_input):
    policy_context = (
    "Cinco Return Policy Guidelines:\n"
    "- 🕒 **Returns Window**: Items may be returned **within 30 days** from the date of purchase.\n"
    "- 🧾 **Proof of Purchase**: A **valid receipt or order confirmation** is required for all returns and refunds.\n"
    "- 📦 **Item Condition**: Returned products must be in **original condition**, meaning **unused, unwashed, with tags still attached**.\n"
    "- 💳 **Refunds**: Refunds will be processed to the **original method of payment** (e.g., credit card, debit card, PayPal).\n"
    "- 🔁 **Exchanges**: Each item is eligible for **a one-time exchange only**, subject to availability.\n"
    "- 🚫 **Non-Returnable Items**: Items that have been **used, washed, or worn**, or returned **without a receipt**, **cannot be accepted** for return or exchange.\n"
    "- 🎁 **Gifts**: Gift returns require a **gift receipt** and will be refunded via **Cinco gift card**.\n"
    "- 🌐 **Online Orders**: For items bought online, customers can return via **mail** or in **any Cinco retail store** with a printed invoice.\n"
    "- ⌛ **Processing Time**: Please allow **5–7 business days** after the item is received for the refund to reflect.\n"
    "- 🛍️ **Final Sale Items**: Some products (e.g., clearance or promotional items) are marked as **final sale** and are **not eligible** for return or exchange.\n\n"
)


    intro = "You are Cinco Assistant, a helpful and polite customer support agent for the clothing brand Cinco.\n"
    
    # Include history if any
    formatted_history = ""
    if history:
        for i, (user_q, bot_a) in enumerate(history):
            formatted_history += f"Customer: {user_q}\nCinco Assistant: {bot_a}\n"

    current_question = f"Customer: {user_input}\nCinco Assistant:"

    # Combine everything
    prompt = (
        intro
        + policy_context
        + "Use the above return policy to answer the customer's question below.\n\n"
        + formatted_history
        + current_question
    )

    return prompt


# 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"]

    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, history  # return history twice: one for state, one for Chatbot UI

# Enhanced UI with better layout and visuals
with gr.Blocks(title="Cinco Returns Chatbot") as demo:
    gr.Markdown("""
    # 🧾 Cinco Returns Chatbot
    Welcome to the Cinco Returns Assistant. Ask your questions about returns, refunds, and exchanges.
    ### 💬 Example Questions:
    - *Can I return a sweater I no longer want?*
    - *What if I don’t have a receipt?*
    - *Can I exchange a used item?*
    """)

    with gr.Box():
        chatbot_ui = gr.Chatbot(label="Cinco Assistant", show_label=False, height=400)

        with gr.Row():
            user_input = gr.Textbox(
                placeholder="Type your question here (e.g., I want to return a sweater)...",
                scale=9,
                show_label=False,
                lines=2
            )
            submit_btn = gr.Button("Send", variant="primary", scale=1)

        state = gr.State([])

        submit_btn.click(
            fn=chat_fn,
            inputs=[user_input, state],
            outputs=[user_input, state, chatbot_ui]  # return updated chat history to both state and UI
        )
        user_input.submit(
            fn=chat_fn,
            inputs=[user_input, state],
            outputs=[user_input, state, chatbot_ui]
        )

    gr.Markdown("""
    ---
    🤖 Powered by [FLAN-T5](https://huggingface.co/google/flan-t5-base) on Hugging Face Spaces.
    """)

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