app.py

import gradio as gr
import numpy as np
import os
from huggingface_hub import login
from sentence_transformers import SentenceTransformer, util

# --- CONFIGURATION ---
class Config:
    """Configuration settings for the application."""
    EMBEDDING_MODEL_ID = "google/embeddinggemma-300M"
    PROMPT_NAME = "STS"
    TOP_K = 5
    HF_TOKEN = os.getenv('HF_TOKEN')

# --- FONT DATA ---
FONT_DATA = [
    {
        "name": "Playfair Display",
        "family": "serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Playfair+Display:wght@400;700&display=swap",
        "description": "Elegant, sophisticated, editorial, high-contrast serif with dramatic flair, perfect for luxury brands and fashion magazines"
    },
    {
        "name": "Inter",
        "family": "sans-serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600&display=swap",
        "description": "Modern, clean, professional, highly legible sans-serif designed for digital interfaces and contemporary design"
    },
    {
        "name": "Amatic SC",
        "family": "handwriting",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Amatic+SC:wght@400;700&display=swap",
        "description": "Playful, casual, handwritten, fun, child-like, informal font perfect for creative and whimsical projects"
    },
    {
        "name": "Crimson Text",
        "family": "serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Crimson+Text:wght@400;600&display=swap",
        "description": "Classical, scholarly, academic, readable serif inspired by old-style typefaces, ideal for books and literature"
    },
    {
        "name": "Roboto",
        "family": "sans-serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;500&display=swap",
        "description": "Friendly, approachable, geometric sans-serif with a mechanical skeleton, widely used in digital applications"
    },
    {
        "name": "Dancing Script",
        "family": "script",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Dancing+Script:wght@400;700&display=swap",
        "description": "Romantic, flowing, elegant script font perfect for wedding invitations, greeting cards, and feminine designs"
    },
    {
        "name": "Oswald",
        "family": "sans-serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Oswald:wght@300;400;600&display=swap",
        "description": "Bold, condensed, impactful sans-serif with strong presence, ideal for headlines and masculine designs"
    },
    {
        "name": "Lora",
        "family": "serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Lora:wght@400;600&display=swap",
        "description": "Warm, friendly, contemporary serif with calligraphic roots, perfect for body text and storytelling"
    },
    {
        "name": "Montserrat",
        "family": "sans-serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Montserrat:wght@300;400;600&display=swap",
        "description": "Urban, modern, versatile sans-serif inspired by Buenos Aires signage, great for branding and corporate use"
    },
    {
        "name": "Pacifico",
        "family": "script",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Pacifico&display=swap",
        "description": "Surfing, California, retro, casual script font with beach vibes and laid-back summer feeling"
    },
]

