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

# --- CONFIGURATION ---
# Centralized place for all settings and constants.

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 ---
# Comprehensive font dataset with descriptions for mood matching.

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"
    },
    {
        "name": "Source Code Pro",
        "family": "monospace",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Source+Code+Pro:wght@300;400;600&display=swap",
        "description": "Technical, programming, coding, monospaced font designed for developers and technical documentation"
    },
    {
        "name": "Merriweather",
        "family": "serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Merriweather:wght@300;400;700&display=swap",
        "description": "Traditional, readable, pleasant serif designed for comfortable reading on screens and in print"
    },
    {
        "name": "Abril Fatface",
        "family": "serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Abril+Fatface&display=swap",
        "description": "Bold, dramatic, high-contrast display serif inspired by French and Italian typography, perfect for headlines"
    },
    {
        "name": "Great Vibes",
        "family": "script",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Great+Vibes&display=swap",
        "description": "Elegant, formal, calligraphic script with sophisticated curves, ideal for luxury and premium branding"
    },
    {
        "name": "Raleway",
        "family": "sans-serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Raleway:wght@300;400;600&display=swap",
        "description": "Sophisticated, thin, elegant sans-serif with distinctive 'W', perfect for fashion and high-end design"
    },
    {
        "name": "Fredoka One",
        "family": "display",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Fredoka+One&display=swap",
        "description": "Friendly, rounded, playful display font perfect for children's content, toys, and fun applications"
    },
    {
        "name": "Libre Baskerville",
        "family": "serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Libre+Baskerville:wght@400;700&display=swap",
        "description": "Classic, traditional, scholarly serif based on American Type Founder's Baskerville, perfect for academic texts"
    },
    {
        "name": "Poppins",
        "family": "sans-serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;600&display=swap",
        "description": "Geometric, modern, friendly sans-serif with circular forms, popular for contemporary web design"
    },
    {
        "name": "Lobster",
        "family": "script",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Lobster&display=swap",
        "description": "Bold, retro, vintage script font with a 1950s diner feel, perfect for nostalgic and Americana designs"
    },
    {
        "name": "Open Sans",
        "family": "sans-serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600&display=swap",
        "description": "Neutral, friendly, optimistic sans-serif designed for legibility across interfaces, print, and web"
    },
    {
        "name": "Shadows Into Light",
        "family": "handwriting",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Shadows+Into+Light&display=swap",
        "description": "Casual, handwritten, personal font that feels like natural handwriting with a marker or pen"
    },
    {
        "name": "Creepster",
        "family": "display",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Creepster&display=swap",
        "description": "Horror, scary, Halloween, gothic font with dripping effect, perfect for spooky and thriller themes"
    },
    {
        "name": "Righteous",
        "family": "display",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Righteous&display=swap",
        "description": "Futuristic, sci-fi, technology, bold display font with unique character shapes for modern designs"
    },
    {
        "name": "Satisfy",
        "family": "script",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Satisfy&display=swap",
        "description": "Casual, relaxed, handwritten script with natural flow, perfect for personal and informal communications"
    },
    {
        "name": "Anton",
        "family": "sans-serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Anton&display=swap",
        "description": "Bold, condensed, impactful sans-serif perfect for headlines, posters, and attention-grabbing text"
    },
    {
        "name": "Courgette",
        "family": "script",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Courgette&display=swap",
        "description": "French, bistro, café, elegant script font with continental European charm and sophistication"
    },
    {
        "name": "Indie Flower",
        "family": "handwriting",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Indie+Flower&display=swap",
        "description": "Indie, hipster, handwritten font with quirky personality, perfect for creative and artistic projects"
    },
    {
        "name": "PT Serif",
        "family": "serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=PT+Serif:wght@400;700&display=swap",
        "description": "Russian, Cyrillic, transitional serif with excellent readability for both Latin and Cyrillic scripts"
    },
    {
        "name": "Questrial",
        "family": "sans-serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Questrial&display=swap",
        "description": "Simple, clean, minimal sans-serif with subtle quirks, perfect for modern and understated designs"
    },
    {
        "name": "Bangers",
        "family": "display",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Bangers&display=swap",
        "description": "Comic book, superhero, pop art font inspired by mid-20th century comic books and advertisements"
    },
    {
        "name": "Sacramento",
        "family": "script",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Sacramento&display=swap",
        "description": "Monoline, cursive script with vintage charm, perfect for elegant and sophisticated branding"
    },
    {
        "name": "Bitter",
        "family": "serif",
        "google_fonts_url": "https://fonts.googleapis.com/css2?family=Bitter:wght@400;700&display=swap",
        "description": "Contemporary, slab serif with slight contrast, designed for comfortable reading in long texts"
    }
]


