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

# --- 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.

COLOR_DATA = [
  {
    "name": "Crimson",
    "hex": "#DC143C",
    "description": "A deep, rich red color, leaning slightly towards purple."
  },
  {
    "name": "Scarlet",
    "hex": "#FF2400",
    "description": "A brilliant, vivid red with a hint of orange."
  },
  {
    "name": "Coral",
    "hex": "#FF7F50",
    "description": "A vibrant pinkish-orange reminiscent of marine invertebrates."
  },
  {
    "name": "Tangerine",
    "hex": "#F28500",
    "description": "A saturated, zesty orange, like the ripe citrus fruit."
  },
  {
    "name": "Gold",
    "hex": "#FFD700",
    "description": "A bright, metallic yellow associated with wealth and luxury."
  },
  {
    "name": "Lemon Chiffon",
    "hex": "#FFFACD",
    "description": "A pale, light yellow, as soft and airy as the dessert."
  },
  {
    "name": "Lime Green",
    "hex": "#32CD32",
    "description": "A bright green color, evoking freshness and zesty energy."
  },
  {
    "name": "Forest Green",
    "hex": "#228B22",
    "description": "A dark, shaded green, like the canopy of a dense forest."
  },
  {
    "name": "Teal",
    "hex": "#008080",
    "description": "A medium blue-green color, often seen as sophisticated and calming."
  },
  {
    "name": "Cyan",
    "hex": "#00FFFF",
    "description": "A vibrant greenish-blue, one of the primary subtractive colors."
  },
  {
    "name": "Sky Blue",
    "hex": "#87CEEB",
    "description": "A light, pale blue, like the color of a clear daytime sky."
  },
  {
    "name": "Royal Blue",
    "hex": "#4169E1",
    "description": "A deep, vivid blue that is both rich and bright."
  },
  {
    "name": "Indigo",
    "hex": "#4B0082",
    "description": "A deep, rich color between blue and violet in the spectrum."
  },
  {
    "name": "Lavender",
    "hex": "#E6E6FA",
    "description": "A light, pale purple with a bluish hue, named after the flower."
  },
  {
    "name": "Plum",
    "hex": "#DDA0DD",
    "description": "A reddish-purple color, like the ripe fruit it's named after."
  },
  {
    "name": "Magenta",
    "hex": "#FF00FF",
    "description": "A purplish-red color that lies between red and violet."
  },
  {
    "name": "Hot Pink",
    "hex": "#FF69B4",
    "description": "A bright, vivid pink that is both bold and energetic."
  },
  {
    "name": "Ivory",
    "hex": "#FFFFF0",
    "description": "An off-white color that resembles the material from tusks and teeth."
  },
  {
    "name": "Beige",
    "hex": "#F5F5DC",
    "description": "A pale sandy fawn color, often used as a warm, neutral tone."
  },
  {
    "name": "Taupe",
    "hex": "#483C32",
    "description": "A dark grayish-brown or brownish-gray color."
  },
  {
    "name": "Slate Gray",
    "hex": "#708090",
    "description": "A medium gray with a slight blue tinge, like the metamorphic rock."
  },
  {
    "name": "Charcoal",
    "hex": "#36454F",
    "description": "A dark, almost black gray, like burnt wood."
  },
  {
    "name": "Onyx",
    "hex": "#353839",
    "description": "A deep, rich black, often with a subtle hint of dark blue."
  },
  {
    "name": "Emerald",
    "hex": "#50C878",
    "description": "A brilliant green, named after the precious gemstone."
  },
  {
    "name": "Sapphire",
    "hex": "#0F52BA",
    "description": "A deep, lustrous blue, reminiscent of the valuable gemstone."
  },
  {
    "name": "Ruby",
    "hex": "#E0115F",
    "description": "A deep red color, inspired by the gemstone of the same name."
  },
  {
    "name": "Amethyst",
    "hex": "#9966CC",
    "description": "A moderate, violet-purple color, like the quartz gemstone."
  },
  {
    "name": "Peridot",
    "hex": "#E6E200",
    "description": "A light olive-green or yellowish-green, named for the gem."
  },
  {
    "name": "Turquoise",
    "hex": "#40E0D0",
    "description": "A greenish-blue color, often associated with tropical waters."
  },
  {
    "name": "Silver",
    "hex": "#C0C0C0",
    "description": "A metallic gray color that resembles polished silver."
  },
  {
    "name": "Bronze",
    "hex": "#CD7F32",
    "description": "A metallic brown color that resembles the alloy of copper and tin."
  },
  {
    "name": "Obsidian",
    "hex": "#000000",
    "description": "A pure, deep black, like the volcanic glass formed from cooled lava."
  }
]


# --- 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
        )
        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()