# --- CORE LOGIC ---
class FontMoodGenerator:
    """Handles model loading, embedding generation, and font palette creation."""

    def __init__(self, config: Config, font_data: list[dict[str, any]]):
        """Initializes the generator, logs in, and loads necessary assets."""
        self.config = config
        self.font_data = font_data
        self._login_to_hf()
        self.embedding_model = self._load_model()
        self.font_embeddings = self._precompute_font_embeddings()

    def _login_to_hf(self):
        """Logs into Hugging Face Hub if a token is provided."""
        if self.config.HF_TOKEN:
            print("Logging into Hugging Face Hub...")
            login(token=self.config.HF_TOKEN)
        else:
            print("HF_TOKEN not found. Proceeding without login.")

    def _load_model(self) -> SentenceTransformer:
        """Loads the Sentence Transformer model."""
        print(f"Initializing embedding model: {self.config.EMBEDDING_MODEL_ID}...")
        try:
            return SentenceTransformer(self.config.EMBEDDING_MODEL_ID)
        except Exception as e:
            print(f"Error loading model: {e}")
            raise

    def _precompute_font_embeddings(self) -> np.ndarray:
        """Generates and stores embeddings for the font descriptions."""
        print("Pre-computing embeddings for font palette...")
        font_texts = [
            f"{font['name']}, {font['description']}"
            for font in self.font_data
        ]
        embeddings = self.embedding_model.encode(
            font_texts,
            prompt_name=self.config.PROMPT_NAME
        )
        print("Embeddings computed successfully.")
        return embeddings

    def _format_palette_as_html(self, top_hits: list[dict[str, any]]) -> str:
        """Formats the top font hits into a displayable HTML string."""
        if not top_hits:
            return "<p>Could not generate a font palette. Please try another mood.</p>"

        sample_texts = [
            "The Quick Brown Fox Jumps Over The Lazy Dog",
            "Sphinx of black quartz, judge my vow",
            "How vexingly quick daft zebras jump!",
            "Pack my box with five dozen liquor jugs",
            "Waltz, bad nymph, for quick jigs vex"
        ]

        cards_html = ""
        for i, hit in enumerate(top_hits):
            font_info = self.font_data[hit['corpus_id']]
            font_name = font_info['name']
            font_family = font_info['family']
            score = hit['score']
            sample_text = sample_texts[i % len(sample_texts)]

            cards_html += f"""
            <div class="font-card">
                <div class="font-header">
                    <h3>{font_name}</h3>
                    <span class="font-score">Score: {score:.2f}</span>
                </div>
                <div class="font-sample" style="font-family: '{font_name}', {font_family};">
                    {sample_text}
                </div>
                <div class="font-details">
                    <span class="font-family">{font_family.title()}</span>
                    <span class="font-description">{font_info['description'][:100]}...</span>
                </div>
            </div>
            """
        return f"<div class='font-palette-container'>{cards_html}</div>"

    def generate_palette(self, mood_text: str) -> tuple[str, list[dict[str, any]]]:
        """Generates font palette and returns both HTML and raw data."""
        if not mood_text or not mood_text.strip():
            return "<p>Please enter a mood or a description.</p>", []

        mood_embedding = self.embedding_model.encode(
            mood_text,
            prompt_name=self.config.PROMPT_NAME
        )
        top_hits = util.semantic_search(
            mood_embedding, self.font_embeddings, top_k=self.config.TOP_K
        )[0]

        palette_html = self._format_palette_as_html(top_hits)
        return palette_html, top_hits

    def generate_css_code(self, top_hits: list[dict[str, any]]) -> str:
        """Generates exportable CSS code."""
        if not top_hits:
            return "/* No fonts generated yet */"
        
        imports = []
        seen_urls = set()
        
        for hit in top_hits:
            font_info = self.font_data[hit['corpus_id']]
            google_fonts_url = font_info['google_fonts_url']
            
            if google_fonts_url not in seen_urls:
                imports.append(f"@import url('{google_fonts_url}');")
                seen_urls.add(google_fonts_url)

        font_imports = "\n".join(imports)
        
        css_code = f"""/* Generated Font Palette CSS */
{font_imports}

/* Font Variables */
:root {{"""
        
        for i, hit in enumerate(top_hits):
            font_info = self.font_data[hit['corpus_id']]
            font_name = font_info['name']
            css_code += f"""
    --font-{i+1}: '{font_name}', {font_info['family']};"""
        
        css_code += """
}

/* Usage Examples */
.heading { font-family: var(--font-1); }
.body-text { font-family: var(--font-2); }
.accent { font-family: var(--font-3); }"""
        
        return css_code

    def apply_theme_css(self, top_hits: list[dict[str, any]]) -> str:
        """Generates CSS to apply fonts to the UI."""
        if not top_hits:
            return ""

        imports = []
        seen_urls = set()
        
        for hit in top_hits:
            font_info = self.font_data[hit['corpus_id']]
            google_fonts_url = font_info['google_fonts_url']
            
            if google_fonts_url not in seen_urls:
                imports.append(f"@import url('{google_fonts_url}');")
                seen_urls.add(google_fonts_url)

        font_imports = "\n".join(imports)
        
        css_rules = []
        
        if len(top_hits) >= 1:
            primary_font = self.font_data[top_hits[0]['corpus_id']]['name'].replace("'", "\\'")
            css_rules.append(f"h1, h2, h3, .gr-button-primary {{ font-family: '{primary_font}', sans-serif !important; }}")
        
        if len(top_hits) >= 2:
            secondary_font = self.font_data[top_hits[1]['corpus_id']]['name'].replace("'", "\\'")
            css_rules.append(f".gr-textbox input, .gr-textbox textarea {{ font-family: '{secondary_font}', sans-serif !important; }}")
        
        if len(top_hits) >= 3:
            tertiary_font = self.font_data[top_hits[2]['corpus_id']]['name'].replace("'", "\\'")
            css_rules.append(f".gr-button-secondary {{ font-family: '{tertiary_font}', sans-serif !important; }}")

        css_rules_str = "\n            ".join(css_rules)

        css = f"""<style>
            {font_imports}
            
            {css_rules_str}
            
            * {{
                transition: font-family 0.3s ease-in-out;
            }}
        </style>"""

        return css


