import gradio as gr
import torch
import os
from transformers import AutoTokenizer, AutoModelForCausalLM
from gradio.themes.utils import colors, sizes
from gradio.themes import Soft

# Define custom themes
class GemmaLightTheme(Soft):
    def __init__(self):
        super().__init__(
            primary_hue=colors.blue,
            secondary_hue=colors.indigo,
            neutral_hue=colors.gray,
            spacing_size=sizes.spacing_md,
            radius_size=sizes.radius_md,
            text_size=sizes.text_md,
        )

class GemmaDarkTheme(Soft):
    def __init__(self):
        super().__init__(
            primary_hue=colors.blue,
            secondary_hue=colors.indigo,
            neutral_hue=colors.gray,
            spacing_size=sizes.spacing_md,
            radius_size=sizes.radius_md,
            text_size=sizes.text_md,
        )
        self.name = "gemma_dark"
        self.background_fill_primary = "#1F1F2E"
        self.background_fill_secondary = "#2A2A3C"
        self.border_color_primary = "#3A3A4C"
        self.border_color_secondary = "#4A4A5C"
        self.color_accent_soft = "#3B5FA3"
        self.color_accent = "#4B82C4"
        self.text_color = "#FFFFFF"
        self.text_color_subdued = "#CCCCCC"
        self.shadow_spread = "8px"
        self.shadow_inset = "0px 1px 2px 0px rgba(0, 0, 0, 0.1) inset"

# [Previous code for helper functions, model loading, etc., remains unchanged]

# Create Gradio interface
with gr.Blocks(theme=GemmaLightTheme()) as demo:
    # Header with theme toggle
    with gr.Row(equal_height=True):
        with gr.Column(scale=6):
            gr.Markdown(
                """
                # 🤖 Gemma Capabilities Demo
                
                This interactive demo showcases Google's Gemma model capabilities across different tasks.
                """
            )
        with gr.Column(scale=1, min_width=150):
            theme_toggle = gr.Radio(
                ["Light", "Dark"],
                value="Light",
                label="Theme",
                info="Switch between light and dark mode",
                elem_id="theme_toggle"
            )

    # Add CSS and JavaScript for themes
    gr.HTML("""
    <style>
    /* Base styles for better UI */
    .gr-group {
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        border-radius: 10px;
        padding: 20px;
        background-color: var(--block-background-fill);
    }
    .gr-tabs {
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        border-radius: 10px;
        overflow: hidden;
    }
    button.primary {
        transition: transform 0.2s, box-shadow 0.2s;
    }
    button.primary:hover {
        transform: translateY(-2px);
        box-shadow: 0 4px 8px rgba(0,0,0,0.2);
    }
    textarea, .input-box {
        font-size: 16px !important;
    }
    
    /* Dark theme overrides using Gradio variables */
    .dark-theme {
        --body-background-fill: #1F1F2E !important;
        --block-background-fill: #2A2A3C !important;
        --input-background-fill: #3A3A4C !important;
        --button-primary-background-fill: #4B82C4 !important;
        --button-secondary-background-fill: #3A3A4C !important;
        --border-color-primary: #4A4A5C !important;
        --body-text-color: #FFFFFF !important;
        --block-label-text-color: #CCCCCC !important;
        --input-text-color: #FFFFFF !important;
        --button-primary-text-color: #FFFFFF !important;
        --button-secondary-text-color: #FFFFFF !important;
        --shadow: 0 1px 2px rgba(0,0,0,0.5) !important;
    }
    .dark-theme .code {
        background-color: #2A2A3C !important;
        color: #FFFFFF !important;
    }
    </style>
    
    <script>
    (function() {
        function applyTheme() {
            const themeToggle = document.getElementById('theme_toggle');
            if (!themeToggle) return;
            const inputs = themeToggle.querySelectorAll('input');
            for (let input of inputs) {
                if (input.checked && input.value === 'Dark') {
                    document.body.classList.add('dark-theme');
                    return;
                }
            }
            document.body.classList.remove('dark-theme');
        }
        
        function setupObserver() {
            const themeToggle = document.getElementById('theme_toggle');
            if (!themeToggle) {
                setTimeout(setupObserver, 100);
                return;
            }
            applyTheme();
            const observer = new MutationObserver((mutations) => {
                mutations.forEach((mutation) => {
                    if (mutation.type === 'attributes' || mutation.type === 'childList') {
                        applyTheme();
                    }
                });
            });
            observer.observe(themeToggle, { attributes: true, childList: true, subtree: true });
            const inputs = themeToggle.querySelectorAll('input');
            inputs.forEach(input => input.addEventListener('change', applyTheme));
        }
        
        if (document.readyState === 'loading') {
            document.addEventListener('DOMContentLoaded', setupObserver);
        } else {
            setupObserver();
        }
    })();
    </script>
    """)

    # Authentication Section
    with gr.Group(elem_id="auth_box"):
        gr.Markdown("## 🔑 Authentication")
        with gr.Row(equal_height=True):
            with gr.Column(scale=3):
                hf_token = gr.Textbox(
                    label="Hugging Face Token",
                    placeholder="Enter your token here...",
                    type="password",
                    value=DEFAULT_HF_TOKEN,
                    info="Get your token from https://huggingface.co/settings/tokens"
                )
            with gr.Column(scale=1):
                auth_button = gr.Button("Authenticate", variant="primary")
        auth_status = gr.Markdown("Please authenticate to use the model.")
        
        def authenticate(token):
            return "⏳ Loading model... Please wait, this may take a minute."
        
        def auth_complete(token):
            return load_model(token)
        
        auth_button.click(fn=authenticate, inputs=[hf_token], outputs=[auth_status], queue=False).then(
            fn=auth_complete, inputs=[hf_token], outputs=[auth_status]
        )
        gr.Markdown(
            """
            ### How to get a token:
            1. Go to [Hugging Face Token Settings](https://huggingface.co/settings/tokens)
            2. Create a new token with read access
            3. Accept the [Gemma model license](https://huggingface.co/google/gemma-3-4b-pt)
            """
        )

    # [Remaining tabs and components remain unchanged for brevity]

    # Footer
    with gr.Group(elem_id="footer"):
        gr.Markdown(
            """
            ## About Gemma
            Gemma is a family of lightweight, state-of-the-art open models from Google...
            [Learn more about Gemma](https://huggingface.co/google/gemma-3-4b-pt)
            <div style="text-align: center; margin-top: 20px; padding: 10px;">
                <p>© 2023 Gemma Capabilities Demo | Made with ❤️ using Gradio</p>
            </div>
            """
        )

# Launch the app
demo.launch(share=False)