Upload app.py

app.py

@@ -1,10 +1,11 @@
 """
 ShopSite AI - Small Business Website Generator
-Block-based editing: structured ops via Python, creative blocks via Qwen Coder
+Block-based page builder: users compose pages from reusable blocks,
+LLM (Qwen Coder) rewrites individual block HTML on natural language instruction.
 """
 
 import gradio as gr
-import requests
+from PIL import Image, ImageDraw, ImageFont
 import json
 import zipfile
 import os
@@ -12,98 +13,212 @@ import base64
 import io
 import re
 import shutil
+import uuid
 import urllib.parse
 from pathlib import Path
 import spaces
 
 try:
     import torch
-    from diffusers import AutoPipelineForText2Image
-    SD_AVAILABLE = True
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    TORCH_AVAILABLE = True
 except ImportError:
+    TORCH_AVAILABLE = False
+    print("⚠️  torch/transformers not installed. LLM features disabled.")
+
+try:
+    from diffusers import AutoPipelineForText2Image
+    SD_AVAILABLE = True and TORCH_AVAILABLE
+except (ImportError, RuntimeError):
     SD_AVAILABLE = False
-    print("⚠️  diffusers/torch not installed. Poster generation disabled.")
+    print("⚠️  diffusers not available. Poster generation disabled.")
 
 # ============================================================
 # CONFIG
 # ============================================================
-OLLAMA_BASE = "http://localhost:11434"
-QWEN_MODEL = "qwen2.5:7b"
-QWEN_CODER_MODEL = "qwen2.5-coder:7b"
-SD_MODEL_ID = "stabilityai/sd-turbo"
+QWEN_MODEL       = "Qwen/Qwen2.5-7B-Instruct"
+QWEN_CODER_MODEL = "Qwen/Qwen2.5-Coder-14B-Instruct"
+SD_MODEL_ID      = "stabilityai/sd-turbo"
 
-WORK_DIR = Path("./workspace")
-WORK_DIR.mkdir(exist_ok=True)
+WORK_DIR     = Path("./workspace");  WORK_DIR.mkdir(exist_ok=True)
 TEMPLATE_DIR = Path("./templates")
 
-# Global state
-current_html = ""
-current_menu_data = {}
-current_site_info = {}
+# ============================================================
+# GLOBAL STATE
+# ============================================================
+current_html        = ""
+current_menu_data   = {}
+current_site_info   = {}
 current_template_key = "warm"
-sd_pipe = None
-
-# Editable block definitions - markers in template HTML
-BLOCKS = {
-    "Hero Banner": {
-        "start": "<!-- BLOCK:HERO:START -->",
-        "end": "<!-- BLOCK:HERO:END -->",
-        "desc": "Shop name, tagline, background decorations",
-        "mode": "coder",  # Creative: Qwen Coder modifies HTML
-    },
-    "Promo / Event": {
-        "start": "<!-- BLOCK:PROMO:START -->",
-        "end": "<!-- BLOCK:PROMO:END -->",
-        "desc": "Promotional banner or event poster area",
-        "mode": "coder",
-    },
-    "About / Story": {
-        "start": "<!-- BLOCK:ABOUT:START -->",
-        "end": "<!-- BLOCK:ABOUT:END -->",
-        "desc": "Brand story and shop description",
-        "mode": "coder",
-    },
-    "Contact Info": {
-        "start": "<!-- BLOCK:CONTACT:START -->",
-        "end": "<!-- BLOCK:CONTACT:END -->",
-        "desc": "Phone, address, opening hours",
-        "mode": "python",  # Structured: Python modifies data
-    },
-    "Menu Items": {
-        "start": "",
-        "end": "",
-        "desc": "Add, remove, or change menu items and prices",
-        "mode": "python",
-    },
-    "Style / Colors": {
-        "start": "",
-        "end": "",
-        "desc": "Change colors, fonts, and visual style",
-        "mode": "python",
-    },
+page_blocks         = []   # [{"id": str, "type": str, "html": str}, ...]
+sd_pipe             = None
+
+# Menu item HTML template — LLM can rewrite this to change structure
+# Placeholders: {name}, {price}, {img_tag}
+MENU_ITEM_TEMPLATE_DEFAULT = """\
+      <div class="menu-item" data-item-name="{name}">
+        {img_tag}
+        <div class="menu-item-info">
+          <div class="menu-item-name">{name}</div>
+          <div class="menu-item-price">{price}</div>
+        </div>
+      </div>"""
+current_menu_item_template = MENU_ITEM_TEMPLATE_DEFAULT
+
+# ============================================================
+# BLOCK DEFAULTS  (warm-theme HTML that works with CSS variables)
+# ============================================================
+BLOCK_DEFAULTS = {
+    "Hero Banner": """\
+  <div class="hero">
+    <div class="hero-badge">
+      <svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="10"/></svg>
+      Est. 2024
+    </div>
+    <h1>Shop Name</h1>
+    <p class="hero-tagline">Welcome to our shop</p>
+  </div>
+  <div class="info-pills">
+    <div class="pill">
+      <svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="10"/><polyline points="12 6 12 12 16 14"/></svg>
+      <span>Open Daily</span>
+    </div>
+    <div class="pill">
+      <svg viewBox="0 0 24 24"><path d="M21 10c0 7-9 13-9 13s-9-6-9-13a9 9 0 0 1 18 0z"/><circle cx="12" cy="10" r="3"/></svg>
+      <span>Visit Us</span>
+    </div>
+  </div>""",
+
+    "Promo / Event": """\
+  <div class="section-title">
+    <h2>What's New</h2>
+    <div class="line"></div>
+  </div>
+  <div class="promo-banner">
+    <div class="promo-placeholder">
+      <span>Coming Soon</span>
+      <small>Stay tuned for updates</small>
+    </div>
+  </div>""",
+
+    "About / Story": """\
+  <div class="section-title">
+    <h2>Our Story</h2>
+    <div class="line"></div>
+  </div>
+  <div class="about-card">
+    <p>Welcome to our shop. We are passionate about quality and great service. Come visit us and experience the difference.</p>
+  </div>""",
+
+    "Contact Info": """\
+  <div class="section-title">
+    <h2>Find Us</h2>
+    <div class="line"></div>
+  </div>
+  <div class="contact-section">
+    <a class="contact-item" href="">
+      <div class="contact-icon">
+        <svg viewBox="0 0 24 24"><path d="M22 16.92v3a2 2 0 0 1-2.18 2 19.79 19.79 0 0 1-8.63-3.07 19.5 19.5 0 0 1-6-6 19.79 19.79 0 0 1-3.07-8.67A2 2 0 0 1 4.11 2h3a2 2 0 0 1 2 1.72c.127.96.361 1.903.7 2.81a2 2 0 0 1-.45 2.11L8.09 9.91a16 16 0 0 0 6 6l1.27-1.27a2 2 0 0 1 2.11-.45c.907.339 1.85.573 2.81.7A2 2 0 0 1 22 16.92z"/></svg>
+      </div>
+      <div class="contact-text"><div class="label">Phone</div><div class="value">—</div></div>
+    </a>
+    <div class="contact-item">
+      <div class="contact-icon">
+        <svg viewBox="0 0 24 24"><path d="M21 10c0 7-9 13-9 13s-9-6-9-13a9 9 0 0 1 18 0z"/><circle cx="12" cy="10" r="3"/></svg>
+      </div>
+      <div class="contact-text"><div class="label">Address</div><div class="value">—</div></div>
+    </div>
+    <div class="contact-item">
+      <div class="contact-icon">
+        <svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="10"/><polyline points="12 6 12 12 16 14"/></svg>
+      </div>
+      <div class="contact-text"><div class="label">Hours</div><div class="value">—</div></div>
+    </div>
+  </div>
+  <div class="spacer-lg"></div>""",
+
+    "Announcement": """\
+  <div style="margin:16px 20px;">
+    <div style="background:var(--bg-warm,#F5EDE3);border-radius:var(--card-radius,18px);padding:20px 24px;border-left:4px solid var(--primary);">
+      <div style="font-size:11px;text-transform:uppercase;letter-spacing:0.08em;color:var(--text-muted);font-weight:700;margin-bottom:8px;">📢 Notice</div>
+      <p style="font-size:15px;line-height:1.6;color:var(--text);">Add your announcement here.</p>
+    </div>
+  </div>""",
+
+    "Menu Preview": """\
+  <div class="section-title">
+    <h2>Our Menu</h2>
+    <div class="line"></div>
+  </div>
+  <div style="padding:0 24px 16px;text-align:center;">
+    <button onclick="switchPage('menu')" style="background:var(--primary);color:var(--bg);border:none;padding:12px 32px;border-radius:100px;font-size:14px;font-weight:600;cursor:pointer;letter-spacing:0.04em;">View Full Menu →</button>
+  </div>""",
 }
 
+# ============================================================
+# BLOCK HELPERS
+# ============================================================
+def _uid():
+    return str(uuid.uuid4())[:6]
+
+def _block_label(block):
+    return f"{block['type']}  [{block['id']}]"
+
+def _find_block(label):
+    for b in page_blocks:
+        if _block_label(b) == label:
+            return b
+    return None
+
+def get_block_choices():
+    return [_block_label(b) for b in page_blocks]
 
 # ============================================================
-# OLLAMA
+# HuggingFace LOCAL INFERENCE
 # ============================================================
-def ollama_chat(model, system_prompt, user_message, temperature=0.3):
-    try:
-        resp = requests.post(
-            f"{OLLAMA_BASE}/api/chat",
-            json={"model": model, "messages": [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_message},
-            ], "stream": False, "options": {"temperature": temperature}},
-            timeout=180,
+_hf_models = {}
+
+def load_hf_model(model_id):
+    if model_id not in _hf_models:
+        print(f"Loading {model_id}...")
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto",
         )
-        resp.raise_for_status()
-        return resp.json()["message"]["content"]
-    except requests.exceptions.ConnectionError:
-        return "ERROR: Cannot connect to Ollama. Run: ollama serve"
+        _hf_models[model_id] = (tokenizer, model)
+        print(f"✅ {model_id} loaded.")
+    return _hf_models[model_id]
+
+def _ollama_chat(model, system_prompt, user_message, temperature=0.3):
+    """Inner implementation — safe to call from within a @spaces.GPU context."""
+    if not TORCH_AVAILABLE:
+        return "ERROR: torch/transformers not installed."
+    try:
+        tokenizer, hf_model = load_hf_model(model)
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user",   "content": user_message},
+        ]
+        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(text, return_tensors="pt").to(hf_model.device)
+        max_new_tokens = 1024 if model == QWEN_CODER_MODEL else 512
+        with torch.no_grad():
+            outputs = hf_model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature if temperature > 0 else None,
+                do_sample=temperature > 0,
+                pad_token_id=tokenizer.eos_token_id,
+            )
+        return tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     except Exception as e:
         return f"ERROR: {e}"
 
