Update app.py

app.py

@@ -1,31 +1,33 @@
 import gradio as gr
-from google_play_scraper import search, app
+from google_play_scraper import search, app as play_app
 import os
 import time
 from datetime import datetime
 import json
-import shutil
+import google.generativeai as genai
 
 from huggingface_hub import HfApi, hf_hub_download, create_repo
 
 # -----------------------------
-# CONFIG & HF CONNECTION
+# CONFIG & HF/GEMINI CONNECTION
 # -----------------------------
 HF_TOKEN = os.getenv("HF_TOKEN")
-APP_PASSWORD = os.getenv("APP_PASSWORD") # Set this in your hf Space Secrets
+APP_PASSWORD = os.getenv("APP_PASSWORD") 
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") 
 DATASET_REPO = "kacapower/lexical-space-data"
 THEME_FILE_NAME = "website_theme.xml"
 
+# -----------------------------
+# CORE FUNCTIONS
+# -----------------------------
 def save_theme_to_dataset(xml_file):
     """Saves the uploaded XML file directly to the private hf dataset."""
     if xml_file is None:
         return "❌ Please upload an XML file first."
         
     try:
-        # Safely get the file path
         file_path = xml_file if isinstance(xml_file, str) else xml_file.name
         
-        # NEW: Automatically create the dataset repository if it doesn't exist
         try:
             create_repo(
                 repo_id=DATASET_REPO, 
@@ -37,7 +39,6 @@ def save_theme_to_dataset(xml_file):
         except Exception as repo_err:
             return f"❌ Failed to access or create repository: {repo_err}. Check your HF_TOKEN permissions."
 
-        # Upload directly using HfApi
         api = HfApi()
         api.upload_file(
             path_or_fileobj=file_path,
@@ -52,18 +53,18 @@ def save_theme_to_dataset(xml_file):
     except Exception as e:
         return f"❌ Error saving theme to dataset: {e}"
 
+
 def generate_blogger_html(app_query, custom_link, mod_features, theme_color, telegram_link, custom_size, custom_date):
+    """Scrapes Google Play and generates the base SEO-optimized Blogger HTML."""
     try:
-        # 1. Agent searches the Play Store for the app
         search_results = search(app_query, lang='en', country='us')
         if not search_results:
-            return "Error: Could not find this app on the Play Store.", "", "0", "0"
+            return "Error: Could not find this app on the Play Store.", "<b>App not found</b>", "0", "0"
             
-        # 2. Extract the exact App ID and fetch full details
         app_id = search_results[0]['appId']
-        app_details = app(app_id, lang='en', country='us')
+        app_details = play_app(app_id, lang='en', country='us')
         
-        # 3. Parse the extended data
+        # Data Extraction
         title = app_details.get('title', app_query)
         icon_url = app_details.get('icon', '')
         rating = round(app_details.get('score', 4.5), 1)
@@ -73,117 +74,77 @@ def generate_blogger_html(app_query, custom_link, mod_features, theme_color, tel
         installs = app_details.get('installs', '10,000+')
         version = app_details.get('version', 'Varies with device')
         
-        # Use custom size or scraped size
         scraped_size = app_details.get('size', 'Varies')
-        final_size = custom_size if custom_size else scraped_size
+        final_size = custom_size.strip() if custom_size else scraped_size
 
-        # Format the updated timestamp or use custom override
         updated_ts = app_details.get('updated', 0)
         scraped_date = datetime.fromtimestamp(updated_ts).strftime('%B %d, %Y') if updated_ts else "Recently"
-        updated_date = custom_date if custom_date else scraped_date
+        updated_date = custom_date.strip() if custom_date else scraped_date
         
-        # Use full description for better SEO and longer content
         description = app_details.get('summary', 'Premium application with enhanced features.')
         full_description = app_details.get('descriptionHTML', description)
         
-        # 4. Format the custom Mod Features into an HTML list
+        # Media Extraction (Enhanced responsiveness)
+        screenshots = app_details.get('screenshots', [])[:4] # Take up to 4 screenshots
+        video_url = app_details.get('video', None)
+
+        # Mod Features Formatting
         features_list = ""
         for feature in mod_features.split(','):
             if feature.strip():
                 features_list += f'<li style="margin-bottom: 8px;"><i class="fa-solid fa-check" style="color: {theme_color}; margin-right: 10px;"></i> {feature.strip()}</li>\n'
 
-        # 5. Build the SEO-Optimized Schemas (App + FAQ + SysReq)
+        # Screenshot HTML Generation
+        screenshots_html = ""
+        if screenshots:
+            screenshots_html = f'<h2 style="border-left: 4px solid {theme_color}; padding-left: 12px; margin-top: 30px;">App Screenshots</h2><div style="display: flex; gap: 15px; overflow-x: auto; padding-bottom: 15px; scroll-snap-type: x mandatory;">'
+            for snap in screenshots:
+                screenshots_html += f'<img src="{snap}" style="height: 280px; border-radius: 12px; border: 1px solid rgba(255,255,255,0.1); scroll-snap-align: start; box-shadow: 0 4px 10px rgba(0,0,0,0.3);">'
+            screenshots_html += '</div>'
+
+        # Video HTML Generation
+        video_html = ""
+        if video_url:
+            embed_url = video_url.replace("watch?v=", "embed/")
+            video_html = f'<h2 style="border-left: 4px solid {theme_color}; padding-left: 12px; margin-top: 30px;">Official Trailer</h2><div style="position: relative; padding-bottom: 56.25%; height: 0; overflow: hidden; border-radius: 16px; box-shadow: 0 5px 15px rgba(0,0,0,0.4);"><iframe src="{embed_url}" style="position: absolute; top: 0; left: 0; width: 100%; height: 100%;" frameborder="0" allowfullscreen></iframe></div>'
+
+        # SEO Schemas (App, FAQ, System Requirements)
         schema_app = {
-            "@context": "https://schema.org",
-            "@type": "SoftwareApplication",
-            "name": title,
-            "operatingSystem": "ANDROID",
-            "applicationCategory": category,
-            "image": icon_url,
-            "author": {
-                "@type": "Organization",
-                "name": developer
-            },
-            "aggregateRating": {
-                "@type": "AggregateRating",
-                "ratingValue": str(rating),
-                "ratingCount": str(reviews_count)
-            },
-            "offers": {
-                "@type": "Offer",
-                "price": "0.00",
-                "priceCurrency": "USD"
-            }
+            "@context": "https://schema.org", "@type": "SoftwareApplication", "name": title, "operatingSystem": "ANDROID", "applicationCategory": category, "image": icon_url,
+            "author": { "@type": "Organization", "name": developer },
+            "aggregateRating": { "@type": "AggregateRating", "ratingValue": str(rating), "ratingCount": str(reviews_count) },
+            "offers": { "@type": "Offer", "price": "0.00", "priceCurrency": "USD" }
         }
-        
-        # Generate FAQ data from the app title
         schema_faq = {
-            "@context": "https://schema.org",
-            "@type": "FAQPage",
+            "@context": "https://schema.org", "@type": "FAQPage",
             "mainEntity": [
-                {
-                    "@type": "Question",
-                    "name": f"Is {title} Mod APK safe?",
-                    "acceptedAnswer": {
-                        "@type": "Answer",
-                        "text": "Yes, our mods are thoroughly scanned and tested on physical devices before being published to the community."
-                    }
-                },
-                {
-                    "@type": "Question",
-                    "name": f"Does {title} require root access?",
-                    "acceptedAnswer": {
-                        "@type": "Answer",
-                        "text": "No, you do not need root access to install our modification. It works on any non-rooted Android device."
-                    }
-                }
+                {"@type": "Question", "name": f"Is {title} Mod APK safe?", "acceptedAnswer": {"@type": "Answer", "text": "Yes, our mods are thoroughly scanned and tested."}},
+                {"@type": "Question", "name": f"Does {title} require root access?", "acceptedAnswer": {"@type": "Answer", "text": "No, it works on any non-rooted Android device."}}
             ]
         }
-        
         schema_sysreq = {
-            "@context": "https://schema.org",
-            "@type": "SoftwareApplication",
-            "name": title,
-            "operatingSystem": "ANDROID",
-            "applicationCategory": category,
-            "author": {
-                "@type": "Organization",
-                "name": developer
-            },
-            "requirements": {
-                "@type": "SoftwareApplicationRequirements",
-                "name": "Minimum System Requirements",
-                "operatingSystem": "Android 9.0 (Pie)+",
-                "ram": "4 GB+",
-                "storage": "100 MB+"
-            }
+            "@context": "https://schema.org", "@type": "SoftwareApplication", "name": title, "operatingSystem": "ANDROID", "applicationCategory": category,
+            "author": { "@type": "Organization", "name": developer },
+            "requirements": { "@type": "SoftwareApplicationRequirements", "name": "Minimum System Requirements", "operatingSystem": "Android 9.0 (Pie)+", "ram": "4 GB+", "storage": "100 MB+" }
         }
-        
         schema_script = f'<script type="application/ld+json">\n[{json.dumps(schema_app)}, {json.dumps(schema_faq)}, {json.dumps(schema_sysreq)}]\n</script>'
 
-        # 6. Read and "Apply" the Theme from the Dataset
+        # Fetch Theme XML if available
         try:
-            downloaded_theme_path = hf_hub_download(
-                repo_id=DATASET_REPO, 
-                filename=THEME_FILE_NAME, 
-                repo_type="dataset", 
-                token=HF_TOKEN
-            )
-            # Since we cannot parse the XML in this environment, we will simulate 
-            # extracting a primary theme color from the XML file. We will assume 
-            # the user uses the `theme_color` picker from the UI as a way to specify
-            # this primary color in their website's theme.
-            theme_background = "#0f172a" # Default
-            # extracted_theme_color = theme_color # Simulated
-
+            downloaded_theme_path = hf_hub_download(repo_id=DATASET_REPO, filename=THEME_FILE_NAME, repo_type="dataset", token=HF_TOKEN)
+            theme_background = "#0f172a" 
         except Exception:
-            # If no theme is found, use defaults
             theme_background = "#0f172a"
-            # extracted_theme_color = "#3b82f6"
 
-        # 7. Build the Final HTML Template, integrated with deep shadows from example
+        # Final Base HTML Template Construction
         html_template = f"""
 {schema_script}
+<style>
+  @keyframes pulse {{ 0% {{ transform: scale(1); }} 50% {{ transform: scale(1.02); box-shadow: 0 0 30px {theme_color}80; }} 100% {{ transform: scale(1); }} }}
+  .lex-btn:hover {{ filter: brightness(1.2); }}
+  .app-post-wrapper * {{ box-sizing: border-box; }}
+</style>
+
 <div style="display: none;">
   <img src="{icon_url}" alt="Download {title} Mod APK" />
 </div>
@@ -204,7 +165,7 @@ def generate_blogger_html(app_query, custom_link, mod_features, theme_color, tel
   </div>
 
   <div style="text-align: center; margin-bottom: 35px; background: #020617; padding: 25px; border-radius: 16px; border: 1px solid rgba(255,255,255,0.02);">
-     <a href="{custom_link}" target="_blank" rel="nofollow noopener" style="display: inline-block; width: 100%; max-width: 400px; background: {theme_color}; color: #fff; text-decoration: none; padding: 16px; border-radius: 50px; font-weight: 800; font-size: 1.15rem; box-shadow: 0 0 25px {theme_color}80; transition: transform 0.3s ease;">
+     <a href="{custom_link}" target="_blank" rel="nofollow noopener" class="lex-btn" style="display: inline-block; width: 100%; max-width: 400px; background: {theme_color}; color: #fff; text-decoration: none; padding: 16px; border-radius: 50px; font-weight: 800; font-size: 1.15rem; box-shadow: 0 0 25px {theme_color}80; transition: transform 0.3s ease; animation: pulse 2s infinite;">
         <i class="fa-solid fa-cloud-arrow-down" style="margin-right: 8px;"></i> Download APK via Secure Link
      </a>
      <p style="font-size: 0.8rem; color: #64748b; margin-top: 15px; margin-bottom: 0;">
@@ -214,7 +175,7 @@ def generate_blogger_html(app_query, custom_link, mod_features, theme_color, tel
   
   <div style="background: rgba(56, 189, 248, 0.1); border: 1px solid rgba(56, 189, 248, 0.2); padding: 15px; border-radius: 12px; text-align: center; margin-bottom: 30px;">
     <strong style="color: #e2e8f0; margin-right: 15px;">Join Lexical Space Community:</strong>
-    <a href="{telegram_link}" target="_blank" style="display: inline-block; background: #0088cc; color: white; padding: 8px 20px; border-radius: 20px; text-decoration: none; font-weight: bold; font-size: 0.9rem;">✈️ Telegram</a>
+    <a href="{telegram_link}" target="_blank" style="display: inline-block; background: #0088cc; color: white; padding: 8px 20px; border-radius: 20px; text-decoration: none; font-weight: bold; font-size: 0.9rem; transition: 0.2s;">✈️ Telegram</a>
   </div>
 
   <h2 style="border-left: 4px solid #38bdf8; padding-left: 12px; margin-bottom: 20px; font-size: 1.3rem;">App Information</h2>
@@ -225,6 +186,9 @@ def generate_blogger_html(app_query, custom_link, mod_features, theme_color, tel
       <div><span style="display: block; font-size: 0.8rem; color: #94a3b8; text-transform: uppercase;">File Size</span><strong>{final_size}</strong></div>
   </div>
 
+  {video_html}
+  {screenshots_html}
+
   <h2 style="border-left: 4px solid {theme_color}; padding-left: 12px; margin-bottom: 15px; font-size: 1.3rem;">Max Store Modifications</h2>
   <div style="background: rgba(16, 185, 129, 0.05); border: 1px solid rgba(16, 185, 129, 0.2); padding: 20px; border-radius: 12px; margin-bottom: 35px; box-shadow: 0 5px 15px rgba(16, 185, 129, 0.1);">
       <ul style="color: #f1f5f9; line-height: 1.8; list-style-type: none; padding-left: 0; margin: 0; font-weight: 600;">
@@ -242,31 +206,83 @@ def generate_blogger_html(app_query, custom_link, mod_features, theme_color, tel
   </div>
 </div>
 """
-        # Calculate Stats
         word_count = str(len(html_template.split()))
-        keyword_density = str(html_template.lower().count(title.lower()))
+        keyword_density = str(html_template.lower().count(title.lower()[:10])) # Safer substring counting
         
-        # Save Draft locally for backup download
         with open("latest_draft.txt", "w", encoding="utf-8") as f:
             f.write(html_template)
 
         return html_template, html_template, word_count, keyword_density
         
     except Exception as e:
-        return f"Error generating HTML: {str(e)}", "<b>Error occurred</b>", "0", "0"
+        return f"Error generating HTML: {str(e)}", f"<b>Error occurred: {str(e)}</b>", "0", "0"
+
+# -----------------------------
+# GEMINI AI ENHANCER
+# -----------------------------
+def enhance_with_gemini(raw_html, style_name, state):
+    """Uses Google Gemini to restructure and theme the HTML while retaining functionality."""
+    if not state.get("logged_in") or (time.time() - state.get("login_time", 0) > 1800):
+        return "Session expired.", "Session Expired.", "0", "0", gr.update(visible=False)
+
+    if not GEMINI_API_KEY:
+        return raw_html, "<b>Error: GEMINI_API_KEY not found in HF secrets.</b>", "0", "0", gr.update(visible=False)
+        
+    if not raw_html or raw_html.startswith("Error"):
+        return raw_html, "<b>Please generate Base HTML first before enhancing.</b>", "0", "0", gr.update(visible=False)
+
+    try:
+        genai.configure(api_key=GEMINI_API_KEY)
+        model = genai.GenerativeModel('gemini-1.5-flash') 
+
+        prompt = f"""
+        You are the lead front-end developer for the 'Lexical Space' blog.
+        I have generated the following raw HTML for an app post. It contains accurate factual data, JSON-LD schemas, exact image links, download URLs, and video iframes.
+        Your task is to enhance and re-style this HTML to perfectly match the '{style_name}' theme.
+
+        STRICT RULES:
+        1. CRITICAL: DO NOT remove or alter the `<script type="application/ld+json">` blocks. They are critical for SEO.
+        2. CRITICAL: DO NOT alter ANY `href="..."` or `src="..."` attributes. The download links, Telegram links, image links, and video iframes MUST remain exactly as they are.
+        3. DO NOT hallucinate or change the factual data (file size, version, rating, text content).
+        4. ONLY output raw HTML. Do not wrap your response in ```html blocks. Do not add conversational intro/outro text.
+
+        Theme Guidelines for '{style_name}':
+        - If 'Lexical Standard': Deep space background (#0f172a), neon accents, glassmorphism filters, glowing borders, and pulse animations.
+        - If 'Lexical Minimalist': Cleaner, fewer borders, flat white/light-gray or pure black backgrounds, high readability, subtle drop shadows instead of neon glow.
+        - If 'Lexical Max Store': Replicate a premium mobile app store UI (like Google Play / App Store). Use grid layouts, highly rounded corners (24px), large action buttons, and emphasize the app icon and download statistics.
+
+        Base HTML Code to Enhance:
+        {raw_html}
+        """
+
+        response = model.generate_content(prompt)
+        enhanced_html = response.text.replace("```html", "").replace("```", "").strip()
+        
+        # Calculate new stats
+        word_count = str(len(enhanced_html.split()))
+        keyword_density = "N/A (AI Enhanced)"
+
+        # Overwrite draft with enhanced version
+        with open("latest_draft.txt", "w", encoding="utf-8") as f:
+            f.write(enhanced_html)
+
+        return enhanced_html, enhanced_html, word_count, keyword_density, gr.update(value="latest_draft.txt", visible=True)
+
+    except Exception as e:
+        return raw_html, f"<b>Gemini API Error:</b> {str(e)}<br><i>Showing unenhanced Base HTML instead.</i>", "0", "0", gr.update(visible=False)
+
 
 # -----------------------------
 # UI & AUTHENTICATION LOGIC
 # -----------------------------
-# Gradio 6.0 fix: Remove theme from constructor
-with gr.Blocks() as agent_ui:
-    # State variable to track authentication session
+with gr.Blocks(title="Lexical Space Internal Tool", theme=gr.themes.Monochrome(primary_hue="blue")) as agent_ui:
+    # State tracking
     session_state = gr.State({"logged_in": False, "login_time": 0})
     
     # --- LOGIN SCREEN ---
     with gr.Group(visible=True) as login_screen:
         gr.Markdown("# 🔒 Lexical Space Internal Tool")
-        gr.Markdown("Please enter the password to unlock the Post Generator Agent.")
+        gr.Markdown("Please enter the password to unlock the Post Generator Agent. Sessions expire after 30 minutes.")
         
         with gr.Row():
             pwd_input = gr.Textbox(label="Password", type="password", scale=4)
@@ -274,43 +290,54 @@ with gr.Blocks() as agent_ui:
             
         login_err = gr.Markdown("❌ Incorrect password.", visible=False)
 
-    # --- MAIN APPLICATION ---
+    # --- MAIN APPLICATION DASHBOARD ---
     with gr.Group(visible=False) as main_app:
         gr.Markdown("# 🚀 Lexical Space V2 Post Generator")
         gr.Markdown("Search the Play Store, inject mods, and generate SEO-optimized HTML for your Blogger post.")
         
         with gr.Row():
+            # --- LEFT COLUMN: SETTINGS ---
             with gr.Column(scale=1):
                 gr.Markdown("### 1. App Parameters")
-                app_input = gr.Textbox(label="App Name (e.g., TeraBox)", placeholder="Search query for Play Store...")
+                app_input = gr.Textbox(label="App Name (e.g., TeraBox, YouTube ReVanced)", placeholder="Search query for Play Store...")
                 link_input = gr.Textbox(label="Your Download Link (Mega, PixelDrain, etc.)", placeholder="https://...")
-                features_input = gr.Textbox(label="Mod Features (Comma separated)", placeholder="Premium Unlocked, No Ads", value="Premium Unlocked, No Ads")
+                features_input = gr.Textbox(label="Mod Features (Comma separated)", placeholder="Premium Unlocked, No Ads", value="Premium Unlocked, No Ads, Analytics Removed")
                 
-                # New Theme & XML Upload Accordion
                 with gr.Accordion("Theme & XML Upload (Optional)", open=False):
                     xml_theme_input = gr.File(label="Upload Website Theme XML")
                     save_theme_btn = gr.Button("Save Theme to Dataset")
                     theme_status_out = gr.Textbox(label="Theme Status", interactive=False)
 
-                # Advanced SEO & Theme Settings Accordion
                 with gr.Accordion("Advanced SEO & Theme Settings", open=False):
                     theme_color = gr.ColorPicker(label="Post Accent Color (Applied from Theme)", value="#3b82f6")
-                    telegram_link = gr.Textbox(label="Your Telegram Channel Link", value="https://t.me/lexicalspace")
+                    telegram_link = gr.Textbox(label="Your Telegram Channel Link", value="https://t.me/lexicalspace") # Fixed raw URL mapping
                     custom_size = gr.Textbox(label="Custom File Size Override (e.g. 150 MB)", placeholder="Optional Override...")
                     custom_date = gr.Textbox(label="Custom Updated Date Override (e.g. January 15, 2026)", placeholder="Optional Override...")
                 
-                generate_btn = gr.Button("⚙️ Generate HTML", variant="primary")
+                generate_btn = gr.Button("⚙️ Generate Base HTML", variant="primary")
                 
+                gr.Markdown("<hr>")
+                gr.Markdown("### 2. AI Style Enhancer (Gemini)")
                 with gr.Row():
-                    word_count_out = gr.Textbox(label="Total Word Count", interactive=False)
-                    kw_density_out = gr.Textbox(label="Keyword Density", interactive=False)
+                    style_dropdown = gr.Dropdown(
+                        choices=["Lexical Standard", "Lexical Minimalist", "Lexical Max Store"], 
+                        label="Target Post Style", 
+                        value="Lexical Standard",
+                        scale=2
+                    )
+                    enhance_btn = gr.Button("✨ Enhance Code", variant="secondary", scale=1)
                 
-                # Crash Fix: Remove the initial 'latest_draft.txt' value
+                gr.Markdown("<hr>")
+                gr.Markdown("### Post Statistics")
+                with gr.Row():
+                    word_count_out = gr.Textbox(label="Total Word Count", interactive=False)
+                    kw_density_out = gr.Textbox(label="Keyword Mentions", interactive=False)
                 
                 file_download = gr.File(label="Download Draft Backup (.txt file)", visible=False, interactive=False)
 
+            # --- RIGHT COLUMN: OUTPUT ---
             with gr.Column(scale=2):
-                gr.Markdown("### 2. Output & Preview")
+                gr.Markdown("### 3. Output & Preview")
                 with gr.Tabs():
                     with gr.Tab("Source Code (HTML)"):
                         html_output = gr.Code(label="Copy & Paste to Blogger HTML View", language="html", interactive=False)
@@ -319,9 +346,8 @@ with gr.Blocks() as agent_ui:
 
     # --- CONTROLLERS ---
     def verify_login(pwd, state):
-        """Checks the password against the hf secret and updates session state."""
         if not APP_PASSWORD:
-            return state, gr.update(), gr.update(), gr.update(value="❌ APP_PASSWORD missing from secrets.", visible=True)
+            return state, gr.update(), gr.update(), gr.update(value="❌ APP_PASSWORD missing from Space secrets.", visible=True)
             
         if pwd == APP_PASSWORD:
             state["logged_in"] = True
@@ -331,60 +357,48 @@ with gr.Blocks() as agent_ui:
             return state, gr.update(visible=True), gr.update(visible=False), gr.update(value="❌ Incorrect password.", visible=True)
 
     def secure_generate(app_query, custom_link, mod_features, t_color, t_link, c_size, c_date, state):
-        """Validates the 30-minute session before running the generation."""
-        # 1800 seconds = 30 minutes
+        # Validate session
         if not state.get("logged_in") or (time.time() - state.get("login_time", 0) > 1800):
             state["logged_in"] = False
-            return (
-                "Session expired. Please refresh and log in again.", # html_output
-                "",                                      # preview_output
-                "0",                                     # word_count
-                "0",                                     # kw_density
-                gr.update(visible=True),                 # login_screen
-                gr.update(visible=False),                # main_app
-                state,                                   # session_state
-                gr.update(visible=False)                 # file download
-            )
+            return ("Session expired.", "Session expired.", "0", "0", gr.update(visible=True), gr.update(visible=False), state, gr.update(visible=False))
 
-        # Run generation if session is valid
+        # Generate HTML
         code, preview, words, keywords = generate_blogger_html(app_query, custom_link, mod_features, t_color, t_link, c_size, c_date)
         
-        # --- Crash Fix: Only show the download button if generation succeeded and file exists ---
+        # Enable file download if successful
         if os.path.exists("latest_draft.txt") and not code.startswith("Error"):
             file_update = gr.update(value="latest_draft.txt", visible=True)
         else:
             file_update = gr.update(value=None, visible=False)
             
-        return (
-            code, preview, words, keywords,
-            gr.update(visible=False), # Keep login screen hidden
-            gr.update(visible=True),  # Keep main app visible
-            state,
-            file_update
-        )
+        return (code, preview, words, keywords, gr.update(visible=False), gr.update(visible=True), state, file_update)
 
     # --- EVENT BINDINGS ---
     login_btn.click(
-        fn=verify_login,
-        inputs=[pwd_input, session_state],
+        fn=verify_login, 
+        inputs=[pwd_input, session_state], 
         outputs=[session_state, login_screen, main_app, login_err]
     )
-
+    
     save_theme_btn.click(
-        fn=save_theme_to_dataset,
-        inputs=[xml_theme_input],
+        fn=save_theme_to_dataset, 
+        inputs=[xml_theme_input], 
         outputs=theme_status_out
     )
     
-    # Target the file download component inside the main_app tree
-    
-    
     generate_btn.click(
         fn=secure_generate,
         inputs=[app_input, link_input, features_input, theme_color, telegram_link, custom_size, custom_date, session_state],
         outputs=[html_output, preview_output, word_count_out, kw_density_out, login_screen, main_app, session_state, file_download]
     )
+    
+    # Show loading icon specifically on the AI enhance button
+    enhance_btn.click(
+        fn=enhance_with_gemini,
+        inputs=[html_output, style_dropdown, session_state],
+        outputs=[html_output, preview_output, word_count_out, kw_density_out, file_download],
+        show_progress="full"
+    )
 
 if __name__ == "__main__":
-    # Gradio 6.0 fix: Move theme into launch()
-    agent_ui.launch(theme=gr.themes.Monochrome(primary_hue="blue"))
+    agent_ui.launch()