import streamlit as st
import requests
from bs4 import BeautifulSoup
from transformers import pipeline
import torch

# --- App Configuration ---
st.set_page_config(
    page_title="CodeStratLabs Tagline Generator",
    page_icon="🚀",
    layout="centered"
)

# --- Model Loading (with caching) ---
# Cache the model loading to prevent reloading on every interaction.
@st.cache_resource
def load_generator_model():
    """Loads the text generation model from Hugging Face."""
    # We use a smaller, faster model perfect for this kind of creative task.
    # It's optimized for instruction-following.
    return pipeline(
        "text-generation",
        model="HuggingFaceH4/zephyr-7b-beta",
        torch_dtype=torch.bfloat16,
        device_map="auto" # Automatically uses GPU if available
    )

generator = load_generator_model()

# --- Web Scraping Function ---
@st.cache_data(ttl=3600) # Cache the scraped data for 1 hour
def get_codestrat_about_text():
    """
    Scrapes the 'About Us' text from the CodeStratLabs website.
    Returns the text content or an error message.
    """
    url = "https://www.codestratlabs.com/#about"
    try:
        response = requests.get(url, headers={'User-Agent': 'Mozilla/5.0'})
        response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)

        soup = BeautifulSoup(response.content, 'html.parser')

        # Find the specific 'about' section by its ID
        about_section = soup.find('section', id='about')

        if about_section:
            # Extract text from all paragraphs within the section for cleaner output
            paragraphs = about_section.find_all('p')
            full_text = ' '.join(p.get_text(strip=True) for p in paragraphs)
            return full_text
        else:
            return "Error: Could not find the '#about' section on the page."

    except requests.exceptions.RequestException as e:
        return f"Error: Could not fetch the website content. Details: {e}"

# --- Main App UI ---

st.title("🚀 AI Tagline Generator for CodeStratLabs")
st.markdown(
    "This app scrapes the live 'About Us' page from the CodeStratLabs website "
    "and uses a Hugging Face LLM to generate creative marketing taglines."
)

# Fetch the text when the app loads
with st.spinner("Scraping the CodeStratLabs 'About Us' page..."):
    about_text = get_codestrat_about_text()

if "Error:" in about_text:
    st.error(about_text)
else:
    st.success("Successfully scraped the website content!")
    with st.expander("Click to see the scraped text"):
        st.write(about_text)

    st.markdown("---")

    # --- Tagline Generation ---
    st.header("Generate Marketing Taglines")
    
    # User can customize the number of taglines
    num_taglines = st.slider("Select number of taglines to generate:", min_value=3, max_value=10, value=5)

    if st.button("✨ Generate Now"):
        with st.spinner("🤖 The AI is thinking... Please wait."):
            # We create a specific, instruction-based prompt for the model.
            # This technique is called "prompt engineering".
            prompt = f"""
            <|system|>
            You are a world-class marketing expert specializing in catchy, professional, and impactful taglines for technology companies. Your goal is to generate taglines based on the company description provided. The taglines should be short, memorable, and reflect the company's focus on AI, innovation, and business impact.</s>
            <|user|>
            Here is the company description for CodeStratLabs:
            "{about_text}"

            Please generate {num_taglines} unique marketing taglines based on this description.</s>
            <|assistant|>
            """

            try:
                # Generate the text using the pipeline
                sequences = generator(
                    prompt,
                    max_new_tokens=150, # Max length of the generated response
                    do_sample=True,
                    temperature=0.7,
                    top_k=50,
                    top_p=0.95,
                    num_return_sequences=1,
                )

                # The model's output is a single string, so we need to parse it.
                if sequences and sequences[0]['generated_text']:
                    # The response starts with the prompt, so we split it to get only the assistant's part.
                    assistant_response = sequences[0]['generated_text'].split("<|assistant|>")[1].strip()
                    
                    # Clean up the output and split into a list
                    all_taglines = [line.strip() for line in assistant_response.split('\n') if line.strip()]
                    
                    # *** THIS IS THE FIX ***
                    # We slice the list to respect the user's choice.
                    limited_taglines = all_taglines[:num_taglines]

                    st.subheader("Here are your AI-generated taglines:")
                    for i, tagline in enumerate(limited_taglines):
                        # Remove potential numbering like "1. " from the model's output
                        clean_tagline = tagline.split('. ', 1)[-1]
                        st.markdown(f"> {i+1}. **{clean_tagline}**")
                else:
                    st.error("The model did not return a valid response. Please try again.")

            except Exception as e:
                st.error(f"An error occurred during AI generation: {e}")

st.markdown("---")
st.info("Created by Parul Kumari. Built as a demonstration project for the Generative AI Engineer application at CodeStratLabs.")