+@spaces.GPU
+def ollama_chat(model, system_prompt, user_message, temperature=0.3):
+    return _ollama_chat(model, system_prompt, user_message, temperature)
 
 def parse_json_from_response(text):
     m = re.search(r'```(?:json)?\s*\n?(.*?)\n?```', text, re.DOTALL)
@@ -120,7 +235,6 @@ def parse_json_from_response(text):
                 pass
     return {}
 
-
 # ============================================================
 # MENU ZIP
 # ============================================================
@@ -128,15 +242,12 @@ def process_menu_zip(zip_file):
     menu = {}
     if zip_file is None:
         return menu
-
     extract_dir = WORK_DIR / "menu_images"
     if extract_dir.exists():
         shutil.rmtree(extract_dir)
     extract_dir.mkdir(parents=True)
-
     with zipfile.ZipFile(zip_file, 'r') as zf:
         zf.extractall(extract_dir)
-
     for root, dirs, files in os.walk(extract_dir):
         rel = Path(root).relative_to(extract_dir)
         if str(rel).startswith(('__', '.')):
@@ -146,11 +257,9 @@ def process_menu_zip(zip_file):
                 continue
             if not fname.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
                 continue
-
             fpath = Path(root) / fname
             parts = fpath.relative_to(extract_dir).parts
             category = parts[-2] if len(parts) >= 2 else "Menu"
-
             stem = fpath.stem
             last_us = stem.rfind('_')
             if last_us > 0:
@@ -158,30 +267,24 @@ def process_menu_zip(zip_file):
                 try:
                     price = float(stem[last_us + 1:].strip())
                 except ValueError:
-                    name_part = stem.replace('_', ' ').strip()
-                    price = 0.0
+                    name_part = stem.replace('_', ' ').strip(); price = 0.0
             else:
-                name_part = stem.replace('_', ' ').strip()
-                price = 0.0
-
+                name_part = stem.replace('_', ' ').strip(); price = 0.0
             with open(fpath, 'rb') as f:
                 img_bytes = f.read()
             img_b64 = base64.b64encode(img_bytes).decode('utf-8')
-            ext = fpath.suffix.lower()
+            ext  = fpath.suffix.lower()
             mime = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.webp': 'image/webp'}
-
             if category not in menu:
                 menu[category] = []
             menu[category].append({
-                "name": name_part,
-                "price": price,
-                "image_base64": f"data:{mime.get(ext, 'image/png')};base64,{img_b64}",
+                "name": name_part, "price": price,
+                "image_base64": f"data:{mime.get(ext,'image/png')};base64,{img_b64}",
             })
     return menu
 
-
 # ============================================================
-# TEMPLATE ENGINE
+# MENU HTML BUILDERS
 # ============================================================
 def build_category_tabs(menu_data):
     return "\n      ".join(
@@ -189,7 +292,6 @@ def build_category_tabs(menu_data):
         for cat in menu_data.keys()
     )
 
-
 def build_menu_html(menu_data):
     html = ""
     for cat, items in menu_data.items():
@@ -197,146 +299,292 @@ def build_menu_html(menu_data):
         html += f'      <div class="menu-category-title">{cat}</div>\n'
         for item in items:
             price_str = f"${item['price']:.2f}" if item['price'] > 0 else ""
-            img_src = item.get('image_base64', '')
-            if img_src:
-                img_tag = f'<img class="menu-item-img" src="{img_src}" alt="{item["name"]}" loading="lazy">'
-            else:
-                img_tag = '<div class="menu-item-img" style="background:linear-gradient(135deg,var(--secondary),var(--primary));opacity:0.3;"></div>'
-            html += f'''      <div class="menu-item" data-item-name="{item['name']}">
-        {img_tag}
-        <div class="menu-item-info">
-          <div class="menu-item-name">{item['name']}</div>
-          <div class="menu-item-price">{price_str}</div>
-        </div>
-      </div>\n'''
+            img_src   = item.get('image_base64', '')
+            img_tag   = (f'<img class="menu-item-img" src="{img_src}" alt="{item["name"]}" loading="lazy">'
+                         if img_src else
+                         '<div class="menu-item-img" style="background:linear-gradient(135deg,var(--secondary),var(--primary));opacity:0.3;"></div>')
+            html += (current_menu_item_template
+                     .replace("{name}", item["name"])
+                     .replace("{price}", price_str)
+                     .replace("{img_tag}", img_tag)
+                     .replace("{img_src}", img_src)
+                     .replace("{description}", item.get("description", "")) + "\n")
         html += '    </div>\n'
     return html
 
-
+# ============================================================
+# TEMPLATE ENGINE
+# ============================================================
 def load_template(key):
     path = TEMPLATE_DIR / f"{key}.html"
     if not path.exists():
         path = TEMPLATE_DIR / "warm.html"
-    if not path.exists():
-        return "<html><body><h1>Template not found</h1></body></html>"
     return path.read_text(encoding='utf-8')
 
-
 def rebuild_html():
-    """Rebuild full HTML from template + current data. Single source of truth."""
     global current_html
-    s = current_site_info
-    address = s.get("address", "")
-
-    template = load_template(current_template_key)
-
-    replacements = {
-        "{{SHOP_NAME}}": s.get("shop_name", "My Shop"),
-        "{{TAGLINE}}": s.get("description", "Welcome!")[:80],
-        "{{DESCRIPTION}}": s.get("description", ""),
-        "{{PHONE}}": s.get("phone", ""),
-        "{{ADDRESS}}": address,
-        "{{ADDRESS_ENCODED}}": urllib.parse.quote_plus(address),
-        "{{HOURS}}": s.get("hours", ""),
-        "{{LOCATION_SHORT}}": address.split(",")[0] if "," in address else address[:30],
-        "{{PROMO_TITLE}}": s.get("promo_title", "Coming Soon"),
-        "{{PROMO_SUBTITLE}}": s.get("promo_subtitle", "Stay tuned for updates"),
-    }
+    template  = load_template(current_template_key)
+    home_html = "\n".join(b["html"] for b in page_blocks)
 
     html = template
-    for key, val in replacements.items():
-        html = html.replace(key, val)
-
+    html = html.replace("<!-- {{HOME_BLOCKS}} -->", home_html)
     html = html.replace("<!-- {{CATEGORY_TABS}} -->", build_category_tabs(current_menu_data))
-    html = html.replace("<!-- {{MENU_ITEMS}} -->", build_menu_html(current_menu_data))
-
-    # Inject poster if exists
-    poster = s.get("poster_base64", "")
-    if poster:
-        pat = r'<div class="promo-placeholder"[^>]*>.*?</div>'
-        repl = f'<img src="{poster}" alt="Poster" style="width:100%;height:200px;object-fit:cover;display:block;">'
-        html = re.sub(pat, repl, html, flags=re.DOTALL, count=1)
+    html = html.replace("<!-- {{MENU_ITEMS}} -->",    build_menu_html(current_menu_data))
+    html = html.replace("{{SHOP_NAME}}", current_site_info.get("shop_name", "My Shop"))
+
+    # Poster carousel injection (into any .promo-placeholder found in blocks)
+    posters = current_site_info.get("posters", [])
+    if posters:
+        if len(posters) == 1:
+            carousel_html = f'<img src="{posters[0]}" alt="Poster" style="width:100%;height:auto;display:block;">'
+        else:
+            slides = "\n".join(
+                f'<div class="ps-slide{"  ps-active" if i==0 else ""}"><img src="{p}" alt="Poster {i+1}" style="width:100%;height:auto;display:block;"></div>'
+                for i, p in enumerate(posters)
+            )
+            dots = "\n".join(
+                f'<span class="ps-dot{"  ps-dot-on" if i==0 else ""}" onclick="psGo({i})"></span>'
+                for i in range(len(posters))
+            )
+            carousel_html = (
+                f'<div class="ps-wrap">{slides}'
+                f'<button class="ps-btn ps-l" onclick="psMove(-1)">&#10094;</button>'
+                f'<button class="ps-btn ps-r" onclick="psMove(1)">&#10095;</button>'
+                f'<div class="ps-dots">{dots}</div></div>'
+            )
+        pat  = r'<div class="promo-placeholder"[^>]*>.*?</div>'
+        html = re.sub(pat, carousel_html, html, flags=re.DOTALL, count=1)
+        carousel_css = (
+            ".ps-wrap{position:relative;overflow:hidden;border-radius:var(--card-radius,12px);}"
+            ".ps-slide{display:none;}.ps-slide.ps-active{display:block;}"
+            ".ps-btn{position:absolute;top:50%;transform:translateY(-50%);background:rgba(0,0,0,0.45);"
+            "color:#fff;border:none;padding:10px 14px;font-size:18px;cursor:pointer;z-index:10;border-radius:6px;}"
+            ".ps-l{left:8px;}.ps-r{right:8px;}"
+            ".ps-dots{position:absolute;bottom:10px;width:100%;text-align:center;}"
+            ".ps-dot{display:inline-block;width:8px;height:8px;background:rgba(255,255,255,0.5);"
+            "border-radius:50%;margin:0 3px;cursor:pointer;}"
+            ".ps-dot.ps-dot-on{background:#fff;}"
+        )
+        carousel_js = (
+            "<script>(function(){var idx=0;"
+            "function show(n){var s=document.querySelectorAll('.ps-slide');"
+            "var d=document.querySelectorAll('.ps-dot');if(!s.length)return;"
+            "idx=(n+s.length)%s.length;"
+            "s.forEach(function(e){e.classList.remove('ps-active');});"
+            "d.forEach(function(e){e.classList.remove('ps-dot-on');});"
+            "s[idx].classList.add('ps-active');"
+            "if(d[idx])d[idx].classList.add('ps-dot-on');}"
+            "window.psMove=function(d){show(idx+d);};"
+            "window.psGo=function(n){show(n);};"
+            "})();</script>"
+        )
+        html = html.replace("</style>", f"\n/* Carousel */\n{carousel_css}\n</style>", 1)
+        html = html.replace("</body>", f"\n{carousel_js}\n</body>", 1)
 
-    # Inject custom CSS overrides
-    css = s.get("custom_css", "")
+    # Custom CSS overrides
+    css = current_site_info.get("custom_css", "")
     if css:
-        html = html.replace("</style>", f"\n/* Custom */\n{css}\n</style>")
-
-    # Apply any saved block overrides
-    for block_name, block_html in s.get("block_overrides", {}).items():
-        block_def = BLOCKS.get(block_name)
-        if block_def and block_def["start"] and block_def["end"]:
-            start = block_def["start"]
-            end = block_def["end"]
-            pat = re.escape(start) + r'.*?' + re.escape(end)
-            replacement = start + "\n" + block_html + "\n" + end
-            html = re.sub(pat, replacement, html, flags=re.DOTALL, count=1)
+        html = html.replace("</style>", f"\n/* Custom */\n{css}\n</style>", 1)
 
     current_html = html
     return html
 