# --- GRADIO UI WITH WALKTHROUGH ---
def create_ui(generator: FontMoodGenerator):
    """Creates the Gradio web interface with Walkthrough."""
    
    with gr.Blocks(theme="ocean") as demo:
        
        gr.Markdown("""
            # 📝 Font Mood Generator
            Follow the steps below to generate and apply a personalized font palette based on your mood or description.
        """)

        with gr.Walkthrough(selected=0) as walkthrough:
            
            # STEP 1: Input Mood
            with gr.Step("🎯 Describe Your Mood", id=0):
                gr.Markdown("""
                    ### Step 1: Tell us about your mood or vision
                    Describe the feeling, atmosphere, or aesthetic you're aiming for. 
                    Be as detailed as you like - the more descriptive, the better the results!
                """)
                
                mood_input = gr.Textbox(
                    value="Elegant wedding invitation with vintage charm",
                    label="Enter Your Mood or Scene",
                    info="Examples: 'Modern tech startup', 'Playful children's book', 'Gothic horror movie'",
                    lines=3
                )
                
                gr.Examples(
                    [
                        "Elegant wedding invitation with vintage charm",
                        "Modern tech startup with clean aesthetics", 
                        "Playful children's book with whimsical characters",
                        "Horror movie poster with scary atmosphere",
                        "Luxury fashion brand with sophisticated appeal"
                    ],
                    inputs=mood_input,
                )
                
                generate_btn = gr.Button("Generate Font Palette →", variant="primary", size="lg")
                
                # Hidden outputs to store results
                palette_html_hidden = gr.HTML(visible=False)
                font_data_hidden = gr.JSON(visible=False)
                
                def generate_and_move(mood_text):
                    palette_html, top_hits = generator.generate_palette(mood_text)
                    # Convert to serializable format
                    font_data_json = [{"corpus_id": hit["corpus_id"], "score": hit["score"]} for hit in top_hits]
                    return palette_html, font_data_json, gr.Walkthrough(selected=1)
                
                generate_btn.click(
                    fn=generate_and_move,
                    inputs=mood_input,
                    outputs=[palette_html_hidden, font_data_hidden, walkthrough]
                )
            
            # STEP 2: Review Generated Fonts
            with gr.Step("🎨 Review Your Font Palette", id=1):
                gr.Markdown("""
                    ### Step 2: Review your generated fonts
                    Here are the fonts that best match your mood, ranked by similarity score.
                """)
                
                palette_display = gr.HTML()
                
                with gr.Row():
                    back_to_input_btn = gr.Button("← Back to Input", variant="secondary")
                    apply_theme_btn = gr.Button("Apply Typography Theme →", variant="primary", size="lg")
                
                back_to_input_btn.click(
                    fn=lambda: gr.Walkthrough(selected=0),
                    outputs=walkthrough
                )
                
                # Update display when entering this step
                def show_generated_palette(palette_html):
                    return palette_html
                
                palette_html_hidden.change(
                    fn=show_generated_palette,
                    inputs=palette_html_hidden,
                    outputs=palette_display
                )
                
                # Hidden CSS output for theming
                theme_css_hidden = gr.HTML(visible=False)
                
                def apply_theme_and_move(font_data_json):
                    # Convert back to the format expected by apply_theme_css
                    top_hits = [{"corpus_id": item["corpus_id"], "score": item["score"]} for item in font_data_json]
                    theme_css = generator.apply_theme_css(top_hits)
                    return theme_css, gr.Walkthrough(selected=2)
                
                apply_theme_btn.click(
                    fn=apply_theme_and_move,
                    inputs=font_data_hidden,
                    outputs=[theme_css_hidden, walkthrough]
                )
            
            # STEP 3: Experience the Typography
            with gr.Step("✨ Experience Your Typography", id=2):
                gr.Markdown("""
                    ### Step 3: See your fonts in action!
                    Notice how the entire interface has transformed to reflect your chosen aesthetic.
                """)
                
                # Apply CSS when entering this step
                theme_css_display = gr.HTML()
                
                theme_css_hidden.change(
                    fn=lambda css: css,
                    inputs=theme_css_hidden,
                    outputs=theme_css_display
                )
                
                gr.Markdown("""
                    **🎉 Your typography theme is now active!** 
                    
                    Look around the interface - the headings, buttons, and text inputs now use fonts from your generated palette.
                    
                    **Font Roles:**
                    - **Primary Font**: Used for headings and primary buttons
                    - **Secondary Font**: Used for input fields and body text  
                    - **Accent Font**: Used for secondary buttons and highlights
                """)
                
                with gr.Row():
                    back_to_palette_btn = gr.Button("← Back to Palette", variant="secondary")
                    get_code_btn = gr.Button("Get CSS Code →", variant="primary", size="lg")
                
                back_to_palette_btn.click(
                    fn=lambda: gr.Walkthrough(selected=1),
                    outputs=walkthrough
                )
                
                get_code_btn.click(
                    fn=lambda: gr.Walkthrough(selected=3),
                    outputs=walkthrough
                )
            
            # STEP 4: Export and Use
            with gr.Step("💾 Export & Use Your Fonts", id=3):
                gr.Markdown("""
                    ### Step 4: Get the code and use your fonts
                    Copy the CSS code below to use your font palette in your own projects.
                """)
                
                css_code_output = gr.Code(
                    language="css", 
                    label="Your Font Palette CSS",
                    value="/* Generate a palette first to see CSS code here */",
                    lines=15
                )
                
                # Update CSS code when font data changes
                def update_css_code(font_data_json):
                    if not font_data_json:
                        return "/* Generate a palette first to see CSS code here */"
                    top_hits = [{"corpus_id": item["corpus_id"], "score": item["score"]} for item in font_data_json]
                    return generator.generate_css_code(top_hits)
                
                font_data_hidden.change(
                    fn=update_css_code,
                    inputs=font_data_hidden,
                    outputs=css_code_output
                )
                
                gr.Markdown("""
                    **🚀 Next Steps:**
                    1. Copy the CSS code above
                    2. Include it in your website's stylesheet
                    3. Apply the font variables to your HTML elements
                    4. Enjoy your new typography!
                """)
                
                start_over_btn = gr.Button("🔄 Start Over", variant="secondary", size="lg")
                
                def restart():
                    return "", [], "", gr.Walkthrough(selected=0)
                
                start_over_btn.click(
                    fn=restart,
                    outputs=[palette_html_hidden, font_data_hidden, theme_css_hidden, walkthrough]
                )

        # Static CSS for font cards
        gr.HTML("""
            <style>
                .font-palette-container {
                    display: flex; flex-direction: column; gap: 15px;
                    align-items: center; width: 100%;
                }
                .font-card {
                    border: 2px solid #e0e0e0; border-radius: 12px;
                    padding: 20px; width: 90%; max-width: 600px;
                    background: linear-gradient(135deg, #f8f9fa 0%, #ffffff 100%);
                    box-shadow: 0 4px 12px rgba(0,0,0,0.1);
                    transition: all 0.3s ease;
                }
                .font-card:hover {
                    transform: translateY(-2px);
                    box-shadow: 0 8px 20px rgba(0,0,0,0.15);
                }
                .font-header {
                    display: flex; justify-content: space-between;
                    align-items: center; margin-bottom: 15px;
                }
                .font-header h3 {
                    margin: 0; color: #2c3e50; font-size: 1.2em;
                }
                .font-score {
                    background: #3498db; color: white; padding: 4px 8px;
                    border-radius: 12px; font-size: 0.8em; font-weight: bold;
                }
                .font-sample {
                    font-size: 24px; line-height: 1.4; margin: 15px 0;
                    padding: 15px; background: #f8f9fa; border-radius: 8px;
                    border-left: 4px solid #3498db; color: #2c3e50;
                    min-height: 60px; display: flex; align-items: center;
                }
                .font-details {
                    display: flex; flex-direction: column; gap: 8px;
                }
                .font-family {
                    font-weight: bold; color: #7f8c8d; font-size: 0.9em;
                    text-transform: uppercase; letter-spacing: 1px;
                }
                .font-description {
                    color: #5d6d7e; font-size: 0.9em; line-height: 1.4;
                }
            </style>
        """)

    return demo


if __name__ == "__main__":
    # Initialize application components
    generator = FontMoodGenerator(config=Config(), font_data=FONT_DATA)
    
    demo = create_ui(generator)
    demo.launch()