import gradio as gr
import numpy as np
import os
from huggingface_hub import login
from sentence_transformers import SentenceTransformer, util
import pandas as pd
from datasets import load_dataset


# --- 1. 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')


# --- 2. COLOR DATA ---
# The color palette data is kept separate for clarity and easy modification.

ds = load_dataset("burkelibbey/colors")
df = pd.DataFrame(ds['train'])

df = df.rename(columns={"color": "hex"})

# Split column into two new ones
df[["name", "description"]] = df["description"].str.split(":", n=1, expand=True)
df["name"] = df["name"].str.title()

df = df.drop_duplicates(subset="name")

COLOR_DATA = df.to_dict(orient="records")


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

class MoodPaletteGenerator:
    """Handles model loading, embedding generation, and palette creation."""

    def __init__(self, config: Config, color_data: list[dict[str, any]]):
        """Initializes the generator, logs in, and loads necessary assets."""
        self.config = config
        self.color_data = color_data
        self._login_to_hf()
        self.embedding_model = self._load_model()
        self.color_embeddings = self._precompute_color_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_color_embeddings(self) -> np.ndarray:
        """Generates and stores embeddings for the color descriptions."""
        print("Pre-computing embeddings for color palette...")
        color_texts = [
            f"{color['name']}, {color['description']}"
            for color in self.color_data
        ]
        embeddings = self.embedding_model.encode(
            color_texts,
            prompt_name=self.config.PROMPT_NAME,
            show_progress_bar=True
        )
        print("Embeddings computed successfully.")
        return embeddings

    def _get_text_color_for_bg(self, hex_color: str) -> str:
        """
        Calculates the luminance of a hex color and returns black ('#000000')
        or white ('#FFFFFF') for the best text contrast.
        """
        hex_color = hex_color.lstrip('#')
        try:
            r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
            luminance = (0.299 * r + 0.587 * g + 0.114 * b)
            return '#000000' if luminance > 150 else '#FFFFFF'
        except (ValueError, IndexError):
            return '#000000' # Default to black on invalid hex

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

        cards_html = ""
        for hit in top_hits:
            color_info = self.color_data[hit['corpus_id']]
            hex_code = color_info['hex']
            name = color_info['name']
            score = hit['score']
            text_color = self._get_text_color_for_bg(hex_code)

            cards_html += f"""
            <div class="color-card" style="background-color: {hex_code}; color: {text_color};">
                {name} | {hex_code} | Score: {score:.2f}
            </div>
            """
        return f"<div class='palette-container'>{cards_html}</div>"

    def _create_dynamic_theme_css(self, top_hits: list[dict[str, any]]) -> str:
        """Generates a <style> block to override Gradio theme variables."""
        theme_colors = []

        if top_hits:
            theme_colors = [
                {
                    "bg": self.color_data[hit['corpus_id']]['hex'],
                    "txt": self._get_text_color_for_bg(self.color_data[hit['corpus_id']]['hex'])
                }
                for hit in top_hits
            ]

        css_map = {
            "button-primary-background-fill": (0, 'bg'),
            "button-primary-text-color": (0, 'txt'),
            "button-secondary-background-fill": (1, 'bg'),
            "button-secondary-text-color": (1, 'txt'),
            "block-background-fill": (2, 'bg'),
            "block-info-text-color": (2, 'txt'),
            "block-title-background-fill": (3, 'bg'),
            "button-primary-background-fill-hover": (3, 'bg'),
            "block-title-text-color": (3, 'txt'),
            "button-primary-text-color-hover": (3, 'txt'),
            "button-secondary-background-fill-hover": (4, 'bg'),
            "button-secondary-text-color-hover": (4, 'txt'),
        }

        css_rules = []
        num_available_colors = len(theme_colors)

        for var_suffix, (index, key) in css_map.items():
            if index < num_available_colors:
                color_value = theme_colors[index][key]
                css_rules.append(f"--{var_suffix}: {color_value};")

        css_rules_str = "\n".join(css_rules)

        # Create CSS variables to inject
        css = f"""
        <style>
            :root {{
{css_rules_str}
            }}

            :root .dark {{
{css_rules_str}
            }}

            .gallery-item .gallery {{
                background: {theme_colors[4]['bg']};
                color: {theme_colors[4]['txt']};
            }}
        </style>
        """

        return css

    def generate_palette_and_theme(self, mood_text: str) -> tuple[str, str]:
        """
        Generates a color 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.color_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 "", ""


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

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

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

        with gr.Row():
            with gr.Column(scale=4):
                mood_input = gr.Textbox(
                    value="Strawberry ice cream",
                    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 Palette", variant="primary")
                clear_button = gr.Button("Clear", variant="secondary")

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

        # Define CSS for palette cards here once, instead of in every update
        gr.HTML("""
            <style>
                .palette-container {
                    display: flex; flex-direction: column; gap: 10px;
                    align-items: center; width: 100%;
                q}
                .color-card {
                    border-radius: 10px; text-align: center; padding: 15px 10px;
                    width: 90%; max-width: 400px;
                    box-shadow: 0 4px 8px 0 rgba(0,0,0,0.2);
                    font-family: sans-serif; font-size: 16px; font-weight: bold;
                    transition: transform 0.2s;
                }
                .color-card:hover { transform: scale(1.05); }
            </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(
            [
                "A futuristic city at night, slick with rain",
                "The feeling of a cozy cabin during a blizzard",
                "Joyful chaos at a summer music festival",
                "Beach sunset with the palm tree",
                "A calm and lonely winter morning",
                "Vintage romance in a dusty library",
                "Cyberpunk alleyway with neon signs",
            ],
            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 **Mood Palette Generator**, transforms your words into a vibrant color palette. Simply describe a mood, a scene, or a feeling and the app will generate a set of matching colors. As a unique touch, the entire user interface dynamically updates its theme to reflect the generated palette, immersing you in your chosen mood.

             ## 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 Color Search**: The application has a pre-defined library of colors, where each color 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 color embeddings to find the closest matches based on a similarity score.
             4.  **Palette Creation**: The colors with the highest similarity scores are selected and presented to you as a complete palette.

             The Mood Palette Generator is a perfect example of how embeddings can be used for creative applications beyond simple text search.
        """)

    return demo

if __name__ == "__main__":
    # Initialize application components
    generator = MoodPaletteGenerator(config=Config(), color_data=COLOR_DATA)
    
    demo = create_ui(generator)
    demo.launch()