+# ============================================================
+# BLOCK MANAGEMENT HANDLERS
+# ============================================================
+def _make_menu_preview_html():
+    if not current_menu_data:
+        return BLOCK_DEFAULTS["Menu Preview"]
+    preview_items = []
+    for items in current_menu_data.values():
+        for item in items:
+            preview_items.append(item)
+            if len(preview_items) >= 3:
+                break
+        if len(preview_items) >= 3:
+            break
+    rows = ""
+    for item in preview_items:
+        price_str = f"${item['price']:.2f}" if item.get('price', 0) > 0 else ""
+        img_src = item.get('image_base64', '')
+        img_html = (f'<img src="{img_src}" alt="{item["name"]}" style="width:56px;height:56px;border-radius:10px;object-fit:cover;flex-shrink:0;">'
+                    if img_src else
+                    '<div style="width:56px;height:56px;border-radius:10px;background:linear-gradient(135deg,var(--secondary),var(--primary));opacity:0.4;flex-shrink:0;"></div>')
+        rows += f"""    <div style="display:flex;align-items:center;gap:12px;padding:10px 0;border-bottom:1px solid var(--border);">
+      {img_html}
+      <div>
+        <div style="font-size:13px;font-weight:600;color:var(--text);">{item['name']}</div>
+        <div style="font-size:13px;color:var(--primary);font-weight:600;">{price_str}</div>
+      </div>
+    </div>\n"""
+    return f"""\
+  <div class="section-title">
+    <h2>Our Menu</h2>
+    <div class="line"></div>
+  </div>
+  <div style="margin:0 24px;padding:0 20px;background:var(--bg-card);border-radius:var(--card-radius);border:1px solid var(--border);">
+{rows}    <div style="padding:14px 0;text-align:center;">
+      <button onclick="switchPage('menu')" style="background:var(--primary);color:var(--bg);border:none;padding:10px 28px;border-radius:100px;font-size:13px;font-weight:600;cursor:pointer;">View Full Menu →</button>
+    </div>
+  </div>"""
+
+def handle_add_block(block_type, selected_label):
+    if block_type == "Menu Preview":
+        initial_html = _make_menu_preview_html()
+    else:
+        initial_html = BLOCK_DEFAULTS.get(block_type, "")
+    new_block = {"id": _uid(), "type": block_type, "html": initial_html}
+    if selected_label and _find_block(selected_label):
+        idx = next((i for i, b in enumerate(page_blocks) if _block_label(b) == selected_label), -1)
+        page_blocks.insert(idx + 1, new_block)
+    else:
+        page_blocks.append(new_block)
+    rebuild_html()
+    choices   = get_block_choices()
+    new_label = _block_label(new_block)
+    return preview(current_html), gr.update(choices=choices, value=new_label)
+
+def handle_remove_block(selected_label):
+    global page_blocks
+    if not selected_label:
+        return preview(current_html), gr.update()
+    page_blocks = [b for b in page_blocks if _block_label(b) != selected_label]
+    rebuild_html()
+    choices = get_block_choices()
+    return preview(current_html), gr.update(choices=choices, value=choices[0] if choices else None)
+
+def handle_move_up(selected_label):
+    idx = next((i for i, b in enumerate(page_blocks) if _block_label(b) == selected_label), -1)
+    if idx > 0:
+        page_blocks[idx - 1], page_blocks[idx] = page_blocks[idx], page_blocks[idx - 1]
+        rebuild_html()
+    return preview(current_html), gr.update(choices=get_block_choices(), value=selected_label)
+
+def handle_move_down(selected_label):
+    idx = next((i for i, b in enumerate(page_blocks) if _block_label(b) == selected_label), -1)
+    if 0 <= idx < len(page_blocks) - 1:
+        page_blocks[idx], page_blocks[idx + 1] = page_blocks[idx + 1], page_blocks[idx]
+        rebuild_html()
+    return preview(current_html), gr.update(choices=get_block_choices(), value=selected_label)
 
 # ============================================================
-# BLOCK EXTRACTION
+# LLM BLOCK EDITOR
 # ============================================================
-def extract_block(html, block_name):
-    """Extract HTML content between block markers."""
-    block_def = BLOCKS.get(block_name)
-    if not block_def or not block_def["start"]:
-        return None
-    start = block_def["start"]
-    end = block_def["end"]
-    pat = re.escape(start) + r'\n?(.*?)\n?' + re.escape(end)
-    m = re.search(pat, html, re.DOTALL)
-    return m.group(1) if m else None
+CODER_SYSTEM = """You are a frontend developer editing a mobile website HTML block.
+
+RULES:
+- Output ONLY the modified HTML. No explanation, no markdown fences, no ```html.
+- Keep the mobile-friendly layout and existing CSS classes.
+- Only change what the instruction asks.
+- Do NOT output <html>, <head>, <body>, or <style> tags — only the inner block content.
+- You may add inline styles for visual adjustments.
+- Preserve existing CSS class names and structure unless explicitly asked to change them."""
+
+def edit_block_with_llm(selected_label, instruction, chat_history):
+    chat_history = chat_history or []
+
+    if not current_html:
+        chat_history.append((None, "⚠️ Generate a website first in the Create tab."))
+        return chat_history, preview(current_html)
 
+    if not selected_label:
+        chat_history.append((None, "⚠️ Select a block from the list first."))
+        return chat_history, preview(current_html)
 
-def replace_block(html, block_name, new_content):
-    """Replace a block's content in HTML."""
-    block_def = BLOCKS.get(block_name)
-    if not block_def or not block_def["start"]:
-        return html
-    start = block_def["start"]
-    end = block_def["end"]
-    pat = re.escape(start) + r'.*?' + re.escape(end)
-    replacement = start + "\n" + new_content + "\n" + end
-    return re.sub(pat, replacement, html, flags=re.DOTALL, count=1)
+    block = _find_block(selected_label)
+    if not block:
+        chat_history.append((None, "❌ Block not found."))
+        return chat_history, preview(current_html)
 
+    user_msg = f"[{block['type']}] {instruction}"
+    prompt = f"""Current HTML block:
+{block['html']}
+
+Instruction: {instruction}
+
+Output the modified HTML block:"""
+
+    raw = ollama_chat(QWEN_CODER_MODEL, CODER_SYSTEM, prompt, temperature=0.3)
+
+    if raw.startswith("ERROR:"):
+        chat_history.append((user_msg, f"❌ {raw}"))
+        return chat_history, preview(current_html)
+
+    new_html = raw.strip()
+    new_html = re.sub(r'^```html?\s*\n?', '', new_html)
+    new_html = re.sub(r'\n?```\s*$', '', new_html)
+
+    if '<!DOCTYPE' in new_html or '<html' in new_html.lower():
+        chat_history.append((user_msg, "❌ Model returned full page. Try a simpler instruction."))
+        return chat_history, preview(current_html)
+
+    block["html"] = new_html
+    rebuild_html()
+    chat_history.append((user_msg, f"✅ {block['type']} updated!"))
+    return chat_history, preview(current_html)
 
 # ============================================================
-# INTENT PARSING - Qwen 2.5 → JSON
+# CUSTOM BLOCK GENERATOR
 # ============================================================
-PARSE_SYSTEM_CONTACT = """Parse the user's instruction about contact/business info into JSON.
+CUSTOM_BLOCK_SYSTEM = """You are generating a new HTML block for a mobile-first restaurant/shop website.
+Generate a single self-contained HTML block based on the user's description.
+
+Available CSS variables (already in the page):
+  --primary, --secondary, --accent
+  --bg, --bg-card, --bg-elevated
+  --text, --text-light, --text-muted
+  --border, --card-radius, --card-shadow
+
+Available CSS classes:
+  .section-title h2   section header with decorative line
+  .about-card         padded content card
+  .contact-section    list container
+  .contact-item       row with icon + text
+  .pill               small rounded badge/tag
+  .spacer-lg          bottom spacer
 
-Output format:
-{"actions": [{"field": "...", "value": "..."}]}
-
-Available fields: shop_name, description, phone, address, hours, promo_title, promo_subtitle
+RULES:
+- Output ONLY the HTML. No explanation, no markdown fences, no ```html.
+- Mobile-friendly layout (max ~480px wide).
+- Use CSS variables for all colors so it works with light and dark themes.
+- Do NOT output <html>, <head>, <body>, or <style> tags."""
 
-Examples:
-- "Change phone to 555-1234" → {"actions":[{"field":"phone","value":"555-1234"}]}
-- "Update hours to Mon-Sat 8am-9pm" → {"actions":[{"field":"hours","value":"Mon-Sat 8am-9pm"}]}
-- "Change shop name to Blue Bottle" → {"actions":[{"field":"shop_name","value":"Blue Bottle"}]}
+def generate_custom_block(description, chat_history):
+    chat_history = chat_history or []
+    if not current_html:
+        chat_history.append((None, "⚠️ Generate a website first."))
+        return chat_history, preview(current_html), gr.update()
 
-Output ONLY JSON."""
+    user_msg = f"[Custom Block] {description}"
+    raw = ollama_chat(QWEN_CODER_MODEL, CUSTOM_BLOCK_SYSTEM,
+                      f"Generate an HTML block for: {description}", temperature=0.4)
 
-PARSE_SYSTEM_MENU = """Parse the user's instruction about menu items into JSON.
+    if raw.startswith("ERROR:"):
+        chat_history.append((user_msg, f"❌ {raw}"))
+        return chat_history, preview(current_html), gr.update()
 
-Output format:
-{"actions": [{"op": "add|remove|change_price", "name": "...", "price": 0.0, "category": "..."}]}
+    new_html = raw.strip()
+    new_html = re.sub(r'^```html?\s*\n?', '', new_html)
+    new_html = re.sub(r'\n?```\s*$', '', new_html)
 
-Examples:
-- "Add Iced Mocha $6 in Coffee" → {"actions":[{"op":"add","name":"Iced Mocha","price":6.00,"category":"Coffee"}]}
-- "Remove Espresso" → {"actions":[{"op":"remove","name":"Espresso"}]}
-- "Change Latte price to $5.50" → {"actions":[{"op":"change_price","name":"Latte","price":5.50}]}
-- "Add Tiramisu $8 in Desserts and remove Muffin" → {"actions":[{"op":"add","name":"Tiramisu","price":8.00,"category":"Desserts"},{"op":"remove","name":"Muffin"}]}
+    if '<!DOCTYPE' in new_html or '<html' in new_html.lower():
+        chat_history.append((user_msg, "❌ Model returned full page. Try a simpler description."))
+        return chat_history, preview(current_html), gr.update()
 
