import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import random

# Initialize GPT-2 model and tokenizer
model_name = "gpt2"
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)

# Define some sample data for the chatbot
products = {
    'T-shirt': {'sizes': ['S', 'M', 'L', 'XL'], 'colors': ['Red', 'Blue', 'Green'], 'price': 25},
    'Jacket': {'sizes': ['M', 'L', 'XL'], 'colors': ['Black', 'Grey'], 'price': 50},
    'Sneakers': {'sizes': ['7', '8', '9', '10'], 'colors': ['White', 'Black'], 'price': 60}
}

# Simulate order tracking
orders = {
    '12345': {'status': 'Shipped', 'arrival_date': '2024-12-22'},
    '67890': {'status': 'In Transit', 'arrival_date': '2024-12-23'}
}

# Simulate some discounts
discounts = {
    'WINTER20': '20% off on winter apparel!',
    'WELCOME10': '10% off for first-time customers!'
}

# Function to generate response using GPT-2
def generate_gpt2_response(user_input):
    inputs = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
    outputs = model.generate(inputs, max_length=150, num_return_sequences=1, no_repeat_ngram_size=2, temperature=0.7)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Chatbot logic
def chatbot():
    st.title("E-Commerce Chatbot 🛍️")
    st.write("Hi, I'm your shopping assistant! How can I help you today?")

    user_input = st.text_input("You: ", "")

    if user_input:
        # Handle product queries
        if "size" in user_input.lower() or "color" in user_input.lower():
            st.write("Sure! Here are the available options:")
            for product, details in products.items():
                st.write(f"**{product}**: Sizes: {', '.join(details['sizes'])}, Colors: {', '.join(details['colors'])}, Price: ${details['price']}")
        
        # Handle order tracking
        elif "track order" in user_input.lower():
            order_number = st.text_input("Please enter your order number:", "")
            if order_number:
                if order_number in orders:
                    order_info = orders[order_number]
                    st.write(f"Your order is {order_info['status']}. It will arrive by {order_info['arrival_date']}.")
                else:
                    st.write("Sorry, we couldn't find that order. Please check your order number.")
        
        # Handle product recommendations
        elif "recommend" in user_input.lower():
            st.write("Here are some products I recommend based on your browsing:")
            recommended = random.sample(list(products.keys()), 2)
            for product in recommended:
                st.write(f"- {product}")
        
        # Handle discounts
        elif "discount" in user_input.lower():
            discount_code = st.text_input("Enter your discount code:", "")
            if discount_code and discount_code in discounts:
                st.write(f"Great! You can use the code {discount_code} for: {discounts[discount_code]}")
            elif discount_code:
                st.write("Sorry, that discount code is invalid.")
        
        # Advanced conversation via GPT-2
        else:
            gpt2_response = generate_gpt2_response(user_input)
            st.write(f"Bot: {gpt2_response}")
    
    else:
        st.write("Start by asking about products, tracking your order, or getting discounts!")

# Run the chatbot function
if __name__ == '__main__':
    chatbot()