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

# Load the pretrained model and tokenizer
MODEL_NAME = "google/flan-t5-large"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

# Function to generate text description
def generate_description(input_text):
    inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
    outputs = model.generate(inputs["input_ids"], max_length=200)
    description = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return description

# Function to scrape brand website
def scrape_brand_website(brand_name, sku):
    search_url = f"https://www.google.com/search?q={brand_name}+{sku}"
    headers = {
        "User-Agent": "Mozilla/5.0"
    }
    response = requests.get(search_url, headers=headers)
    soup = BeautifulSoup(response.text, 'html.parser')
    
    # For demonstration, just fetch the first URL (modify as per your requirement)
    link = soup.find('a')['href']
    return link

# Streamlit App
st.title("Watch Description Generator & Brand Scout")

# Inputs for watch details
brand_name = st.text_input("Enter Brand Name")
watch_name = st.text_input("Enter Watch Name")
sku = st.text_input("Enter Watch SKU")

# Button to generate description
if st.button("Generate Description"):
    if brand_name and watch_name and sku:
        # Generate description using the model
        input_text = f"Watch Name: {watch_name}, Brand: {brand_name}, SKU: {sku}"
        generated_description = generate_description(input_text)
        st.write("Generated Description:")
        st.write(generated_description)

        # Scrape brand website for relevant information
        st.write("Fetching Brand Website...")
        brand_link = scrape_brand_website(brand_name, sku)
        st.write(f"Brand Website: {brand_link}")
    else:
        st.warning("Please fill in all fields.")