import streamlit as st
import requests
from bs4 import BeautifulSoup
import json
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load model and tokenizer
@st.cache_resource
def load_model():
    model = AutoModelForSeq2SeqLM.from_pretrained("shreyanshjha0709/watch-description-generator")
    tokenizer = AutoTokenizer.from_pretrained("shreyanshjha0709/watch-description-generator")
    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])))

# Web Scraping function to get product description from Ethos website
def scrape_ethos_description(product_link):
    try:
        response = requests.get(product_link)
        soup = BeautifulSoup(response.text, 'html.parser')

        # Assuming product description is within a specific class or ID, find it
        description = soup.find('div', class_='product-description')  # Example class, adjust based on actual site structure
        if description:
            return description.get_text(strip=True)
        else:
            return "No detailed description available on Ethos site."
    except Exception as e:
        return f"Error fetching details from Ethos site: {str(e)}"

# 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:
        # Display the image from the JSON
        image_url = watch_data.get("image", None)
        if image_url:
            st.image(image_url, caption=f"{watch_data['name']} Image")

        # Get the Ethos product link for web scraping (assuming 'url' key exists in your JSON)
        product_link = watch_data.get("url", None)
        if product_link:
            st.write(f"Fetching details from: [Product Page]({product_link})")

            # Scrape Ethos product page for more description
            ethos_description = scrape_ethos_description(product_link)
            st.write("### Ethos Product Description")
            st.write(ethos_description)
        else:
            st.warning("No Ethos link available for this SKU.")

        # Generate watch description based on attributes and scraped content
        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"),
            "water_resistance": watch_data.get("water_resistance", "Unknown Water Resistance"),
            "power_reserve": watch_data.get("power_reserve", "Unknown Power Reserve"),
            "dial_color": watch_data.get("dial_color", "Unknown Dial Color"),
            "strap_material": watch_data.get("strap_material", "Unknown Strap Material")
        }

        # Create a detailed description prompt combining scraped content and attributes
        input_text = f"""Generate a detailed 100-word description for the following watch:
Brand: {attributes['brand']}
Name: {attributes['name']}
SKU: {attributes['sku']}
Features: {attributes['features']}
Case Size: {attributes['casesize']}
Movement: {attributes['movement']}
Gender: {attributes['gender']}
Water Resistance: {attributes['water_resistance']}
Power Reserve: {attributes['power_reserve']}
Dial Color: {attributes['dial_color']}
Strap Material: {attributes['strap_material']}

Additional details from Ethos:
{ethos_description}

Description: Provide a luxurious, detailed description focusing on the craftsmanship, innovation, and design, similar to a high-end editorial style."""

        # Tokenize input and generate description
        inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
        outputs = model.generate(
            **inputs,
            max_length=150,  # Limit output to 100-150 words
            num_return_sequences=1,
            temperature=0.7,
            top_k=50,
            top_p=0.95,
            do_sample=True,
            repetition_penalty=1.2  # Prevent repetition in descriptions
        )

        # 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.")

# Add some information about the app
st.sidebar.title("About")
st.sidebar.info(
    "This app uses a fine-tuned AI model to generate descriptions for watches. "
    "Select a brand and a watch to get started. The model will generate a unique "
    "description based on the watch's attributes and additional details."
)

# Add a footer
st.markdown(
    """
    <style>
    .footer {
        position: fixed;
        left: 0;
        bottom: 0;
        width: 100%;
        background-color: #f1f1f1;
        color: black;
        text-align: center;
    }
    </style>
    <div class="footer">
        <p>Developed with ❤️ by Shreyansh Jha</p>
    </div>
    """,
    unsafe_allow_html=True
)