import streamlit as st
import requests
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Load model and tokenizer
@st.cache_resource
def load_model():
    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
    return model, tokenizer

model, tokenizer = load_model()

# Load the JSON file from a URL
@st.cache_data
def load_json_from_url(url):
    response = requests.get(url)
    return response.json()

# Provide your JSON URL here
json_url = "https://www.ethoswatches.com/feeds/holbox_ai.json"
data = load_json_from_url(json_url)

# Extract unique brands
brands = sorted(list(set([item["brand"] for item in data])))

# Streamlit UI
st.title("Watch Description Generator")

# Select brand
selected_brand = st.selectbox("Select a Brand", ["Select"] + brands)

# Filter watches and SKUs by the selected brand
if selected_brand != "Select":
    watches = [item["name"] for item in data if item["brand"] == selected_brand]
    skus = [item["sku"] for item in data if item["brand"] == selected_brand]

    selected_watch = st.selectbox("Select Watch Name (Optional)", ["Select"] + watches)
    selected_sku = st.selectbox("Select SKU (Optional)", ["Select"] + skus)

    # Get the selected watch data from the JSON
    watch_data = None
    if selected_watch != "Select":
        watch_data = next((item for item in data if item["name"] == selected_watch), None)
    elif selected_sku != "Select":
        watch_data = next((item for item in data if item["sku"] == selected_sku), None)

    if watch_data:
        # Generate description based on attributes
        if st.button("Generate Description"):
            attributes = {
                "brand": watch_data["brand"],
                "name": watch_data.get("name", "Unknown Watch"),
                "sku": watch_data.get("sku", "Unknown SKU"),
                "features": watch_data.get("features", "Unknown Features"),
                "casesize": watch_data.get("casesize", "Unknown Case Size"),
                "movement": watch_data.get("movement", "Unknown Movement"),
                "gender": watch_data.get("gender", "Unknown Gender"),
                # Add more attributes as needed
            }

            input_text = f"Brand: {attributes['brand']}, Watch Name: {attributes['name']}, SKU: {attributes['sku']}, Features: {attributes['features']}, Case Size: {attributes['casesize']}, Movement: {attributes['movement']}, Gender: {attributes['gender']}"
            
            # Tokenize input and generate description
            inputs = tokenizer(input_text, return_tensors="pt")
            outputs = model.generate(**inputs)
            
            # Decode generated text
            description = tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # Display the result
            st.write("### Generated Description")
            st.write(description)
else:
    st.warning("Please select a brand.")