# --- CORE LOGIC ---
# Encapsulated in a class to manage state (model, embeddings) cleanly.

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.")
            print("Note: This may fail if the model is gated.")

    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 for different font types
        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 _create_font_imports_css(self, top_hits: list[dict[str, any]]) -> str:
        """Generates CSS imports for the selected fonts."""
        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)

        return "\n".join(imports)

    def _create_dynamic_theme_css(self, top_hits: list[dict[str, any]]) -> str:
        """Generates a <style> block to apply fonts to different UI elements."""
        if not top_hits:
            return ""

        font_imports = self._create_font_imports_css(top_hits)
        
        css_rules = []
        
        # Apply different fonts to different elements
        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}
            
            /* Custom animations for font changes */
            * {{
                transition: font-family 0.3s ease-in-out;
            }}
        </style>"""

        return css

    def generate_palette_and_theme(self, mood_text: str) -> tuple[str, str]:
        """
        Generates a font palette HTML and a dynamic theme CSS string.
        """
        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)
        theme_css = self._create_dynamic_theme_css(top_hits)

        return palette_html, theme_css

    def clear_theme(self) -> tuple[str, str]:
        return "", ""


# --- GRADIO UI ---
# Defines and launches the web interface.

def create_ui(generator: FontMoodGenerator):
    """Creates the Gradio web interface."""
    with gr.Blocks(theme="ocean") as demo:
        # This invisible component will hold our dynamic CSS
        dynamic_css_output = gr.HTML()

        gr.Markdown("""
            # 📝 Font Mood Generator
            Describe a mood, a scene, or a feeling, and get a matching font palette.<br>
            **The UI fonts will update to match your mood!**
        """)

        with gr.Row():
            with gr.Column(scale=4):
                mood_input = gr.Textbox(
                    value="Elegant wedding invitation",
                    label="Enter Your Mood or Scene",
                    info="Be as descriptive as you like!"
                )
            with gr.Column(scale=1, min_width=150):
                submit_button = gr.Button("Generate Fonts", variant="primary")
                clear_button = gr.Button("Clear", variant="secondary")

        palette_output = gr.HTML(label="Your Generated Font Palette")

        # Define CSS for font cards here once
        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>
        """)

        # Define the function to be called by events
        event_handler = generator.generate_palette_and_theme
        outputs_list = [palette_output, dynamic_css_output]

        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",
                "Retro 1950s diner with nostalgic vibes",
                "Academic research paper with scholarly tone",
                "Surf shop with California beach culture",
                "Gothic medieval manuscript with ancient feel",
                "Futuristic sci-fi interface with cyber aesthetics"
            ],
            inputs=mood_input,
            outputs=outputs_list,
            fn=event_handler,
            run_on_click=True,
        )

        submit_button.click(
            fn=event_handler,
            inputs=mood_input,
            outputs=outputs_list,
        )
        clear_button.click(
            fn=generator.clear_theme,
            outputs=outputs_list,
        )
        # Also allow submitting by pressing Enter in the textbox
        mood_input.submit(
            fn=event_handler,
            inputs=mood_input,
            outputs=outputs_list,
        )

        gr.Markdown("""
            ----
            ## What is this?
            This interactive application, the **Font Mood Generator**, transforms your words into a curated typography palette. Simply describe a mood, a scene, or a feeling and the app will generate a set of matching fonts. As a unique touch, the entire user interface dynamically updates its typography to reflect the generated palette, immersing you in your chosen aesthetic.
             
            ## How It Works?
            At its core, this tool is powered by [**EmbeddingGemma**](http://huggingface.co/google/embeddinggemma-300M), a state-of-the-art text embedding model. The process works in a few simple steps:
            
            1. **Text to Vector**: When you enter a description, EmbeddingGemma converts your text into a numerical representation called an **embedding**. This embedding captures the semantic essence, or the "vibe" of your words.
            2. **Semantic Font Search**: The application has a pre-defined library of fonts, where each font is associated with its own descriptive text and a pre-computed embedding.
            3. **Finding the Match**: Your input embedding is compared against the entire library of font embeddings to find the closest matches based on a similarity score.
            4. **Palette Creation**: The fonts with the highest similarity scores are selected and presented to you as a complete typography palette.
            
            The Font Mood Generator demonstrates how embeddings can be used for creative applications in typography and design, going beyond simple text search to understand the emotional and aesthetic qualities of typefaces.
        """)

    return demo

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