import streamlit as st
from transformers import pipeline
import torch

# Set Streamlit page configuration
st.set_page_config(page_title="Brand Identity Generator", layout="wide")

# Load FLAN-T5 model pipeline
@st.cache_resource
def load_model():
    return pipeline("text2text-generation", model="google/flan-t5-base", device=-1)

generator = load_model()

# Custom Styling (simulate CSS)
st.markdown(
    """
    <style>
    .stTextInput>div>div>input {
        background-color: #e6f0fa;
        color: #003366;
        font-size: 16px;
    }
    .stButton>button {
        background-color: #4A90E2;
        color: white;
        border-radius: 8px;
        font-size: 16px;
        padding: 10px 16px;
    }
    .stButton>button:hover {
        background-color: #357ABD;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# Page Title and Description
st.markdown(
    "<h1 style='text-align: center; color: #4A90E2;'>🌟 Brand Identity Generator</h1>",
    unsafe_allow_html=True
)

st.markdown(
    """
    Generate **brand names**, **taglines**, **personality descriptions**, **colour palettes**, **target audiences**, and **slogans**
    for your startup idea using the powerful **FLAN-T5-BASE** model.
    """
)

# Form Input for Startup Idea
with st.form("input_form"):
    st.subheader("📌 Describe Your Startup")
    startup_idea = st.text_area("Startup Idea", placeholder="e.g. an eco-friendly fish delivery platform for urban areas")
    submitted = st.form_submit_button("Generate")

# On form submission
if submitted:
    if not startup_idea.strip():
        st.warning("Please enter a valid startup idea.")
    else:
        st.markdown("---")
        st.subheader("🔍 Generated Brand Identity")

        prompts = {
            "Brand Name Suggestions": f"Suggest 5 short, smart, brandable names for this startup. Avoid generic or repeated words. Idea: {startup_idea}",
            "Tagline": f"Give a single compelling tagline without repeating 'coaching', 'e-learning', or 'startup'. Use the idea: {startup_idea}",
            "Brand Personality Description": f"Describe the brand's personality in 1 creative sentence. Avoid clichés. Based on: {startup_idea}",
            "Suggested Brand Colours": f"Suggest a unique pair of primary and secondary brand colours with HEX codes for this startup: {startup_idea}",
            "Target Audience": f"Describe the target audience clearly and concisely (age, interests, needs) for: {startup_idea}",
            "Slogan Options": f"Give 3 short, catchy slogans that creatively reflect this idea: {startup_idea}. Avoid repetition or generic words."
        }

        # Generate and display each section
        for section, prompt in prompts.items():
            try:
                with st.spinner(f"Generating {section}..."):
                    result = generator(prompt, max_length=80, num_return_sequences=1)[0]["generated_text"]
                    st.markdown(f"### {section}")
                    st.success(result.strip())
            except Exception as e:
                st.markdown(f"### {section}")
                st.error("❌ Error generating content. Please check your connection or try again later.")