-Output ONLY JSON."""
+    new_block = {"id": _uid(), "type": "Custom", "html": new_html}
+    page_blocks.append(new_block)
+    rebuild_html()
+    choices   = get_block_choices()
+    new_label = _block_label(new_block)
+    chat_history.append((user_msg, "✅ Custom block added! Select it to edit further."))
+    return chat_history, preview(current_html), gr.update(choices=choices, value=new_label)
 
+# ============================================================
+# STYLE EDITOR  (CSS variable injection, no block needed)
+# ============================================================
 PARSE_SYSTEM_STYLE = """Parse the user's instruction about visual style into JSON.
 
-Output format:
-{"actions": [{"prop": "...", "value": "..."}]}
+Output format: {"actions": [{"prop": "...", "value": "..."}]}
 
 Available props:
 - primary, secondary, accent, bg, text (CSS color values like #006400)
@@ -345,55 +593,200 @@ Available props:
 
 Examples:
 - "Change primary color to forest green" → {"actions":[{"prop":"primary","value":"#228B22"}]}
-- "Make it dark theme, black background" → {"actions":[{"prop":"bg","value":"#1a1a1a"},{"prop":"text","value":"#f0f0f0"}]}
+- "Dark theme, black background" → {"actions":[{"prop":"bg","value":"#1a1a1a"},{"prop":"text","value":"#f0f0f0"}]}
 - "Use Georgia for headings" → {"actions":[{"prop":"font_heading","value":"Georgia, serif"}]}
 
 Output ONLY JSON."""
 
-CODER_BLOCK_SYSTEM = """You are a frontend developer. You will receive a small HTML block from a mobile website and an instruction to modify it.
+def handle_style_edit(instruction, chat_history):
+    chat_history = chat_history or []
+    user_msg = f"[Style] {instruction}"
 
-RULES:
-- Output ONLY the modified HTML block. No explanations, no markdown code blocks, no ```html.
-- Keep the same HTML structure and CSS classes.
-- Only change content/styling as requested.
-- Do NOT add <html>, <head>, <body>, or <style> tags. Only output the inner block content.
-- Keep it concise and mobile-friendly.
-- You may add inline styles if needed for visual changes.
-- Preserve all existing CSS classes and data attributes."""
+    raw = ollama_chat(QWEN_MODEL, PARSE_SYSTEM_STYLE, instruction)
+    if raw.startswith("ERROR:"):
+        chat_history.append((user_msg, f"❌ {raw}"))
+        return chat_history, preview(current_html)
 
+    parsed  = parse_json_from_response(raw)
+    actions = parsed.get("actions", [])
+    if not actions:
+        chat_history.append((user_msg, f"❌ Could not parse. Raw: {raw[:150]}"))
+        return chat_history, preview(current_html)
 
-# ============================================================
-# PYTHON EXECUTORS for structured operations
-# ============================================================
-def execute_contact_actions(actions):
-    """Update site_info fields."""
-    messages = []
+    css_var_map = {
+        "primary": "--primary", "secondary": "--secondary",
+        "accent": "--accent",   "bg": "--bg", "text": "--text",
+        "card_radius": "--card-radius",
+    }
+    lines, messages = [], []
     for a in actions:
-        field = a.get("field", "")
-        value = a.get("value", "")
-        if field in current_site_info:
-            current_site_info[field] = value
-            messages.append(f"Updated {field}")
+        prop, value = a.get("prop", ""), a.get("value", "")
+        if prop == "font_heading":
+            lines.append(f".hero h1, .section-title h2, .menu-header h1, .menu-category-title {{ font-family: {value} !important; }}")
+            messages.append(f"Changed font_heading")
+        elif prop == "font_body":
+            lines.append(f"body {{ font-family: {value} !important; }}")
+            messages.append(f"Changed font_body")
+        elif prop in css_var_map:
+            lines.append(f":root {{ {css_var_map[prop]}: {value}; }}")
+            messages.append(f"Changed {prop} → {value}")
         else:
-            messages.append(f"Unknown field: {field}")
+            messages.append(f"Unknown prop: {prop}")
+
+    if lines:
+        current_site_info["custom_css"] = current_site_info.get("custom_css", "") + "\n" + "\n".join(lines)
     rebuild_html()
-    return messages
+    chat_history.append((user_msg, "🎨 " + " · ".join(messages) + "\n\n✅ Done!"))
+    return chat_history, preview(current_html)
+
+# ============================================================
+# CREATE WEBSITE
+# ============================================================
+def _make_hero_html(shop_name, desc, hours, addr):
+    location_short = addr.split(",")[0][:25] if "," in addr else addr[:25]
+    return f"""\
+  <div class="hero">
+    <div class="hero-badge">
+      <svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="10"/></svg>
+      Est. 2024
+    </div>
+    <h1>{shop_name}</h1>
+    <p class="hero-tagline">{desc[:80]}</p>
+  </div>
+  <div class="info-pills">
+    <div class="pill">
+      <svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="10"/><polyline points="12 6 12 12 16 14"/></svg>
+      <span>{hours or 'Open Daily'}</span>
+    </div>
+    <div class="pill">
+      <svg viewBox="0 0 24 24"><path d="M21 10c0 7-9 13-9 13s-9-6-9-13a9 9 0 0 1 18 0z"/><circle cx="12" cy="10" r="3"/></svg>
+      <span>{location_short or 'Visit Us'}</span>
+    </div>
+  </div>"""
+
+def _make_about_html(desc):
+    return f"""\
+  <div class="section-title">
+    <h2>Our Story</h2>
+    <div class="line"></div>
+  </div>
+  <div class="about-card">
+    <p>{desc}</p>
+  </div>"""
+
+def _make_contact_html(phone, addr, hours):
+    addr_enc = urllib.parse.quote_plus(addr)
+    return f"""\
+  <div class="section-title">
+    <h2>Find Us</h2>
+    <div class="line"></div>
+  </div>
+  <div class="contact-section">
+    <a class="contact-item" href="tel:{phone}">
+      <div class="contact-icon">
+        <svg viewBox="0 0 24 24"><path d="M22 16.92v3a2 2 0 0 1-2.18 2 19.79 19.79 0 0 1-8.63-3.07 19.5 19.5 0 0 1-6-6 19.79 19.79 0 0 1-3.07-8.67A2 2 0 0 1 4.11 2h3a2 2 0 0 1 2 1.72c.127.96.361 1.903.7 2.81a2 2 0 0 1-.45 2.11L8.09 9.91a16 16 0 0 0 6 6l1.27-1.27a2 2 0 0 1 2.11-.45c.907.339 1.85.573 2.81.7A2 2 0 0 1 22 16.92z"/></svg>
+      </div>
+      <div class="contact-text"><div class="label">Phone</div><div class="value">{phone or '—'}</div></div>
+    </a>
+    <a class="contact-item" href="https://maps.google.com/?q={addr_enc}" target="_blank">
+      <div class="contact-icon">
+        <svg viewBox="0 0 24 24"><path d="M21 10c0 7-9 13-9 13s-9-6-9-13a9 9 0 0 1 18 0z"/><circle cx="12" cy="10" r="3"/></svg>
+      </div>
+      <div class="contact-text"><div class="label">Address</div><div class="value">{addr or '—'}</div></div>
+    </a>
+    <div class="contact-item">
+      <div class="contact-icon">
+        <svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="10"/><polyline points="12 6 12 12 16 14"/></svg>
+      </div>
+      <div class="contact-text"><div class="label">Hours</div><div class="value">{hours or '—'}</div></div>
+    </div>
+  </div>
+  <div class="spacer-lg"></div>"""
+
+def create_website(shop_name, desc, phone, addr, hours, style, menu_zip, progress=gr.Progress()):
+    global page_blocks, current_menu_data, current_site_info, current_template_key, current_menu_item_template
+    current_menu_item_template = MENU_ITEM_TEMPLATE_DEFAULT  # reset on each new site
+
+    progress(0.1, desc="Processing menu...")
+    current_menu_data = process_menu_zip(menu_zip)
+    current_site_info = {
+        "shop_name": shop_name or "My Shop",
+        "description": desc or "Welcome to our shop!",
+        "phone": phone or "", "address": addr or "", "hours": hours or "",
+        "custom_css": "", "posters": [], "poster_base64": "",
+    }
+    current_template_key = {"Warm & Cozy": "warm", "Dark & Elegant": "dark"}.get(style, "warm")
 
+    progress(0.4, desc="Building blocks...")
+    page_blocks = [
+        {"id": _uid(), "type": "Hero Banner",   "html": _make_hero_html(shop_name or "My Shop", desc or "Welcome!", hours or "", addr or "")},
+        {"id": _uid(), "type": "Promo / Event", "html": BLOCK_DEFAULTS["Promo / Event"]},
+        {"id": _uid(), "type": "About / Story", "html": _make_about_html(desc or "Welcome to our shop!")},
+        {"id": _uid(), "type": "Contact Info",  "html": _make_contact_html(phone or "", addr or "", hours or "")},
+    ]
+
+    progress(0.8, desc="Rendering...")
+    rebuild_html()
+
+    n_items = sum(len(v) for v in current_menu_data.values())
+    n_cats  = len(current_menu_data)
+    progress(1.0)
+    print(f"[create_website] page_blocks={len(page_blocks)}, html_len={len(current_html)}")
+    return preview(current_html), f"✅ {n_items} items · {n_cats} categories"
+
+# ============================================================
+# MENU ITEM DATA EDITOR  (add / remove / change price)
+# ============================================================
+PARSE_SYSTEM_MENU = """Parse the user's instruction about menu items into JSON.
+
+Output format:
+{"actions": [{"op": "add|remove|change_price", "name": "...", "price": 0.0, "category": "..."}]}
+
+Examples:
+- "Add Iced Mocha $6 in Coffee" → {"actions":[{"op":"add","name":"Iced Mocha","price":6.00,"category":"Coffee"}]}
+- "Remove Espresso" → {"actions":[{"op":"remove","name":"Espresso"}]}
+- "Change Latte price to $5.50" → {"actions":[{"op":"change_price","name":"Latte","price":5.50}]}
+
+Output ONLY JSON."""
+
+def handle_menu_data_edit(instruction, uploaded_image, chat_history):
+    chat_history = chat_history or []
+    if not current_html:
+        chat_history.append((None, "⚠️ Generate a website first."))
+        return chat_history, preview(current_html)
+
+    user_msg = f"[Menu Items] {instruction}"
+    raw = ollama_chat(QWEN_MODEL, PARSE_SYSTEM_MENU, instruction)
+    if raw.startswith("ERROR:"):
+        chat_history.append((user_msg, f"❌ {raw}"))
+        return chat_history, preview(current_html)
+
+    parsed  = parse_json_from_response(raw)
+    actions = parsed.get("actions", [])
+    if not actions:
+        chat_history.append((user_msg, f"❌ Could not parse. Raw: {raw[:150]}"))
+        return chat_history, preview(current_html)
+
+    # Handle optional uploaded image
+    img_b64 = None
+    if uploaded_image:
+        try:
+            with open(uploaded_image, 'rb') as f:
+                img_b64 = f"data:image/png;base64,{base64.b64encode(f.read()).decode()}"
+        except Exception:
+            pass
 
-def execute_menu_actions(actions, uploaded_image_b64=None):
-    """Add/remove/modify menu items in current_menu_data."""
     messages = []
     for a in actions:
-        op = a.get("op", "")
+        op   = a.get("op", "")
         name = a.get("name", "")
 
         if op == "add":
-            cat = a.get("category", "Menu")
+            cat   = a.get("category", "Menu")
             price = a.get("price", 0.0)
             if cat not in current_menu_data:
                 current_menu_data[cat] = []
-            new_item = {"name": name, "price": price, "image_base64": uploaded_image_b64 or ""}
-            current_menu_data[cat].append(new_item)
+            current_menu_data[cat].append({"name": name, "price": price, "image_base64": img_b64 or ""})
             messages.append(f"Added {name} (${price:.2f}) to {cat}")
 
         elif op == "remove":
@@ -401,14 +794,10 @@ def execute_menu_actions(actions, uploaded_image_b64=None):
             for cat, items in current_menu_data.items():
                 for i, item in enumerate(items):
                     if item["name"].lower() == name.lower():
-                        items.pop(i)
-                        messages.append(f"Removed {name} from {cat}")
-                        found = True
-                        break
-                if found:
-                    break
+                        items.pop(i); messages.append(f"Removed {name}"); found = True; break
+                if found: break
             if not found:
-                messages.append(f"Item '{name}' not found")
+                messages.append(f"'{name}' not found")
 
         elif op == "change_price":
             price = a.get("price", 0.0)
@@ -416,177 +805,133 @@ def execute_menu_actions(actions, uploaded_image_b64=None):
             for cat, items in current_menu_data.items():
                 for item in items:
                     if item["name"].lower() == name.lower():
-                        item["price"] = price
-                        messages.append(f"Changed {name} price to ${price:.2f}")
-                        found = True
-                        break
-                if found:
-                    break
+                        item["price"] = price; messages.append(f"{name} → ${price:.2f}"); found = True; break
+                if found: break
             if not found:
-                messages.append(f"Item '{name}' not found")
-
-    rebuild_html()
-    return messages
-
-
-def execute_style_actions(actions):
-    """Modify CSS variables via custom_css injection."""
-    css_var_map = {
-        "primary": "--primary",
-        "secondary": "--secondary",
-        "accent": "--accent",
-        "bg": "--bg",
-        "text": "--text",
-        "font_heading": None,  # special handling
-        "font_body": None,
-        "card_radius": "--card-radius",
-    }
-    lines = []
-    messages = []
-    for a in actions:
-        prop = a.get("prop", "")
-        value = a.get("value", "")
-
-        if prop in ("font_heading", "font_body"):
-            # Override via more specific CSS
-            if prop == "font_heading":
-                lines.append(f".hero h1, .section-title h2, .menu-header h1, .menu-category-title, .menu-item-price {{ font-family: {value} !important; }}")
-            else:
-                lines.append(f"body, .menu-item-name, .about-card p, .contact-text .value {{ font-family: {value} !important; }}")
-            messages.append(f"Changed {prop} to {value}")
-        elif prop in css_var_map and css_var_map[prop]:
-            lines.append(f":root {{ {css_var_map[prop]}: {value}; }}")
-            messages.append(f"Changed {prop} to {value}")
-        else:
-            messages.append(f"Unknown style prop: {prop}")
-
-    if lines:
-        existing = current_site_info.get("custom_css", "")
-        current_site_info["custom_css"] = existing + "\n" + "\n".join(lines)
+                messages.append(f"'{name}' not found")
 
     rebuild_html()
-    return messages
-
+    extra = " (with photo)" if img_b64 else ""
+    chat_history.append((user_msg, "🔧 " + " · ".join(messages) + extra + "\n\n✅ Done!"))
+    return chat_history, preview(current_html)
 
 # ============================================================
-# CODER BLOCK EDITING - Qwen Coder on isolated block only
+# MENU STRUCTURE EDITOR
 # ============================================================
-def edit_block_with_coder(block_name, user_instruction):
-    """Extract block → send to Qwen Coder with instruction → replace block."""
-    global current_html
-
-    block_html = extract_block(current_html, block_name)
-    if block_html is None:
-        return f"❌ Could not find block: {block_name}"
+MENU_STRUCTURE_SYSTEM = """You are editing the HTML template for a single menu item card on a mobile website.
+
+Available placeholders:
+  {name}        — item name text  [REQUIRED — must keep]
+  {price}       — item price text e.g. "$4.50"  [REQUIRED — must keep]
+  {img_tag}     — full <img> element with class="menu-item-img" (76×76 px square)
+  {img_src}     — raw image URL/base64 string — MUST be used as: <img src="{img_src}" ...>  NEVER put {img_src} as text content inside a div or span
+  {description} — per-item text description string (may be empty)  [use ONLY when adding a prose text description field]
+
+IMAGE RULES (CRITICAL):
+- {img_src} is a raw URL string. You MUST embed it like this: <img src="{img_src}" style="..." alt="{name}">
+- NEVER write just {img_src} alone as element content — it will render as raw text on the page
+- When changing image layout/size, use {img_src} with a custom <img> tag and omit {img_tag}
+
+MULTI-COLUMN LAYOUT RULES:
+- To make items display in 2 columns, add a <style> tag at the TOP of the template:
+  <style>.menu-category { display:grid; grid-template-columns:1fr 1fr; gap:12px; } .menu-category-title { grid-column:1/-1; }</style>
+- Then make the item itself a compact vertical block (display:block, not flex row)
+
+GENERAL RULES:
+- Output ONLY the modified HTML template. No explanation, no markdown fences, no ```html.
+- Keep {name}, {price}, and at least one of {img_tag} or {img_src}
+- For UI decorations (stars, spice icons, badges): write them as static HTML/emoji, NOT as a placeholder
+- Use CSS variables for colors (--primary, --bg-card, --text, --text-muted, --border)
+- Do NOT add <html>, <head>, or <body> tags. A single <style> tag at the top is allowed for layout.
+- Use inline styles freely to override any layout constraints"""
+
+def edit_menu_structure(instruction, chat_history):
+    global current_menu_item_template
+    chat_history = chat_history or []
 
-    prompt = f"""Here is the current HTML block for "{block_name}":
+    if not current_html:
+        chat_history.append((None, "⚠️ Generate a website first."))
+        return chat_history, preview(current_html)
 
-{block_html}
+    user_msg = f"[Menu Structure] {instruction}"
+    prompt = f"""Current menu item template:
+{current_menu_item_template}
 
-User instruction: {user_instruction}
+Instruction: {instruction}
 
-Output the modified HTML block:"""
+Output the modified template (keep {{name}}, {{price}}, {{img_tag}} placeholders):"""
 
-    raw = ollama_chat(QWEN_CODER_MODEL, CODER_BLOCK_SYSTEM, prompt, temperature=0.3)
+    raw = ollama_chat(QWEN_CODER_MODEL, MENU_STRUCTURE_SYSTEM, prompt, temperature=0.3)
 
     if raw.startswith("ERROR:"):
-        return f"❌ {raw}"
-
-    # Clean response
-    new_block = raw.strip()
-    new_block = re.sub(r'^```html?\s*\n?', '', new_block)
-    new_block = re.sub(r'\n?```\s*$', '', new_block)
+        chat_history.append((user_msg, f"❌ {raw}"))
+        return chat_history, preview(current_html)
 
-    # Sanity check: block shouldn't contain full document tags
-    if '<!DOCTYPE' in new_block or '<html' in new_block.lower():
-        return "❌ Coder returned a full page instead of a block. Try again with a simpler instruction."
+    new_template = raw.strip()
+    new_template = re.sub(r'^```html?\s*\n?', '', new_template)
+    new_template = re.sub(r'\n?```\s*$', '', new_template)
+
+    # Auto-fix: {img_src} as bare text content (not inside any attribute or CSS url())
+    # Valid uses: src="{img_src}", url({img_src}), url('{img_src}'), url("{img_src}")
+    # Invalid (shows raw base64 on page): <div>{img_src}</div>
+    if "{img_src}" in new_template and not re.search(
+        r'(?:src\s*=\s*["\']?|url\s*\(\s*["\']?)\{img_src\}', new_template
+    ):
+        new_template = new_template.replace(
+            "{img_src}",
+            '<img src="{img_src}" style="width:100%;height:160px;object-fit:cover;border-radius:8px;" alt="{name}">'
+        )
 
-    # Save override and rebuild
-    if "block_overrides" not in current_site_info:
-        current_site_info["block_overrides"] = {}
-    current_site_info["block_overrides"][block_name] = new_block
+    # Validate required placeholders — {img_tag} can be replaced by {img_src}
+    missing = [p for p in ["{name}", "{price}"] if p not in new_template]
+    if "{img_tag}" not in new_template and "{img_src}" not in new_template:
+        missing.append("{img_tag} or {img_src}")
+    if missing:
+        chat_history.append((user_msg, f"❌ Model dropped placeholders: {missing}. Try again."))
+        return chat_history, preview_menu(current_html)
+
+    current_menu_item_template = new_template
+
+    # If {description} was added and items don't have descriptions yet, auto-generate them
+    if "{description}" in new_template:
+        items_needing_desc = [
+            item for items in current_menu_data.values() for item in items
+            if not item.get("description")
+        ]
+        if items_needing_desc:
+            chat_history.append((user_msg, "⏳ Generating item descriptions..."))
+            all_names = [item["name"] for item in items_needing_desc]
+            prompt = f"Generate descriptions for these menu items: {json.dumps(all_names)}"
+            raw2 = ollama_chat(QWEN_MODEL, DESCRIBE_SYSTEM, prompt, temperature=0.7)
+            parsed2 = parse_json_from_response(raw2)
+            desc_map = {d["name"]: d["description"] for d in parsed2.get("items", []) if "name" in d and "description" in d}
+            for items in current_menu_data.values():
+                for item in items:
+                    if item["name"] in desc_map:
+                        item["description"] = desc_map[item["name"]]
+            count = len(desc_map)
+            chat_history.append((None, f"✅ Structure updated + {count} descriptions generated!"))
+        else:
+            chat_history.append((user_msg, "✅ Menu item structure updated!"))
+    else:
+        chat_history.append((user_msg, "✅ Menu item structure updated!"))
 
     rebuild_html()
-    return f"✅ {block_name} updated!"
-
-
-# ============================================================
-# MAIN EDIT HANDLER
-# ============================================================
-def handle_edit(block_choice, user_message, uploaded_image, chat_history):
-    """Route edit to the correct handler based on block type."""
-    global current_html
-
-    chat_history = chat_history or []
-
-    if not current_html:
-        chat_history.append({"role": "user", "content": f"[{block_choice}] {user_message}"})
-        chat_history.append({"role": "assistant", "content": "⚠️ Generate a website first in the Create tab."})
-        return chat_history, preview(current_html)
+    return chat_history, preview_menu(current_html)
 
-    chat_history.append({"role": "user", "content": f"[{block_choice}] {user_message}"})
-
-    block_def = BLOCKS[block_choice]
-
-    # ---- Python-handled blocks ----
-    if block_choice == "Contact Info":
-        raw = ollama_chat(QWEN_MODEL, PARSE_SYSTEM_CONTACT, user_message)
-        if raw.startswith("ERROR:"):
-            chat_history.append({"role": "assistant", "content": f"❌ {raw}"})
-            return chat_history, preview(current_html)
-        parsed = parse_json_from_response(raw)
-        actions = parsed.get("actions", [])
-        if not actions:
-            chat_history.append({"role": "assistant", "content": f"❌ Could not parse. Raw: {raw[:150]}"})
-            return chat_history, preview(current_html)
-        results = execute_contact_actions(actions)
-        chat_history.append({"role": "assistant", "content": "🔧 " + " · ".join(results) + "\n\n✅ Done!"})
-
-    elif block_choice == "Menu Items":
-        raw = ollama_chat(QWEN_MODEL, PARSE_SYSTEM_MENU, user_message)
-        if raw.startswith("ERROR:"):
-            chat_history.append({"role": "assistant", "content": f"❌ {raw}"})
-            return chat_history, preview(current_html)
-        parsed = parse_json_from_response(raw)
-        actions = parsed.get("actions", [])
-        if not actions:
-            chat_history.append({"role": "assistant", "content": f"❌ Could not parse. Raw: {raw[:150]}"})
-            return chat_history, preview(current_html)
-
-        # Handle uploaded image
-        img_b64 = None
-        if uploaded_image:
-            try:
-                with open(uploaded_image, 'rb') as f:
-                    img_b64 = f"data:image/png;base64,{base64.b64encode(f.read()).decode()}"
-            except Exception:
-                pass
+DESCRIBE_SYSTEM = """Generate short, appetising descriptions for menu items.
+Given a list of item names, output ONLY JSON:
+{"items": [{"name": "...", "description": "..."}]}
+Each description: 1 sentence, ≤12 words, mouth-watering and relevant to the item.
+Output ONLY JSON."""
 
-        results = execute_menu_actions(actions, img_b64)
-        extra = " (with photo)" if img_b64 else ""
-        chat_history.append({"role": "assistant", "content": "🔧 " + " · ".join(results) + extra + "\n\n✅ Done!"})
-
-    elif block_choice == "Style / Colors":
-        raw = ollama_chat(QWEN_MODEL, PARSE_SYSTEM_STYLE, user_message)
-        if raw.startswith("ERROR:"):
-            chat_history.append({"role": "assistant", "content": f"❌ {raw}"})
-            return chat_history, preview(current_html)
-        parsed = parse_json_from_response(raw)
-        actions = parsed.get("actions", [])
-        if not actions:
-            chat_history.append({"role": "assistant", "content": f"❌ Could not parse. Raw: {raw[:150]}"})
-            return chat_history, preview(current_html)
-        results = execute_style_actions(actions)
-        chat_history.append({"role": "assistant", "content": "🎨 " + " · ".join(results) + "\n\n✅ Done!"})
-
-    # ---- Coder-handled blocks ----
-    elif block_def["mode"] == "coder":
-        chat_history.append({"role": "assistant", "content": f"🔧 Sending to Coder: editing {block_choice}..."})
-        result = edit_block_with_coder(block_choice, user_message)
-        chat_history.append({"role": "assistant", "content": result})
-
-    return chat_history, preview(current_html)
 
+def reset_menu_structure(chat_history):
+    global current_menu_item_template
+    current_menu_item_template = MENU_ITEM_TEMPLATE_DEFAULT
+    rebuild_html()
+    chat_history = (chat_history or []) + [(None, "↩️ Menu structure reset to default.")]
+    return chat_history, preview_menu(current_html)
 
 # ============================================================
 # POSTER
@@ -597,7 +942,6 @@ def load_sd_model():
         return sd_pipe
     if not SD_AVAILABLE:
         return None
-    print("Loading SD Turbo...")
     sd_pipe = AutoPipelineForText2Image.from_pretrained(
         SD_MODEL_ID,
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
@@ -608,42 +952,173 @@ def load_sd_model():
         sd_pipe.enable_attention_slicing()
     return sd_pipe
 
+PARSE_SYSTEM_POSTER = """Extract key promotional text from the user's description for a poster.
+Output ONLY JSON: {"title": "...", "offer": "...", "detail": "...", "tagline": "..."}
+- title: main headline (2-5 words)
+- offer: the key offer (e.g. "20% OFF")
+- detail: conditions or date (e.g. "Valid Mon-Fri")
+- tagline: short catchy phrase
+Output ONLY JSON."""
 
-def gen_poster_prompt(desc):
-    sys = """Generate a Stable Diffusion prompt for a poster background.
-ONLY output the prompt. Under 50 words. No text/letters in image. Focus on mood, color, atmosphere. Add "high quality, professional, 4k"."""
-    raw = ollama_chat(QWEN_MODEL, sys, desc, temperature=0.7)
-    return raw.strip().strip('"\'') if not raw.startswith("ERROR:") else "beautiful poster background, professional, soft lighting, 4k"
-
-
-def generate_poster(desc, progress=gr.Progress()):
-    if not SD_AVAILABLE:
-        return None, "SD not installed."
-    progress(0.1, desc="Generating prompt...")
-    prompt = gen_poster_prompt(desc)
-    progress(0.3, desc="Loading SD Turbo...")
-    pipe = load_sd_model()
-    if not pipe:
-        return None, "Failed to load SD."
-    progress(0.5, desc="Generating...")
-    with torch.no_grad():
-        img = pipe(prompt=prompt, num_inference_steps=4, guidance_scale=0.0, width=512, height=768).images[0]
-    img.save(WORK_DIR / "latest_poster.png")
-    progress(1.0)
-    return img, f"**Prompt:** {prompt}"
+def extract_poster_info(desc):
+    raw = _ollama_chat(QWEN_MODEL, PARSE_SYSTEM_POSTER, desc, temperature=0.4)
+    if raw.startswith("ERROR:"):
+        return {"title": "Special Offer", "offer": "", "detail": "", "tagline": ""}
+    parsed = parse_json_from_response(raw)
+    return parsed if parsed else {"title": "Special Offer", "offer": "", "detail": "", "tagline": ""}
+
+def gen_sd_prompt(desc, bg_style):
+    if bg_style == "Restaurant atmosphere":
+        sys = "Generate a Stable Diffusion prompt for a restaurant/cafe poster background. Show warm interior, food, bokeh. Under 40 words, no quality tags."
+        fallback = "warm cozy restaurant interior, wooden table, bokeh lights, appetizing food, soft warm lighting"
+        suffix = "no text, no words, no people, professional food photography, 4k, bokeh"
+    else:
+        sys = "Generate a Stable Diffusion prompt that visually matches the promotion. Style: vibrant, professional. Under 40 words, no quality tags."
+        fallback = "happy people enjoying food, warm atmosphere, vibrant colors"
+        suffix = "no text, no words, high quality, 4k, professional photography"
+    raw = _ollama_chat(QWEN_MODEL, sys, desc, temperature=0.5)
+    if raw.startswith("ERROR:"):
+        raw = fallback
+    return raw.strip().strip('"\'') + ", " + suffix
+
+def compose_poster(bg_img, poster_info):
+    img = bg_img.copy().convert("RGBA")
+    w, h = img.size
+    overlay = Image.new("RGBA", (w, h), (0, 0, 0, 0))
+    draw_ov = ImageDraw.Draw(overlay)
+    for i in range(h // 3, h):
+        alpha = int(210 * (i - h // 3) / (h * 2 // 3))
+        draw_ov.line([(0, i), (w, i)], fill=(0, 0, 0, min(alpha, 210)))
+    img  = Image.alpha_composite(img, overlay)
+    draw = ImageDraw.Draw(img)
+
+    def get_font(size):
+        for fp in ["/System/Library/Fonts/Helvetica.ttc", "/System/Library/Fonts/Arial.ttf",
+                   "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
+                   "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf"]:
+            try:
+                return ImageFont.truetype(fp, size=size)
+            except Exception:
+                continue
+        return ImageFont.load_default()
+
+    def fit_font(text, base_size, min_size=12):
+        max_w = int(w * 0.88)
+        size  = base_size
+        while size >= min_size:
+            font = get_font(size)
+            bbox = draw.textbbox((0, 0), text, font=font)
+            if bbox[2] - bbox[0] <= max_w:
+                return font
+            size = max(min_size, int(size * 0.82))
+        return get_font(min_size)
+
+    def wrap_text(text, font, max_w):
+        words, lines, cur = text.split(), [], []
+        for word in words:
+            test = " ".join(cur + [word])
+            if draw.textbbox((0, 0), test, font=font)[2] <= max_w:
+                cur.append(word)
+            else:
+                if cur: lines.append(" ".join(cur))
+                cur = [word]
+        if cur: lines.append(" ".join(cur))
+        return lines or [text]
+
+    def draw_centered(text, font, y, color):
+        if not text: return 0
+        lines  = wrap_text(text, font, int(w * 0.88))
+        line_h = draw.textbbox((0, 0), "A", font=font)[3] + 4
+        for i, line in enumerate(lines):
+            tw = draw.textbbox((0, 0), line, font=font)[2]
+            draw.text(((w - tw) // 2, y + i * line_h), line, font=font, fill=color)
+        return line_h * len(lines)
+
+    y = h - int(h * 0.06)
+    if poster_info.get("detail"):
+        f = fit_font(poster_info["detail"], int(h * 0.040))
+        draw_centered(poster_info["detail"], f, y - int(h * 0.045), (200, 200, 200, 255))
+        y -= int(h * 0.07)
+    if poster_info.get("offer"):
+        f = fit_font(poster_info["offer"], int(h * 0.11))
+        draw_centered(poster_info["offer"], f, y - int(h * 0.11), (255, 215, 50, 255))
+        y -= int(h * 0.13)
+    if poster_info.get("title"):
+        f = fit_font(poster_info["title"], int(h * 0.065))
+        draw_centered(poster_info["title"], f, y - int(h * 0.07), (255, 255, 255, 255))
+        y -= int(h * 0.09)
+    if poster_info.get("tagline"):
+        f = fit_font(poster_info["tagline"], int(h * 0.038))
+        draw_centered(poster_info["tagline"], f, y - int(h * 0.04), (160, 230, 160, 255))
+    return img.convert("RGB")
+
+def _poster_status():
+    n = len(current_site_info.get("posters", []))
+    return f"🖼️ {n} poster{'s' if n != 1 else ''} in carousel" if n else "*No posters added yet*"
+
+@spaces.GPU
+def _generate_poster_gpu(desc, bg_prompt, uploaded_bg, bg_style):
+    """GPU-only: LLM + SD inference + compose. Returns (img, img_b64, info) or (None, None, error_msg)."""
+    poster_info = extract_poster_info(desc)
+    info = f"📝 **Extracted:** {poster_info.get('title','')} · {poster_info.get('offer','')} · {poster_info.get('detail','')}"
+
+    if uploaded_bg is not None:
+        try:
+            bg_img = Image.open(uploaded_bg).convert("RGB").resize((512, 768), Image.LANCZOS)
+            info += "\n\n🖼️ Using uploaded background"
+        except Exception as e:
+            return None, None, f"❌ Failed to load image: {e}"
+    else:
+        sd_prompt = bg_prompt.strip() if bg_prompt and bg_prompt.strip() else gen_sd_prompt(desc, bg_style)
+        pipe = load_sd_model()
+        if not pipe:
+            return None, None, "❌ Failed to load SD model."
+        with torch.no_grad():
+            bg_img = pipe(prompt=sd_prompt, num_inference_steps=4, guidance_scale=0.0, width=512, height=768).images[0]
+        info += f"\n\n🎨 SD prompt: {sd_prompt}"
+
+    final_img = compose_poster(bg_img, poster_info)
+    buf = io.BytesIO()
+    final_img.save(buf, format='PNG')
+    img_b64 = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
+    return final_img, img_b64, info
+
+def poster_flow(desc, bg_prompt, uploaded_bg, add_to_site, bg_style="Restaurant atmosphere"):
+    """Non-GPU wrapper: calls GPU function then updates global state in main process."""
+    if not SD_AVAILABLE and uploaded_bg is None:
+        return None, "❌ SD not available. Upload a background image.", preview(current_html) if current_html else "", _poster_status()
+
+    final_img, img_b64, info = _generate_poster_gpu(desc, bg_prompt, uploaded_bg, bg_style)
+
+    if final_img is None:
+        return None, info, preview(current_html) if current_html else "", _poster_status()
+
+    final_img.save(WORK_DIR / "latest_poster.png")
+
+    if add_to_site and current_html:
+        if "posters" not in current_site_info:
+            current_site_info["posters"] = []
+        current_site_info["posters"].append(img_b64)
+        rebuild_html()
+        n = len(current_site_info["posters"])
+        info += f"\n\n✅ Added! {n} poster{'s' if n > 1 else ''} in carousel."
 
+    return final_img, info, preview(current_html) if current_html else "", _poster_status()
 
-def poster_flow(desc, add_to_site):
-    global current_html
-    img, info = generate_poster(desc)
-    if add_to_site and img and current_html:
-        buf = io.BytesIO()
-        img.save(buf, format='PNG')
-        current_site_info["poster_base64"] = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-        rebuild_html()
-        info += "\n\n✅ Added to banner!"
-    return img, info, preview(current_html) if current_html else ""
+def remove_last_poster():
+    posters = current_site_info.get("posters", [])
+    if not posters:
+        return "*No posters to remove*", preview(current_html)
+    posters.pop()
+    current_site_info["posters"] = posters
+    rebuild_html()
+    return _poster_status(), preview(current_html)
 
+def clear_all_posters():
+    current_site_info["posters"] = []
+    current_site_info["poster_base64"] = ""
+    rebuild_html()
+    return "*No posters added yet*", preview(current_html)
 
 # ============================================================
 # HELPERS
@@ -656,30 +1131,18 @@ def preview(html):
         return '<div style="padding:40px;text-align:center;color:#999;">No website yet. Create one first.</div>'
     return f'<div style="max-width:480px;margin:0 auto;height:700px;overflow-y:auto;border:1px solid #ddd;border-radius:12px;"><iframe srcdoc="{esc(html)}" style="width:100%;height:100%;border:none;"></iframe></div>'
 
-
-# ============================================================
-# HANDLERS
-# ============================================================
-def create_website(shop_name, desc, phone, addr, hours, style, menu_zip, progress=gr.Progress()):
-    global current_menu_data, current_site_info, current_template_key
-
-    progress(0.1, desc="Processing menu...")
-    current_menu_data = process_menu_zip(menu_zip)
-    current_site_info = {
-        "shop_name": shop_name or "My Shop",
-        "description": desc or "Welcome to our shop!",
-        "phone": phone or "", "address": addr or "", "hours": hours or "",
-        "promo_title": "Coming Soon", "promo_subtitle": "Stay tuned for updates",
-        "custom_css": "", "poster_base64": "", "block_overrides": {},
-    }
-    current_template_key = {"Warm & Cozy": "warm", "Dark & Elegant": "dark"}.get(style, "warm")
-
-    progress(0.5, desc="Building...")
-    rebuild_html()
-    n = sum(len(v) for v in current_menu_data.values())
-    progress(1.0)
-    return preview(current_html), f"✅ Done! {n} items, {len(current_menu_data)} categories."
-
+def preview_menu(html):
+    """Render preview with the Menu page active instead of Home."""
+    if not html:
+        return preview(html)
+    html_m = html.replace(
+        '<div class="page active" id="page-home">',
+        '<div class="page" id="page-home">'
+    ).replace(
+        '<div class="page" id="page-menu">',
+        '<div class="page active" id="page-menu">'
+    )
+    return preview(html_m)
 
 def download():
     if not current_html:
@@ -688,7 +1151,6 @@ def download():
     p.write_text(current_html, encoding='utf-8')
     return str(p)
 
-
 # ============================================================
 # UI
 # ============================================================
@@ -701,101 +1163,180 @@ def build_app():
             with gr.Row():
                 with gr.Column(scale=1):
                     gr.Markdown("### Business Info")
-                    shop_name = gr.Textbox(label="Shop Name", placeholder="The Cozy Bean")
-                    desc = gr.Textbox(label="Description", placeholder="A neighborhood coffee shop...", lines=3)
-                    phone = gr.Textbox(label="Phone", placeholder="+44 20 7946 0958")
-                    addr = gr.Textbox(label="Address", placeholder="42 High Street, London")
-                    hours = gr.Textbox(label="Hours", placeholder="Mon-Fri 7am-7pm")
-                    style = gr.Dropdown(label="Template", choices=["Warm & Cozy", "Dark & Elegant"], value="Warm & Cozy")
+                    shop_name = gr.Textbox(label="Shop Name",    placeholder="The Cozy Bean")
+                    desc      = gr.Textbox(label="Description",  placeholder="A neighborhood coffee shop...", lines=3)
+                    phone     = gr.Textbox(label="Phone",        placeholder="+44 20 7946 0958")
+                    addr      = gr.Textbox(label="Address",      placeholder="42 High Street, London")
+                    hours     = gr.Textbox(label="Hours",        placeholder="Mon-Fri 7am-7pm")
+                    style     = gr.Dropdown(label="Template", choices=["Warm & Cozy", "Dark & Elegant"], value="Warm & Cozy")
                     gr.Markdown("### Menu ZIP\n`Category/Name_Price.png`")
-                    menu_zip = gr.File(label="Upload ZIP", file_types=[".zip"])
+                    menu_zip  = gr.File(label="Upload ZIP", file_types=[".zip"])
                     create_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
-                    status = gr.Markdown("")
+                    status    = gr.Markdown("")
                 with gr.Column(scale=1):
                     gr.Markdown("### 📱 Preview")
                     create_prev = gr.HTML()
-            create_btn.click(create_website, [shop_name, desc, phone, addr, hours, style, menu_zip], [create_prev, status])
 
         # ---- Tab 2: Edit ----
         with gr.Tab("📝 Edit Website"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    gr.Markdown("### Select a section to edit")
-
-                    block_choice = gr.Dropdown(
-                        label="Section",
-                        choices=list(BLOCKS.keys()),
-                        value="Hero Banner",
-                    )
-
-                    # Dynamic hint based on selection
-                    block_hint = gr.Markdown("*Hero Banner: Shop name, tagline, background decorations*")
-
-                    def update_hint(choice):
-                        b = BLOCKS[choice]
-                        mode_label = "🤖 AI Coder" if b["mode"] == "coder" else "⚡ Instant"
-                        examples = {
-                            "Hero Banner": "e.g., *Make the tagline more inviting* · *Add a warm welcome message* · *Make it feel luxurious*",
-                            "Promo / Event": "e.g., *Add text: Grand Opening Feb 28* · *Make it feel festive and exciting*",
-                            "About / Story": "e.g., *Write a cozy brand story about our 10-year journey* · *Make it shorter and punchier*",
-                            "Contact Info": "e.g., *Change phone to 555-1234* · *Update hours to Mon-Sat 8am-9pm*",
-                            "Menu Items": "e.g., *Add Iced Mocha $6 in Coffee* · *Remove Espresso* · *Change Latte to $5.50*",
-                            "Style / Colors": "e.g., *Change primary color to forest green* · *Use Georgia for headings* · *Make background darker*",
-                        }
-                        return f"**{mode_label}** — {b['desc']}\n\n{examples.get(choice, '')}"
-
-                    block_choice.change(update_hint, block_choice, block_hint)
-
-                    chatbot = gr.Chatbot(height=300)
 
+                    gr.Markdown("### Page Blocks")
+                    block_radio = gr.Radio(label="Select a block to edit", choices=[], interactive=True, value=None)
+
+                    with gr.Row():
+                        move_up_btn = gr.Button("⬆️ Up",     size="sm")
+                        move_dn_btn = gr.Button("⬇️ Down",   size="sm")
+                        remove_btn  = gr.Button("🗑️ Remove", size="sm")
+
+                    with gr.Row():
+                        add_type = gr.Dropdown(
+                            choices=list(BLOCK_DEFAULTS.keys()),
+                            value="Hero Banner", label="New block type", scale=2
+                        )
+                        add_btn = gr.Button("➕ Add", size="sm", scale=1)
+
+                    gr.Markdown("---")
+                    gr.Markdown("### ✨ Generate Custom Block")
+                    gr.Markdown("*Describe a new block — LLM generates the HTML from scratch*")
+                    with gr.Row():
+                        custom_block_msg = gr.Textbox(placeholder='e.g. "A loyalty card section with 10 stamp slots" · "Opening hours table"', lines=2, scale=3, label="")
+                        custom_block_btn = gr.Button("Generate", variant="primary", scale=1)
+
+                    gr.Markdown("---")
+                    gr.Markdown("### Edit Selected Block")
+                    gr.Markdown("*Describe what to change — the LLM rewrites the block HTML*")
+                    chatbot  = gr.Chatbot(height=280)
+                    with gr.Row():
+                        edit_msg  = gr.Textbox(placeholder='e.g. "Add a description field to each menu item"', lines=2, scale=3, label="")
+                        edit_send = gr.Button("Send", variant="primary", scale=1)
+
+                    gr.Markdown("### Style / Colors")
+                    with gr.Row():
+                        style_msg  = gr.Textbox(placeholder='e.g. "Make primary color forest green"', scale=3, label="")
+                        style_send = gr.Button("Apply", scale=1)
+
+                    gr.Markdown("### Menu Items")
+                    gr.Markdown("*Add, remove, or change prices of individual items*")
+                    with gr.Row():
+                        menu_data_msg   = gr.Textbox(placeholder='e.g. "Add Iced Mocha $6 in Coffee" · "Remove Espresso" · "Change Latte to $5.50"', lines=2, scale=3, label="")
+                        menu_data_image = gr.Image(label="📷 Photo", type="filepath", scale=1, height=120)
+                        menu_data_send  = gr.Button("Send", variant="primary", scale=1)
+
+                    gr.Markdown("### Menu Item Structure")
+                    gr.Markdown("*Rewrite the layout of every menu card — add descriptions, ratings, badges, etc.*")
                     with gr.Row():
-                        msg = gr.Textbox(label="Instruction", placeholder="Describe what to change...", lines=2, scale=3)
-                        img = gr.Image(label="📷 Photo", type="filepath", scale=1, height=100)
-                    send_btn = gr.Button("Send", variant="primary")
+                        menu_struct_msg  = gr.Textbox(placeholder='e.g. "Add a short description field below the item name"', lines=2, scale=3, label="")
+                        menu_struct_send = gr.Button("Send", variant="primary", scale=1)
+                    menu_reset_btn = gr.Button("↩️ Reset menu to default", size="sm")
 
                 with gr.Column(scale=1):
                     gr.Markdown("### 📱 Live Preview")
                     edit_prev = gr.HTML()
 
-            send_btn.click(
-                handle_edit, [block_choice, msg, img, chatbot], [chatbot, edit_prev]
-            ).then(lambda: ("", None), outputs=[msg, img])
-            msg.submit(
-                handle_edit, [block_choice, msg, img, chatbot], [chatbot, edit_prev]
-            ).then(lambda: ("", None), outputs=[msg, img])
+            # Custom block generation
+            custom_block_btn.click(
+                generate_custom_block, [custom_block_msg, chatbot], [chatbot, edit_prev, block_radio]
+            ).then(lambda: "", outputs=custom_block_msg)
+            custom_block_msg.submit(
+                generate_custom_block, [custom_block_msg, chatbot], [chatbot, edit_prev, block_radio]
+            ).then(lambda: "", outputs=custom_block_msg)
+
+            # Wire up block management
+            add_btn.click(handle_add_block,    [add_type, block_radio],  [edit_prev, block_radio])
+            remove_btn.click(handle_remove_block, block_radio,           [edit_prev, block_radio])
+            move_up_btn.click(handle_move_up,  block_radio,              [edit_prev, block_radio])
+            move_dn_btn.click(handle_move_down, block_radio,             [edit_prev, block_radio])
+
+            # Wire up LLM editing
+            edit_send.click(
+                edit_block_with_llm, [block_radio, edit_msg, chatbot], [chatbot, edit_prev]
+            ).then(lambda: "", outputs=edit_msg)
+            edit_msg.submit(
+                edit_block_with_llm, [block_radio, edit_msg, chatbot], [chatbot, edit_prev]
+            ).then(lambda: "", outputs=edit_msg)
+
+            # Style
+            style_send.click(
+                handle_style_edit, [style_msg, chatbot], [chatbot, edit_prev]
+            ).then(lambda: "", outputs=style_msg)
+
+            # Menu item data editing
+            menu_data_send.click(
+                handle_menu_data_edit, [menu_data_msg, menu_data_image, chatbot], [chatbot, edit_prev]
+            ).then(lambda: ("", None), outputs=[menu_data_msg, menu_data_image])
+            menu_data_msg.submit(
+                handle_menu_data_edit, [menu_data_msg, menu_data_image, chatbot], [chatbot, edit_prev]
+            ).then(lambda: ("", None), outputs=[menu_data_msg, menu_data_image])
+
+            # Menu structure editing
+            menu_struct_send.click(
+                edit_menu_structure, [menu_struct_msg, chatbot], [chatbot, edit_prev]
+            ).then(lambda: "", outputs=menu_struct_msg)
+            menu_struct_msg.submit(
+                edit_menu_structure, [menu_struct_msg, chatbot], [chatbot, edit_prev]
+            ).then(lambda: "", outputs=menu_struct_msg)
+            menu_reset_btn.click(reset_menu_structure, chatbot, [chatbot, edit_prev])
 
         # ---- Tab 3: Poster ----
         with gr.Tab("🎨 Poster"):
             with gr.Row():
                 with gr.Column(scale=1):
                     gr.Markdown("### Event / Promotion")
-                    edesc = gr.Textbox(label="Description", placeholder="Spring launch of Sakura Latte...", lines=4)
-                    add2site = gr.Checkbox(label="Add to website banner", value=True)
-                    pbtn = gr.Button("🎨 Generate Poster", variant="primary", size="lg")
-                    pinfo = gr.Markdown("")
+                    edesc    = gr.Textbox(label="Promotion description", placeholder="Student discount 20% off, valid Mon-Fri with student ID", lines=2)
+                    bg_prompt = gr.Textbox(label="Background prompt (optional)", placeholder="Leave empty to auto-generate", lines=2)
+                    bg_style  = gr.Radio(choices=["Restaurant atmosphere", "Match the promotion"], value="Restaurant atmosphere", label="Auto-generate style")
+                    gr.Markdown("*Or upload your own background photo*")
+                    poster_bg = gr.Image(label="Background Image (optional)", type="filepath")
+                    add2site  = gr.Checkbox(label="Add to website carousel", value=True)
+                    pbtn      = gr.Button("🎨 Generate Poster", variant="primary", size="lg")
+                    pinfo     = gr.Markdown("")
+                    gr.Markdown("### Manage Carousel")
+                    poster_status = gr.Markdown("*No posters added yet*")
+                    with gr.Row():
+                        remove_poster_btn = gr.Button("🗑️ Remove last", size="sm")
+                        clear_poster_btn  = gr.Button("🗑️ Clear all",   size="sm", variant="stop")
                 with gr.Column(scale=1):
                     gr.Markdown("### Result")
-                    pimg = gr.Image(label="Poster", type="pil")
+                    pimg  = gr.Image(label="Poster", type="pil")
                     gr.Markdown("### 📱 Preview")
                     pprev = gr.HTML()
-            pbtn.click(poster_flow, [edesc, add2site], [pimg, pinfo, pprev])
+
+            pbtn.click(poster_flow, [edesc, bg_prompt, poster_bg, add2site, bg_style], [pimg, pinfo, pprev, poster_status])
+            remove_poster_btn.click(remove_last_poster, outputs=[poster_status, pprev])
+            clear_poster_btn.click(clear_all_posters,   outputs=[poster_status, pprev])
 
         # ---- Tab 4: Download ----
         with gr.Tab("💾 Download"):
             gr.Markdown("### Download as single HTML\nFully self-contained with embedded images.")
-            dbtn = gr.Button("💾 Download", variant="primary")
+            dbtn  = gr.Button("💾 Download", variant="primary")
             dfile = gr.File(label="Website file")
             dbtn.click(download, outputs=dfile)
 
-    return app
+        # Create: first update preview+status, then refresh block list
+        create_btn.click(
+            create_website,
+            [shop_name, desc, phone, addr, hours, style, menu_zip],
+            [create_prev, status]
+        ).then(
+            lambda: gr.update(choices=get_block_choices(), value=None),
+            inputs=None,
+            outputs=block_radio
+        ).then(
+            lambda: preview(current_html),
+            inputs=None,
+            outputs=edit_prev
+        )
 
+    return app
 
 # ============================================================
 # MAIN
 # ============================================================
 if __name__ == "__main__":
     print("=" * 50)
-    print("  ShopSite AI")
+    print("  ShopSite AI  —  Block-based page builder")
     print("=" * 50)
 
     if not TEMPLATE_DIR.exists():
@@ -803,19 +1344,24 @@ if __name__ == "__main__":
     else:
         print(f"\n✅ Templates: {', '.join(t.stem for t in TEMPLATE_DIR.glob('*.html'))}")
 
-    try:
-        r = requests.get(f"{OLLAMA_BASE}/api/tags", timeout=5)
-        models = [m["name"] for m in r.json().get("models", [])]
-        print(f"✅ Ollama: {', '.join(models)}")
-    except Exception:
-        print("❌ Ollama not connected. Run: ollama serve")
+    if TORCH_AVAILABLE:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"✅ torch available. Device: {device}")
+        print(f"   General model: {QWEN_MODEL}")
+        print(f"   Coder model:   {QWEN_CODER_MODEL}")
+        print("   (Models downloaded from HuggingFace on first use)")
+    else:
+        print("❌ torch/transformers not installed. Run: pip install torch transformers accelerate")
 
     if SD_AVAILABLE:
-        print(f"✅ SD Turbo. CUDA: {torch.cuda.is_available()}")
-        if torch.cuda.is_available():
-            print(f"   GPU: {torch.cuda.get_device_name(0)}, {torch.cuda.get_device_properties(0).total_memory/1024**3:.1f}GB")
+        print(f"✅ SD Turbo ready. CUDA: {torch.cuda.is_available()}")
     else:
-        print("⚠️  Poster generation disabled.")
+        print("⚠️  Poster generation disabled (diffusers not available).")
 
     print("\n🚀 http://127.0.0.1:7860\n")
-    build_app().launch(server_name="127.0.0.1", server_port=7860)
+    in_colab = "google.colab" in str(__import__("sys").modules)
+    demo = build_app()
+    if in_colab:
+        demo.launch(share=True)
+    else:
+        demo.launch(server_name="0.0.0.0", server_port=7860)