ui/layout.py

import datetime
import gradio as gr
import json
import time
from pathlib import Path
# from logic.supabase_client import auth_handler
from huggingface_hub import HfApi
from urllib.parse import urlencode, parse_qs

from logic.data_utils import CustomHFDatasetSaver, mark_sample_excluded
from data.lang2eng_map import lang2eng_mapping
from gradio_modal import Modal
from logic.handlers import *
from logic.image_transforms import rotate_image_90_left, rotate_image_90_right, reset_image_to_original
from config.settings import *
from functools import partial

from .selection_page import build_selection_page, prefill_country_language
from .main_page import build_main_page
from .main_page import sort_with_pyuca
from .user_page import build_user_page, render_comparative_stats, render_user_centric_stats, render_missing_concepts
from logic.user_stats import rebuild_user_profile, compute_comparative_stats, compute_user_centric_stats

# js_code = """
# function() {
#     // Get the full URL with the fragment
#     const url = window.location.href;
#     const fragment = url.split('#')[1];

#     if (!fragment) {
#         return "";
#     }

#     // Parse the fragment into an object
#     const params = new URLSearchParams(fragment);
#     const access_token = params.get('access_token');
#     const refresh_token = params.get('refresh_token');

#     // Create a JSON string with the tokens
#     const tokens = JSON.stringify({
#         access_token: access_token,
#         refresh_token: refresh_token
#     });

#     // Return the JSON string to the Gradio output component
#     return tokens;
# }
# """



def load_examples_raw(json_path: Path):
    """
    Load static examples from JSON as a list of dicts, preserving all fields
    (image, image_url, caption, country, language, category, concept,
    additional_concepts, id, notes).
    """
    if not json_path.exists():
        print(f"Examples file not found at {json_path}")
        return []

    try:
        with open(json_path, "r", encoding="utf-8") as f:
            raw_examples = json.load(f)
    except Exception as e:
        print(f"Failed to read examples file: {e}")
        return []

    if not isinstance(raw_examples, list):
        return []
    return raw_examples


def _html_escape(text):
    if text is None:
        return ""
    text = str(text)
    return (
        text.replace("&", "&")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )


def render_examples_modal_html(examples, metadata_dict, country, language):
    """Build an HTML string displaying example cards for the current locale."""
    meta = {}
    try:
        meta = metadata_dict.get(country, {}).get(language, {}) or {}
    except Exception:
        meta = {}

    # Labels
    image_label       = meta.get("Examples_Image_Label", "Image")
    desc_label        = meta.get("Examples_Description_Label", "Description")
    country_label     = meta.get("Country", "Country")
    language_label    = meta.get("Language", "Language")
    category_label    = meta.get("Examples_Category_Label", "Main category")
    concept_label     = meta.get("Examples_Concept_Label", "Main concept")
    additional_label  = meta.get("Examples_Additional_Label", "Additional concepts")
    example_label     = meta.get("Examples_Example_Label", "Example")
    intro_text        = meta.get("Examples_Intro", "Below are reference examples. Each field is labeled so you can see exactly what to put where.")
    no_examples_text  = meta.get("Examples_Empty_Msg", "No examples configured yet. Add entries to <code>data/examples.json</code>.")
    intro = (
        '<div class="examples-intro">'
        f'<p>{_html_escape(intro_text)}</p>'
        '</div>'
    )

    display_examples = [ex for ex in (examples or []) if isinstance(ex, dict)]

    if not display_examples:
        return intro + f'<p class="examples-empty"><em>{no_examples_text}</em></p>'

    cards = []
    for idx, ex in enumerate(display_examples, start=1):
        concept = ex.get("concept", "") or ""
        image_url = ex.get("image_url") or ""
        caption = ex.get("caption", "") or ""
        ex_country = ex.get("country", "") or ""
        ex_language = ex.get("language", "") or ""
        ex_category = ex.get("category", "") or ""
        additional = ex.get("additional_concepts", []) or []
        if isinstance(additional, list):
            additional_chips = "".join(
                f'<span class="chip">{_html_escape(str(c).strip())}</span>'
                for c in additional if c
            ) or '<span class="chip chip-empty">—</span>'
        else:
            additional_chips = (
                f'<span class="chip">{_html_escape(str(additional).strip())}</span>'
                if str(additional).strip() else '<span class="chip chip-empty">—</span>'
            )
        notes = ex.get("notes", "") or ""
        title = concept.strip() or f"{example_label} {idx}"

        image_html = ""
        if image_url:
            src = image_url if image_url.startswith(("http://", "https://")) else f"/gradio_api/file={image_url}"
            image_html = (
                '<div class="example-image">'
                f'<img src="{_html_escape(src)}" alt="{_html_escape(title)}" />'
                '</div>'
            )

        meta_rows = (
            '<dl class="example-fields">'
            f'<dt>📍 {_html_escape(country_label)}</dt><dd>{_html_escape(ex_country) or "—"}</dd>'
            f'<dt>💬 {_html_escape(language_label)}</dt><dd>{_html_escape(ex_language) or "—"}</dd>'
            f'<dt>📂 {_html_escape(category_label)}</dt><dd>{_html_escape(ex_category) or "—"}</dd>'
            f'<dt>🎯 {_html_escape(concept_label)}</dt><dd><strong>{_html_escape(concept) or "—"}</strong></dd>'
            f'<dt>➕ {_html_escape(additional_label)}</dt><dd class="chip-row">{additional_chips}</dd>'
            '</dl>'
        )

        description_html = (
            '<div class="example-description">'
            f'<span class="block-label">{_html_escape(desc_label)}</span>'
            f'<p>{_html_escape(caption) or "<em>—</em>"}</p>'
            '</div>'
        )

        notes_html = ""
        if notes:
            is_positive = ex.get("is_positive")
            if is_positive is True:
                bg_color = "rgba(34, 197, 94, 0.12)"
                border_color = "rgba(34, 197, 94, 0.35)"
                emoji = "✅"
            elif is_positive is False:
                bg_color = "rgba(239, 68, 68, 0.12)"
                border_color = "rgba(239, 68, 68, 0.35)"
                emoji = "❌"
            else:
                bg_color = "rgba(250, 204, 21, 0.12)"
                border_color = "rgba(250, 204, 21, 0.35)"
                emoji = "💡"

            notes_html = (
                f'<div class="example-notes" style="background: {bg_color}; border: 1px solid {border_color};">'
                f'<span class="notes-icon" aria-hidden="true">{emoji}</span>'
                f'<p>{_html_escape(notes)}</p>'
                '</div>'
            )

        cards.append(
            '<article class="example-card">'
            '<header class="example-header">'
            f'<span class="example-number">{_html_escape(example_label)} {idx}</span>'
            f'<h3>{_html_escape(title)}</h3>'
            '</header>'
            '<div class="example-body">'
            f'{image_html}'
            '<div class="example-content">'
            f'{meta_rows}'
            f'{description_html}'
            f'{notes_html}'
            '</div>'
            '</div>'
            '</article>'
        )

    return intro + '<div class="examples-grid">' + "".join(cards) + '</div>'


_EXAMPLES_FALLBACK_PATH = Path(__file__).resolve().parent.parent / "data" / "guidelines" / "USA" / "English" / "examples.json"


def build_examples_modal_updates(local_storage, metadata_dict):
    """
    Build (button_label, modal_title_md, modal_content_md) updates for the
    current country/language stored in local_storage.
    """
    country, language, _ = extract_country_language_username(local_storage)
    meta = {}
    try:
        meta = metadata_dict.get(country, {}).get(language, {}) or {}
    except Exception:
        meta = {}

    # Load locale-specific examples, fall back to USA/English
    examples_path = Path(meta.get("Examples", "")) if meta.get("Examples") else None
    if not examples_path or not examples_path.exists():
        examples_path = _EXAMPLES_FALLBACK_PATH
    examples = load_examples_raw(examples_path)

    btn_label = meta.get("Examples_btn") or meta.get("Tab_Examples", "🖼️ See Examples")
    if not btn_label.startswith(("🖼", "📘", "📗", "📚")):
        btn_label = f"🖼️ {btn_label}"

    modal_title = meta.get(
        "Examples_Modal_Title", "## Examples — how to fill in each field"
    )
    close_label = meta.get("Examples_Modal_Close_btn") or meta.get("Modal_Cancel_btn", "Close")

    content = render_examples_modal_html(examples, metadata_dict, country, language)
    return (
        gr.update(value=btn_label),
        gr.update(value=modal_title),
        gr.update(value=content),
        gr.update(value=close_label),
    )


scroll_to_top_js = """
function() {
    window.scrollTo({
        top: 0,
        behavior: 'smooth'
    });
    return "";
}
"""

# def login_user(email, password):
#     result = auth_handler.login(email, password)
#     if result['success']:
#         session_data = result['data']
#         persistent_data = {
#             "refresh_token": session_data['refresh_token'],
#             "user_email": session_data['user_email']
#         }
#         return session_data['client'], persistent_data, result['message']
#     else:
#         persistent_data = {
#             "refresh_token": "",
#             "user_email": ""
#         }
#         return None, persistent_data, result['message']

# def login_user_recovery(session_data: str):
#     """
#     This function receives session data (tokens as a JSON string) from the frontend,
#     retrieves the session, and returns data in a format similar to login_user.
#     """
#     try:
#         import json
#         tokens = json.loads(session_data)
#         access_token = tokens.get("access_token")
#         refresh_token = tokens.get("refresh_token")

#         if not access_token or not refresh_token:
#             return None, gr.skip(), "Invalid session data provided."

#         result = auth_handler.retrieve_session_from_tokens(access_token, refresh_token)
        
#         if result['success']:
#             session_data_result = result['data']
#             persistent_data = {
#                 "refresh_token": session_data_result['refresh_token'],
#                 "user_email": session_data_result['user_email']
#             }
#             return session_data_result['client'], persistent_data, result['message']
#         else:
#             persistent_data = {
#                 "refresh_token": "",
#                 "user_email": ""
#             }
#             return None, persistent_data, result['message']

#     except Exception as e:
#         return None, gr.skip(), f"Failed to process recovery login: {e}"

# def sign_up(email, password):
#     result = auth_handler.sign_up(email, password)
#     return result['message']

# def reset_password(email):
#     result = auth_handler.reset_password_for_email(email)
#     return result['message']

# def log_out(supabase_user_client, persistent_session):
#     """
#     Logs out the user and clears the session. If error occurs, it returns an empty persistent session (logging out user).
#     """
#     persistent_session = {
#         "refresh_token": "",
#         "user_email": ""
#     }
#     if supabase_user_client:
#         result = auth_handler.logout(supabase_user_client)
#         if result['success']:
#             print("User logged out successfully.")
#             return persistent_session
#         else:
#             print(f"Error logging out: {result['message']}")
#             return persistent_session
#     else:
#         print("No user client provided to log out.")
#         return persistent_session

# def restore_user_session(session_data, login_status=None):
#     print("Restoring user session with data:", session_data)
#     # defualt values if the user is not logged in
#     # or the session data is not valid
#     login_status_update = gr.update(value= login_status if login_status else "")
#     proceed_button_update = gr.update(value="Proceed as Anonymous User", interactive=True)
#     login_button_update = gr.update(visible=True)
#     sign_up_button_update = gr.update(visible=True)
#     reset_password_button_update = gr.update(visible=True)
#     logout_button_update = gr.update(visible=False)
#     change_password_field_update = gr.update(visible=False)
#     change_password_field_confirm_update = gr.update(visible=False)
#     change_password_button_update = gr.update(visible=False)
#     change_password_status_update = gr.update(value="")
#     persistent_data = {
#         "refresh_token": "",
#         "user_email": ""
#     }
#     if not session_data or not session_data.get('refresh_token', ''):
#         print("No session data found, proceeding as anonymous user.")
#         return None, persistent_data, login_status_update, proceed_button_update, login_button_update, sign_up_button_update, reset_password_button_update, logout_button_update, change_password_field_update, change_password_field_confirm_update, change_password_button_update, change_password_status_update
    
#     result = auth_handler.restore_session(session_data['refresh_token'])
#     if result['success']:
#         restored_session = result['data']
#         new_persistent_data = {
#             "refresh_token": restored_session['refresh_token'],
#             "user_email": restored_session['user_email']
#         }
#         login_status_update = gr.update(value=result['message'])
#         proceed_button_update = gr.update(value="Proceed", interactive=True)
#         login_button_update = gr.update(visible=False)
#         sign_up_button_update = gr.update(visible=False)
#         reset_password_button_update = gr.update(visible=False)
#         logout_button_update = gr.update(visible=True)
#         change_password_field_update = gr.update(visible=True)
#         change_password_field_confirm_update = gr.update(visible=True)
#         change_password_button_update = gr.update(visible=True)
#         return restored_session['client'], new_persistent_data, login_status_update, proceed_button_update, login_button_update, sign_up_button_update, reset_password_button_update, logout_button_update, change_password_field_update, change_password_field_confirm_update, change_password_button_update, change_password_status_update
#     else:
#         return None, persistent_data, login_status_update, proceed_button_update, login_button_update, sign_up_button_update, reset_password_button_update, logout_button_update, change_password_field_update, change_password_field_confirm_update, change_password_button_update, change_password_status_update

# def change_password(supabase_user_client, new_password, confirm_password):
#     """
#     Changes the user's password.
#     """
#     if new_password != confirm_password:
#         return "Passwords do not match. Please try again."
#     result = auth_handler.change_password(supabase_user_client, new_password)
#     return result['message']


def get_key_by_value(dictionary, value):
    for key, val in dictionary.items():
        if val == value:
            return key
    return None


def _category_display_value_choices(categories_list, language):
    translations = words_mapping.get(language) or {}
    return [(translations.get(category, category), category) for category in categories_list]


def _canonical_category_value(category, language, categories):
    if not category:
        return None
    if category in categories:
        return category
    reverse_translations = {
        display: canonical
        for canonical, display in (words_mapping.get(language) or {}).items()
    }
    canonical = reverse_translations.get(category)
    if canonical in categories:
        return canonical
    return None


def handle_click_example_with_states(
    user_examples,
    vlm_captions,
    vlm_models,
    vlm_feedbacks,
    vlm_usages,
    additional_concepts_map,
    concepts_dict,
):
    """Load an example while preserving its original VLM metadata in state."""
    result = list(
        handle_click_example(
            user_examples,
            vlm_captions,
            vlm_models,
            concepts_dict,
            additional_concepts_map,
        )
    )
    img, url, caption, example_id, cat, concept = result[:6]
    ac1, ac2, ac3, ac4, ac5 = result[6:11]
    vlm_cap = result[12]
    saved_model = (vlm_models or {}).get(example_id) or result[13]
    saved_feedback = (vlm_feedbacks or {}).get(example_id) or None
    saved_usage = (vlm_usages or {}).get(example_id) if vlm_cap else None
    needs_feedback = bool(vlm_cap and not saved_feedback)
    feedback_update = (
        gr.update(value=None, visible=True)
        if needs_feedback
        else gr.update(visible=False)
    )
    ts = datetime.datetime.now().timestamp()
    vlm_sig = image_signature(img) if vlm_cap else None
    return result + [
        feedback_update,
        img,
        url,
        caption,
        cat,
        concept,
        ac1,
        ac2,
        ac3,
        ac4,
        ac5,
        example_id,
        vlm_cap,
        saved_model,
        saved_feedback,
        saved_usage,
        ts,
        vlm_sig,
    ]


def extract_country_language_username(local_storage):
    if not local_storage:
        return None, None, None
    email = (local_storage[2] or "").strip() if len(local_storage) > 2 else ""
    return local_storage[0], local_storage[1], email

def _is_anonymous_email(email):
    email = (email or "").strip().lower()
    return (not email) or email == "anonymous"


def set_step(step, show_browse_data=True, email=""):
    is_anonymous = _is_anonymous_email(email)
    show_data_row = (step == 1) and (not is_anonymous)
    show_data_table = show_data_row and bool(show_browse_data)
    return (
        gr.update(visible=step == 1),
        gr.update(visible=step == 2),
        gr.update(visible=step == 3),
        gr.update(visible=show_data_row),
        gr.update(visible=show_data_table),
    )

def set_step_from_storage(step, show_browse_data, local_storage):
    _, _, email = extract_country_language_username(local_storage)
    return set_step(step, show_browse_data, email)


def sync_preview(image):
    return image, image

def toggle_browse_data(show_browse, current_btn_value):
    new_show = not show_browse
    base_text = current_btn_value.lstrip("▼▲").strip()
    return (
        new_show,
        gr.update(visible=new_show),
        gr.update(value=f"▼ {base_text}" if new_show else f"▲ {base_text}"),
    )

def load_concepts(category, concept_btn, local_storage, loading_example, concepts):
    country, lang, _ = local_storage
    if country is None or lang is None:
        return gr.update(choices=[], value=None), False
    eng_lang = lang2eng_mapping.get(lang, lang)
    if category:
        category = _canonical_category_value(category, lang, concepts[country][eng_lang])
        if category is None:
            return gr.update(choices=[], value=None), False
        choices = concepts[country][eng_lang][category]
        # Keep custom concept values when loading an existing sample.
        if loading_example:
            concept_val = concept_btn
        else:
            concept_val = concept_btn if concept_btn in choices else None
        return gr.update(choices=sort_with_pyuca(choices), value=concept_val), False
    else:
        concepts_list = []
        for cat in concepts[country][eng_lang]:
            for concept in concepts[country][eng_lang][cat]:
                concepts_list.append(concept)
        return gr.update(choices=sort_with_pyuca(concepts_list), value=None), False

def update_category(category, concept_btn, local_storage, loading_example, concepts):
    country, lang, _ = local_storage
    if country is None or lang is None:
        return gr.update(choices=[], value=None), False
    eng_lang = lang2eng_mapping.get(lang, lang)
    all_eng_cats = sort_with_pyuca(list(concepts[country][eng_lang].keys()))
    category_choices = _category_display_value_choices(all_eng_cats, lang)
    if not category and concept_btn:
        found_eng_cat = None
        for cat in concepts[country][eng_lang]:
            for concept in concepts[country][eng_lang][cat]:
                if concept == concept_btn:
                    found_eng_cat = cat
        return gr.update(choices=category_choices, value=found_eng_cat), True
    else:
        category = _canonical_category_value(category, lang, concepts[country][eng_lang])
        return gr.update(choices=category_choices, value=category), loading_example

def load_profile(profile: gr.OAuthProfile | None):
    # if the user is not logged in, profile will be None
    if profile is None:
        proceed_msg = "Continue"
        anno_msg = "Annotate"
        username = ""
        user_name = ""
        hf_username = ""
        hf_email = ""
        flag = False
    else:
        user_name = profile.name
        proceed_msg = f"Continue as {user_name}"
        anno_msg = f"Annotate as {user_name}"
        username = profile.username
        hf_username = profile.username
        # The HF "email" OAuth scope (see README.md) makes the email claim
        # available in the userinfo. OAuthProfile is a dict, so access it via
        # .get(). Falls back to "" if the claim is absent (e.g. local mock).
        hf_email = (profile.get("email") or "").strip()
        flag = True
    return (
        anno_msg,
        username,
        user_name,
        gr.update(value=proceed_msg, visible=flag, interactive=False),
        hf_username,
        hf_email,
    )

def get_username(request: gr.Request):
    if request.user:
        return request.user
    return "anonymous"

def update_profile_proceed(profile: gr.OAuthProfile | None, current_email=""):
    # Uses OAuth profile instead of request.user (no auth middleware).
    current_email = (current_email or "").strip()
    if profile is None:
        # Not signed in with Hugging Face. If the user typed an email manually,
        # proceed with that identity instead of forcing anonymous (and without
        # clobbering the email they entered).
        if current_email:
            return "Annotate", "", "", ""
        gr.Warning("User is not logged in. Proceeding as anonymous.", duration=3)
        return "Annotate as Anonymous", "anonymous", "", ""
    api = HfApi()
    username = profile.username
    hf_username = profile.username
    hf_email = (profile.get("email") or "").strip()
    user_obj = api.get_user_overview(username)
    user_name = user_obj.fullname or username
    return f"Annotate as {user_name}", username, hf_username, hf_email

def load_anonymous():
    proceed_msg = "Annotate as Anonymous"
    username = "anonymous"  # TODO: Make random anonymous users
    return proceed_msg, username, "", gr.update(value="")


def switch_ui(country, language, email, flag=False, metadata_dict=None):
    # wait for 2 seconds
    print(f"Language: {language}, Country: {country}")
    email = (email or "").strip()

    if not flag and (not country or not language):
        gr.Warning(f"⚠️ Please select Country and Language first.",  duration=3)
        return (gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip())

    # Hide the selection page and show the main UI.
    meta = (metadata_dict or {}).get(country, {}).get(language, {}) if metadata_dict else {}
    if flag:
        gr.Info(meta.get("Info_Loading_Selection", "Loading the Language Selection UI"), duration=3)
        time.sleep(2)
        local_storage = [None, None, ""]
        countries = sort_with_pyuca(list(metadata_dict.keys()))
        return (local_storage, gr.update(visible=flag), gr.update(visible=not flag),
                gr.update(visible=not flag),
                gr.update(visible=not flag),
                gr.update(choices=countries, label="Country", value=None),
                gr.update(value=None, choices=[], label="Language", allow_custom_value=False, interactive=False),
                gr.update(value="## Please select your profile"),
                gr.update(value="## Please select your country"),
                gr.update(value="Annotate"))
    
    else:
        gr.Info(meta.get("Info_Loading_Main", "Loading the Main UI"), duration=3)
        time.sleep(2)
        local_storage = [country, language, email]
        return (local_storage, gr.update(visible=flag), gr.update(visible=not flag), 
                gr.update(visible=not flag), gr.update(visible=not flag),
                gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip())


def build_ui(concepts_dict, metadata_dict, HF_API_TOKEN, HF_DATASET_NAME, prefs_store=None):
    hf_writer = CustomHFDatasetSaver(HF_API_TOKEN, HF_DATASET_NAME, private=True)
    custom_css = """
        .compact-container {
            max-width: 600px;
            margin: auto;
            padding: 20px;
        }
        .compact-btn {
            min-width: 0 !important;
        }
        #user_page_btn,
        #user_page_btn button {
            max-width: 260px;
            min-width: 0 !important;
            overflow: hidden;
            text-overflow: ellipsis;
            white-space: nowrap;
        }
        .step-description-box {
            border: 1px solid var(--border-color-primary, #e5e7eb) !important;
            background: var(--block-background-fill, #f9fafb) !important;
            border-radius: 8px !important;
            padding: 0.85rem 1.15rem !important;
            margin-bottom: 1.25rem !important;
            box-shadow: var(--block-shadow, none);
        }
        .step-description-box .step-description-box {
            border: none !important;
            background: transparent !important;
            padding: 0 !important;
            margin: 0 !important;
            box-shadow: none !important;
        }
        #image_instruction_accordion,
        #description_instruction_accordion {
            border: 2px solid rgba(37, 99, 235, 0.35) !important;
            background: rgba(37, 99, 235, 0.06) !important;
            border-radius: 8px !important;
            margin: 0.75rem 0 1rem;
        }
        #image_instruction_accordion .label-wrap,
        #description_instruction_accordion .label-wrap {
            font-weight: 700 !important;
            color: #1d4ed8 !important;
        }
        #image_inp img {
            object-fit: contain;        /* make sure the full image shows */
            height: 460px;           /* set a fixed height */
        }
        #vlm_output .input-container {
        position: relative;
        }
        #vlm_output .input-container::before {
        content: "";
        position: absolute;
        top: 0; left: 0; right: 0; bottom: 0;
        z-index: 10;             /* sits above the textarea */
        background: transparent;
        }
        /* === Examples modal === */
        #examples_modal_content { padding: 0; }
        #examples_modal_content .examples-intro { margin-bottom: 1.25rem; }
        #examples_modal_content .examples-intro > p {
            font-size: 0.95rem;
            color: var(--body-text-color);
            margin: 0 0 0.75rem 0;
        }
        #examples_modal_content .examples-tips {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(240px, 1fr));
            gap: 0.6rem;
        }
        #examples_modal_content .tip-card {
            display: flex;
            gap: 0.6rem;
            padding: 0.7rem 0.85rem;
            border-radius: 10px;
            background: var(--background-fill-secondary, #f6f7fb);
            border: 1px solid var(--border-color-primary, #e3e6ee);
        }
        #examples_modal_content .tip-card .tip-icon { font-size: 1.25rem; line-height: 1.2; }
        #examples_modal_content .tip-card strong { display: block; margin-bottom: 0.2rem; }
        #examples_modal_content .tip-card p { margin: 0; font-size: 0.85rem; opacity: 0.9; }
        #examples_modal_content .examples-grid {
            display: flex;
            flex-direction: column;
            gap: 1.25rem;
            margin-top: 0.5rem;
        }
        #examples_modal_content .example-card {
            border: 1px solid var(--border-color-primary, #e3e6ee);
            border-radius: 14px;
            overflow: hidden;
            background: var(--background-fill-primary, #fff);
            box-shadow: 0 1px 3px rgba(0,0,0,0.04);
        }
        #examples_modal_content .example-header {
            display: flex;
            align-items: center;
            gap: 0.75rem;
            padding: 0.7rem 1rem;
            background: linear-gradient(90deg, rgba(99, 102, 241, 0.08), rgba(168, 85, 247, 0.06));
            border-bottom: 1px solid var(--border-color-primary, #e3e6ee);
        }
        #examples_modal_content .example-number {
            font-size: 0.75rem;
            font-weight: 600;
            text-transform: uppercase;
            letter-spacing: 0.04em;
            background: rgba(99, 102, 241, 0.15);
            color: #4338ca;
            padding: 0.2rem 0.55rem;
            border-radius: 999px;
        }
        #examples_modal_content .example-header h3 {
            margin: 0;
            font-size: 1.05rem;
            font-weight: 600;
        }
        #examples_modal_content .example-body {
            display: grid;
            grid-template-columns: minmax(200px, 280px) 1fr;
            gap: 1rem;
            padding: 1rem;
        }
        @media (max-width: 720px) {
            #examples_modal_content .example-body { grid-template-columns: 1fr; }
        }
        #examples_modal_content .example-image img {
            width: 100%;
            height: auto;
            max-height: 280px;
            object-fit: cover;
            border-radius: 10px;
            display: block;
        }
        #examples_modal_content .example-content { display: flex; flex-direction: column; gap: 0.85rem; }
        #examples_modal_content .example-fields {
            display: grid;
            grid-template-columns: max-content 1fr;
            row-gap: 0.35rem;
            column-gap: 0.75rem;
            margin: 0;
            font-size: 0.9rem;
        }
        #examples_modal_content .example-fields dt {
            font-weight: 500;
            opacity: 0.75;
        }
        #examples_modal_content .example-fields dd { margin: 0; }
        #examples_modal_content .chip-row {
            display: flex;
            flex-wrap: wrap;
            gap: 0.3rem;
        }
        #examples_modal_content .chip {
            display: inline-block;
            padding: 0.15rem 0.55rem;
            border-radius: 999px;
            background: rgba(99, 102, 241, 0.1);
            color: #4338ca;
            font-size: 0.8rem;
            font-weight: 500;
        }
        #examples_modal_content .chip-empty {
            background: transparent;
            color: var(--body-text-color);
            opacity: 0.6;
            font-weight: 400;
        }
        #examples_modal_content .example-description {
            background: var(--background-fill-secondary, #f6f7fb);
            border-left: 3px solid #6366f1;
            border-radius: 0 8px 8px 0;
            padding: 0.6rem 0.85rem;
        }
        #examples_modal_content .block-label {
            display: inline-block;
            font-size: 0.7rem;
            font-weight: 700;
            text-transform: uppercase;
            letter-spacing: 0.05em;
            color: #4f46e5;
            margin-bottom: 0.25rem;
        }
        #examples_modal_content .example-description p {
            margin: 0;
            font-size: 0.92rem;
            line-height: 1.45;
        }
        #examples_modal_content .example-notes {
            display: flex;
            gap: 0.55rem;
            padding: 0.55rem 0.75rem;
            background: rgba(250, 204, 21, 0.12);
            border: 1px solid rgba(250, 204, 21, 0.35);
            border-radius: 10px;
            font-size: 0.88rem;
        }
        #examples_modal_content .example-notes .notes-icon { font-size: 1.05rem; }
        #examples_modal_content .example-notes p { margin: 0; line-height: 1.45; }
        #examples_modal_content .examples-empty {
            padding: 1rem;
            text-align: center;
            opacity: 0.7;
        }
    """
    custom_js = """
    function setMissingConcept(concept) {
        const element = document.querySelector('#hidden_concept_trigger input');
        if (element) {
            element.value = concept;
            element.dispatchEvent(new Event('input'));
        }
    }
    """
    ############################################################################
    with gr.Blocks(css=custom_css, js=custom_js) as ui:
        supabase_user_client = gr.State(None)
        # session_state = gr.State({})
        persistent_session = gr.BrowserState(None)

        local_storage = gr.State([None, None, ""])
        country_state = gr.State(None)
        language_state = gr.State(None)
        username_state = gr.State(None)
        category_state = gr.State(None)
        concept_state = gr.State(None)
        caption_state = gr.State(None)
        image_url_state = gr.State(None)
        quick_submit_state = gr.State(False)
        category_1_concepts_state = gr.State(None)
        category_2_concepts_state = gr.State(None)
        category_3_concepts_state = gr.State(None)
        category_4_concepts_state = gr.State(None)
        category_5_concepts_state = gr.State(None)
        timestamp_state = gr.State(None)
        exampleid_state = gr.State(None)
        vlm_caption_state = gr.State(None)
        vlm_feedback_state = gr.State(None)
        vlm_model_state = gr.State(None)
        vlm_usage_state = gr.State(None)
        vlm_image_signature_state = gr.State(None)
        hf_username_state = gr.State("")
        hf_email_state = gr.State("")
        pref_country_state = gr.State(None)
        pref_language_state = gr.State(None)
        loading_example = gr.State(False) # to check if the values are loaded from a user click on an example in 
        current_step_state = gr.State(1)
        
        # First page: selection
        (
            selection_page, country_choice, language_choice, proceed_btn, anon_btn, login_btn, prof_btn,
            username, user_name, email_inp, intro_title, intro_markdown, signin_title, selection_title,
            signin_section, selection_section,
        ) = build_selection_page(metadata_dict)
        # Output order: proceed button, displayed username, display name, profile button, HF username state, email field.
        # Clearing the email field is important so a previously HF-authenticated email is not
        # reused when the user explicitly chooses to proceed anonymously.
        anon_continue_event = anon_btn.click(fn=load_anonymous, outputs=[proceed_btn, username, hf_username_state, email_inp])

        # Helper-text shown under the email field depending on the identity mode.
        _EMAIL_INFO_EDITABLE = "Add your email to save your work and review or edit it later."
        _EMAIL_INFO_LOCKED = "Using your Hugging Face email. Log out to use a different email."

        def _email_update(value, locked):
            """Build an email_inp update that locks the field (read-only) when the
            user is signed in with a verified Hugging Face email, and keeps it
            editable otherwise."""
            return gr.update(
                value=value,
                interactive=not locked,
                info=_EMAIL_INFO_LOCKED if locked else _EMAIL_INFO_EDITABLE,
            )

        def _prefs_from_hf_or_session(hf_username, hf_email, current_email, session):
            session = session or {}
            session_email = (session.get("email") or "").strip()
            session_country = session.get("country")
            session_language = session.get("language")
            prev_hf_username = (session.get("hf_username") or "").strip()
            current_hf_username = (hf_username or "").strip()
            hf_email = (hf_email or "").strip()
            current_email = (current_email or "").strip()

            # If the user was previously signed in with Hugging Face but is no
            # longer signed in (e.g., they just clicked "Logout"), clear the email
            # and make the field editable again.
            if prev_hf_username and not current_hf_username:
                return _email_update("", locked=False), session_country, session_language

            # Signed in with Hugging Face: force the verified HF email and lock the
            # field, so the contribution identity is always the verified account.
            if current_hf_username:
                country = session_country
                language = session_language
                if prefs_store:
                    by_hf = prefs_store.get_by_hf_username(current_hf_username)
                    if by_hf:
                        country = by_hf.get("last_country") or session_country
                        language = by_hf.get("last_language") or session_language
                if hf_email:
                    # Verified email available -> lock to it.
                    return _email_update(hf_email, locked=True), country, language
                # No verified email (e.g., missing scope / local mock): fall back to
                # a known email but keep the field editable so identity isn't lost.
                fallback_email = session_email or current_email
                return _email_update(fallback_email, locked=False), country, language

            # Not signed in with HF: keep the email editable. Respect a value that
            # is already in the field (manually typed or kept from a prefill).
            if current_email:
                country = session_country
                language = session_language
                if prefs_store:
                    by_email = prefs_store.get_by_email(current_email)
                    if by_email:
                        country = by_email.get("last_country") or session_country
                        language = by_email.get("last_language") or session_language
                return _email_update(current_email, locked=False), country, language

            # Empty field, not logged in: restore a previously persisted email if any.
            email = session_email
            country = session_country
            language = session_language
            if prefs_store and session_email:
                by_email = prefs_store.get_by_email(session_email)
                if by_email:
                    country = by_email.get("last_country") or session_country
                    language = by_email.get("last_language") or session_language
            return _email_update(email, locked=False), country, language

        def _sync_session_after_logout(hf_username, session):
            """Persist the cleared email/hf_username when we detect a logout so
            subsequent reloads don't re-populate the field from stale state."""
            session = session or {}
            prev_hf_username = (session.get("hf_username") or "").strip()
            current_hf_username = (hf_username or "").strip()
            if prev_hf_username and not current_hf_username:
                return dict(session, email="", hf_username="")
            return session

        ui.load(
            load_profile,
            inputs=None,
            outputs=[proceed_btn, username, user_name, prof_btn, hf_username_state, hf_email_state],
        ).then(
            fn=_prefs_from_hf_or_session,
            inputs=[hf_username_state, hf_email_state, email_inp, persistent_session],
            outputs=[email_inp, pref_country_state, pref_language_state],
        ).then(
            fn=_sync_session_after_logout,
            inputs=[hf_username_state, persistent_session],
            outputs=[persistent_session],
        ).then(
            fn=partial(prefill_country_language, metadata=metadata_dict),
            inputs=[pref_country_state, pref_language_state, country_choice, language_choice],
            outputs=[country_choice, language_choice, intro_title, intro_markdown],
        )

        profile_continue_event = prof_btn.click(
            fn=update_profile_proceed,
            inputs=[email_inp],
            outputs=[proceed_btn, username, hf_username_state, hf_email_state],
        ).then(
            fn=_prefs_from_hf_or_session,
            inputs=[hf_username_state, hf_email_state, email_inp, persistent_session],
            outputs=[email_inp, pref_country_state, pref_language_state],
        ).then(
            fn=partial(prefill_country_language, metadata=metadata_dict),
            inputs=[pref_country_state, pref_language_state, country_choice, language_choice],
            outputs=[country_choice, language_choice, intro_title, intro_markdown],
        )

        # def save_choices(d1, d2, state):
        #     state["d1"] = d1
        #     state["d2"] = d2
        #     return state
        # def restore_choices(state):
        #     return state.get("d1"), state.get("d2")
        # saved = gr.State({})
        # d1 = gr.Dropdown(["A", "B", "C"], label="Dropdown 1")
        # d2 = gr.Dropdown(["X", "Y", "Z"], label="Dropdown 2")
        # with gr.Row():
        #     anon = gr.Button("Continue anonymously")
        #     login = gr.Button("Login with Hugging Face")
        # login_real = gr.LoginButton(visible=False)

        # save.click(
        #     save_state,
        #     [d1, d2],
        #     session_state
        # ).then(
        #     lambda: gr.update(visible=True),
        #     None,
        #     login_btn
        # )

        # Second page
        cmp_main_ui = build_main_page(concepts_dict, metadata_dict, local_storage)

        main_ui_placeholder = cmp_main_ui["main_ui_placeholder"]
        step1_section = cmp_main_ui["step1_section"]
        step2_section = cmp_main_ui["step2_section"]
        step3_section = cmp_main_ui["step3_section"]
        step1_title = cmp_main_ui["step1_title"]
        step2_title = cmp_main_ui["step2_title"]
        step3_title = cmp_main_ui["step3_title"]
        step1_next_btn = cmp_main_ui["step1_next_btn"]
        step2_back_btn = cmp_main_ui["step2_back_btn"]
        step2_next_btn = cmp_main_ui["step2_next_btn"]
        step3_back_btn = cmp_main_ui["step3_back_btn"]
        country_inp = cmp_main_ui["country_inp"]
        language_inp = cmp_main_ui["language_inp"]
        image_inp = cmp_main_ui["image_inp"]
        image_preview = cmp_main_ui["image_preview"]
        image_preview_step3 = cmp_main_ui["image_preview_step3"]
        image_url_inp = cmp_main_ui["image_url_inp"]
        load_image_url_btn = cmp_main_ui["load_image_url_btn"]
        long_caption_inp = cmp_main_ui["long_caption_inp"]
        num_words_inp = cmp_main_ui["num_words_inp"]
        category_btn = cmp_main_ui["category_btn"]
        concept_btn = cmp_main_ui["concept_btn"]
        category_concept_dropdowns = cmp_main_ui["category_concept_dropdowns"]
        back_btn = cmp_main_ui["back_btn"]
        examples_btn = cmp_main_ui["examples_btn"]
        examples_modal = cmp_main_ui["examples_modal"]
        examples_modal_title_md = cmp_main_ui["examples_modal_title_md"]
        examples_modal_content_md = cmp_main_ui["examples_modal_content_md"]
        examples_modal_close_btn = cmp_main_ui["examples_modal_close_btn"]
        clear_btn = cmp_main_ui["clear_btn"]
        hide_faces_btn = cmp_main_ui["hide_faces_btn"]
        hide_all_faces_btn = cmp_main_ui["hide_all_faces_btn"]
        unhide_faces_btn = cmp_main_ui["unhide_faces_btn"]
        submit_btn = cmp_main_ui["submit_btn"]
        timestamp_btn = cmp_main_ui["timestamp_btn"]
        exampleid_btn = cmp_main_ui["exampleid_btn"]
        username_inp = cmp_main_ui["username_inp"]
        # password_inp = cmp_main_ui["password_inp"]
        modal_saving = cmp_main_ui["modal_saving"]
        modal_data_saved = cmp_main_ui["modal_data_saved"]
        modal = cmp_main_ui["modal"]
        exit_btn = cmp_main_ui["exit_btn"]
        intro_text_inp = cmp_main_ui["intro_text_inp"] # Note intro_text is actually the "Task" text from metadata
        intro_text_inst_inp = cmp_main_ui["intro_text_inst_inp"] # This is the "Instructions" text from metadata
        modal_saving_text = cmp_main_ui["modal_saving_text"]
        modal_data_saved_text = cmp_main_ui["modal_data_saved_text"]
        modal_exclude_confirm = cmp_main_ui["modal_exclude_confirm"]
        cancel_exclude_btn = cmp_main_ui["cancel_exclude_btn"]
        confirm_exclude_btn = cmp_main_ui["confirm_exclude_btn"]
        vlm_output = cmp_main_ui["vlm_output"]
        gen_button = cmp_main_ui["gen_button"]
        vlm_feedback = cmp_main_ui["vlm_feedback"]
        vlm_model_dropdown = cmp_main_ui["vlm_model_dropdown"]
        user_page_btn = cmp_main_ui["user_page_btn"]
        
        # Image rotation buttons
        rotate_left_btn = cmp_main_ui["rotate_left_btn"]
        rotate_right_btn = cmp_main_ui["rotate_right_btn"]
        reset_image_btn = cmp_main_ui["reset_image_btn"]

        # Multilingual text variables
        step1_md_cmp = cmp_main_ui["step1_md_cmp"]
        step2_md_cmp = cmp_main_ui["step2_md_cmp"]
        step3_md_cmp = cmp_main_ui["step3_md_cmp"]
        additional_concepts_title = cmp_main_ui["additional_concepts_title"]
        additional_concepts_instructions = cmp_main_ui["additional_concepts_instructions"]
        exclude_modal_title = cmp_main_ui["exclude_modal_title"]
        exclude_modal_text = cmp_main_ui["exclude_modal_text"]
        step1_title = cmp_main_ui["step1_title"]
        step2_title = cmp_main_ui["step2_title"]
        step3_title = cmp_main_ui["step3_title"]
        vlm_description = cmp_main_ui["vlm_description"]
        vlm_accordion = cmp_main_ui["vlm_accordion"]
        quick_submit_checkbox = cmp_main_ui["quick_submit_checkbox"]
        image_accordion = cmp_main_ui["image_accordion"]
        image_accordion_content = cmp_main_ui["image_accordion_content"]
        description_accordion = cmp_main_ui["description_accordion"]
        description_accordion_content = cmp_main_ui["description_accordion_content"]
        app_title = cmp_main_ui["app_title"]

        # ── Third page: User Profile ──
        cmp_user_page = build_user_page(concepts_dict)
        user_page_placeholder = cmp_user_page["user_page_placeholder"]
        user_page_back_btn = cmp_user_page["user_page_back_btn"]
        user_page_title = cmp_user_page["user_page_title"]
        comparative_stats_html = cmp_user_page["comparative_stats_html"]
        user_centric_stats_html = cmp_user_page["user_centric_stats_html"]
        missing_concepts_html = cmp_user_page["missing_concepts_html"]
        up_country_filter = cmp_user_page["up_country_filter"]
        up_language_filter = cmp_user_page["up_language_filter"]
        up_category_filter = cmp_user_page["up_category_filter"]
        up_concept_filter = cmp_user_page["up_concept_filter"]
        up_clear_filters_btn = cmp_user_page["up_clear_filters_btn"]
        up_data_table = cmp_user_page["up_data_table"]
        up_delete_btn = cmp_user_page["up_delete_btn"]
        up_vlm_feedback_radio = cmp_user_page["up_vlm_feedback_radio"]
        up_modal_delete_confirm = cmp_user_page["up_modal_delete_confirm"]
        up_cancel_delete_btn = cmp_user_page["up_cancel_delete_btn"]
        up_confirm_delete_btn = cmp_user_page["up_confirm_delete_btn"]
        up_hidden_concept_trigger = cmp_user_page["up_hidden_concept_trigger"]
        up_comparative_stats_title = cmp_user_page["up_comparative_stats_title"]
        up_comparative_stats_desc = cmp_user_page["up_comparative_stats_desc"]
        up_contribution_metrics_title = cmp_user_page["up_contribution_metrics_title"]
        up_missing_concepts_accordion = cmp_user_page["up_missing_concepts_accordion"]
        up_contributed_data_title = cmp_user_page["up_contributed_data_title"]
        up_delete_modal_title = cmp_user_page["up_delete_modal_title"]
        up_delete_modal_text = cmp_user_page["up_delete_modal_text"]


        # State for user page raw data (list of dicts)
        up_all_samples_raw = gr.State([])
        up_selected_example_id = gr.State(None)
        up_selected_country = gr.State(None)
        up_selected_language = gr.State(None)

        # VLM metadata indexed by example ID for browse-table editing.
        vlm_captions = gr.State(None)
        vlm_models = gr.State(None)
        vlm_feedbacks = gr.State(None)
        vlm_usages = gr.State(None)
        additional_concepts_map = gr.State(None)
        ### Category button
        category_btn.change(
            fn=partial(load_concepts, concepts=concepts_dict),
            inputs=[category_btn, concept_btn, local_storage, loading_example],
            outputs=[concept_btn, loading_example]
        )
        concept_btn.change(
            fn=partial(update_category, concepts=concepts_dict),
            inputs=[category_btn, concept_btn, local_storage, loading_example],
            outputs=[category_btn, loading_example]
        )
        def _update_category_state(category):
            return category

        def _update_concept_state(concept):
            return concept

        gr.on(
            triggers=[category_btn.change],
            fn=_update_category_state,
            inputs=[category_btn],
            outputs=[category_state],
        )
        gr.on(
            triggers=[concept_btn.change],
            fn=_update_concept_state,
            inputs=[concept_btn],
            outputs=[concept_state],
        )
        def _passthru(x):
            return x
        for i, dd in enumerate(category_concept_dropdowns):
            state_var = [category_1_concepts_state, category_2_concepts_state, category_3_concepts_state,
                         category_4_concepts_state, category_5_concepts_state][i]
            gr.on(triggers=[dd.change], fn=_passthru, inputs=[dd], outputs=[state_var])
        gr.on(triggers=[exampleid_btn.change], fn=lambda x: x, inputs=[exampleid_btn], outputs=[exampleid_state])
        # .then(
        #     fn=partial(load_concepts, concepts=concepts_dict),
        #     inputs=[category_btn, concept_btn, local_storage, gr.State(True)],
        #     outputs=[concept_btn, loading_example]
        # )
        ####

        # is_blurred = gr.State(False)  # Initialize as False

        # Notification image is false
        with Modal(visible=False) as modal_img_url:
            gr.Markdown("The image URL is not valid. Please provide a valid image URL.")
            # gr.Warning(f"⚠️ The image URL is not valid. Please provide a valid image URL.")

        # Event listeners
        gr.on(
            triggers=[country_choice.change, language_choice.change],
            fn=validate_metadata,
            inputs=[country_choice, language_choice],
            outputs=[proceed_btn],
        )

        ori_img = gr.State(None)
        image_tool_buttons = [rotate_left_btn, rotate_right_btn, reset_image_btn]

        def _copy_image_for_original_state(image):
            if image is None:
                return None
            return image.copy() if hasattr(image, "copy") else image

        def _image_tool_updates(image):
            has_image = image is not None
            return tuple(gr.update(interactive=has_image) for _ in image_tool_buttons)

        def _disable_image_tools():
            return tuple(gr.update(interactive=False) for _ in image_tool_buttons)

        image_replacement_outputs = [
            exampleid_btn,
            exampleid_state,
            vlm_output,
            vlm_feedback,
            gen_button,
            vlm_model_dropdown,
            vlm_caption_state,
            vlm_feedback_state,
            vlm_model_state,
            vlm_image_signature_state,
            vlm_usage_state,
        ]

        def _reset_image_replacement_state():
            return (
                None,
                None,
                gr.update(value=None),
                gr.update(value=None, visible=False),
                gr.update(interactive=True, visible=False),
                gr.update(value=None, interactive=False, visible=False),
                None,
                None,
                None,
                None,
                None,
            )

        def _clear_image_url_state():
            return gr.update(value=None), None

        load_image_url_btn.click(
            fn=_disable_image_tools,
            inputs=None,
            outputs=image_tool_buttons,
        ).then(
            fn=update_image,
            inputs=[image_url_inp],
            outputs=[image_inp, modal_img_url],
        ).then(
            fn=_copy_image_for_original_state,
            inputs=[image_inp],
            outputs=[ori_img],
        ).then(
            fn=_reset_image_replacement_state,
            outputs=image_replacement_outputs,
        ).then(
            fn=sync_preview,
            inputs=[image_inp],
            outputs=[image_preview, image_preview_step3],
        ).then(
            fn=_image_tool_updates,
            inputs=[image_inp],
            outputs=image_tool_buttons,
        )
        image_url_inp.submit(
            fn=_disable_image_tools,
            inputs=None,
            outputs=image_tool_buttons,
        ).then(
            fn=update_image,
            inputs=[image_url_inp],
            outputs=[image_inp, modal_img_url],
        ).then(
            fn=_copy_image_for_original_state,
            inputs=[image_inp],
            outputs=[ori_img],
        ).then(
            fn=_reset_image_replacement_state,
            outputs=image_replacement_outputs,
        ).then(
            fn=sync_preview,
            inputs=[image_inp],
            outputs=[image_preview, image_preview_step3],
        ).then(
            fn=_image_tool_updates,
            inputs=[image_inp],
            outputs=image_tool_buttons,
        )
        image_inp.upload(
            fn=_copy_image_for_original_state,
            inputs=[image_inp],
            outputs=[ori_img],
        ).then(
            fn=_reset_image_replacement_state,
            outputs=image_replacement_outputs,
        ).then(
            fn=_clear_image_url_state,
            outputs=[image_url_inp, image_url_state],
        )
        image_inp.clear(
            fn=lambda: None,
            outputs=[ori_img],
        ).then(
            fn=_reset_image_replacement_state,
            outputs=image_replacement_outputs,
        ).then(
            fn=_clear_image_url_state,
            outputs=[image_url_inp, image_url_state],
        )
        gr.on(
            triggers=[image_url_inp.change],
            fn=lambda x: x,
            inputs=[image_url_inp],
            outputs=[image_url_state],
        )
        gr.on(
            triggers=[language_choice.change],
            fn=partial(update_intro_language, metadata=metadata_dict),
            inputs=[country_choice, language_choice],
            outputs=[intro_title, intro_markdown]
        )

        gr.on(
            triggers=[language_choice.change],
            fn=partial(update_selection_language, metadata=metadata_dict),
            inputs=[country_choice, language_choice, hf_username_state],
            outputs=[signin_title, selection_title, anon_btn, country_choice, language_choice, email_inp, login_btn, proceed_btn],
        )

        def _update_timestamp_with_state():
            ts = datetime.datetime.now().timestamp()
            return gr.update(value=ts), ts

        gr.on(
            triggers=[image_inp.change, long_caption_inp.change],
            fn=_update_timestamp_with_state,
            outputs=[timestamp_btn, timestamp_state]
        )
        image_state = gr.State(None)
        gr.on(
            triggers=[
                image_inp.change, category_btn.change, concept_btn.change,
                long_caption_inp.change, quick_submit_checkbox.change
            ],
            fn=validate_inputs,
            inputs=[image_inp, ori_img, category_btn, concept_btn, long_caption_inp, quick_submit_checkbox], # is_blurred
            outputs=[submit_btn, image_inp, ori_img], # is_blurred
        )
        gr.on(
            triggers=[image_inp.change],
            fn=lambda img: img,
            inputs=[image_inp],
            outputs=[image_state],
        )
        gr.on(
            triggers=[image_inp.change],
            fn=lambda img: gr.update(interactive=(img is not None)),
            inputs=[image_inp],
            outputs=[step1_next_btn],
        )
        gr.on(
            triggers=[image_inp.change],
            fn=_image_tool_updates,
            inputs=[image_inp],
            outputs=image_tool_buttons,
        )

        def validate_step2_next(category, concept):
            category_empty = category is None or (isinstance(category, str) and category.strip() == "")
            concept_empty = concept is None or (isinstance(concept, str) and concept.strip() == "")
            return gr.update(interactive=not (category_empty or concept_empty))

        gr.on(
            triggers=[category_btn.change, concept_btn.change],
            fn=validate_step2_next,
            inputs=[category_btn, concept_btn],
            outputs=[step2_next_btn],
        )
        gr.on(
            triggers=[image_inp.change],
            fn=sync_preview,
            inputs=[image_inp],
            outputs=[image_preview, image_preview_step3],
        )

        def _log_caption_and_count(caption, language):
            return count_words(caption, language)

        gr.on(
            triggers=[long_caption_inp.change],
            fn=_log_caption_and_count,
            inputs=[long_caption_inp, language_choice],
            outputs=[num_words_inp],
        )
        def _update_caption_state(caption):
            return caption

        gr.on(
            triggers=[long_caption_inp.change],
            fn=_update_caption_state,
            inputs=[long_caption_inp],
            outputs=[caption_state],
        )

        #============= Face Blurring ============= #
        with Modal(visible=False) as modal_faces:
            with gr.Column():
                face_img = gr.Image(label="Image Faces", elem_id="image_faces", format="png", height=512, width=768)
                with gr.Row():
                    faces_count = gr.Textbox(label="Face Counts", elem_id="face_counts", interactive=False)
                    blur_faces_ids = gr.Dropdown(
                        [], value=[], multiselect=True, label="Please select the faces IDs you want to blur.", elem_id="blur_faces_ids")
                    # blur_faces_ids = gr.Textbox(label="Specify faces ids to blur by comma", elem_id="blur_faces_ids", interactive=True)
                    submit_btn_face = gr.Button("Submit", variant="primary", interactive=True, elem_id="submit_btn_face")
        
        faces_info = gr.State(None)
        hide_faces_btn.click(
            fn=select_faces_to_hide,
            inputs=[image_inp, blur_faces_ids],
            outputs=[image_inp, modal_faces, face_img, faces_count, faces_info, blur_faces_ids]
        )

        submit_btn_face.click(
            fn=blur_selected_faces,
            inputs=[image_inp, blur_faces_ids, faces_info, face_img, faces_count], # is_blurred 
            outputs=[image_inp, modal_faces, face_img, faces_count, blur_faces_ids] # is_blurred
        )

        hide_all_faces_btn.click(
            fn=blur_all_faces,
            inputs=[image_inp],
            outputs=[image_inp]
        )

        unhide_faces_btn.click(
            fn=unhide_faces,
            inputs=[image_inp, ori_img], # is_blurred
            outputs=[image_inp] # is_blurred
        )

        # ============= Image Rotation ============= #
        rotate_left_btn.click(
            fn=rotate_image_90_left,
            inputs=[image_inp],
            outputs=[image_inp]
        )

        rotate_right_btn.click(
            fn=rotate_image_90_right,
            inputs=[image_inp],
            outputs=[image_inp]
        )

        reset_image_btn.click(
            fn=reset_image_to_original,
            inputs=[image_inp, ori_img],
            outputs=[image_inp]
        )

        # ===============================
        show_browse_data_state = gr.State(True)
        with gr.Column(visible=False, elem_id="browse_toggle_row") as browse_toggle_row:
            with gr.Row(equal_height=True):
                clear_top_btn = gr.Button("Clear", variant="huggingface", elem_id="clear_top_btn")
                exclude_btn = gr.Button("Delete", variant="stop", elem_id="exclude_btn", interactive=False)
            with gr.Row(equal_height=True):
                toggle_browse_btn = gr.Button("▼ Show/Hide Data", elem_id="toggle_browse_btn")
        with gr.Column(visible=False, elem_id="browse_data") as browse_data_placeholder:

            loading_msg = gr.Markdown("**Loading your data, please wait ...**")
            with gr.Tab("Your data") as tab_your_data:
                # Show user's past data points
                user_examples = gr.Dataset(
                    samples=[],
                    show_label=False,
                    components=['image','textbox','textbox','textbox','textbox',
                               'textbox','textbox','textbox', 'textbox'],
                    headers=['Image', 'Image URL (Optional, if not uploading an image)', 'Description', 'Country', 'Language',
                            'Category', 'Concept', 'Additional Concepts', 'ID'],
                )
                # Handle clicking on an example
                user_examples.click(
                    fn=partial(handle_click_example_with_states, concepts_dict=concepts_dict),
                    inputs=[
                        user_examples,
                        vlm_captions,
                        vlm_models,
                        vlm_feedbacks,
                        vlm_usages,
                        additional_concepts_map,
                    ],
                    outputs=[
                        image_inp, image_url_inp, long_caption_inp, exampleid_btn,
                        category_btn, concept_btn,
                        category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2],
                        category_concept_dropdowns[3], category_concept_dropdowns[4], loading_example, vlm_output, vlm_model_dropdown,
                        vlm_feedback,
                        image_state, image_url_state, caption_state, category_state, concept_state,
                        category_1_concepts_state, category_2_concepts_state, category_3_concepts_state,
                        category_4_concepts_state, category_5_concepts_state,
                        exampleid_state, vlm_caption_state, vlm_model_state, vlm_feedback_state, vlm_usage_state,
                        timestamp_state, vlm_image_signature_state,
                    ],
                ).then(
                    fn=_copy_image_for_original_state,
                    inputs=[image_inp],
                    outputs=[ori_img],
                ).then(
                    fn=sync_preview,
                    inputs=[image_inp],
                    outputs=[image_preview, image_preview_step3],
                ).then(
                    fn=lambda ts: gr.update(value=ts, label="Timestamp", visible=False),
                    inputs=[timestamp_state],
                    outputs=[timestamp_btn],
                ).then(
                    fn=lambda: gr.update(interactive=True),
                    outputs=[exclude_btn],
                )

        toggle_browse_btn.click(
            fn=toggle_browse_data,
            inputs=[show_browse_data_state, toggle_browse_btn],
            outputs=[show_browse_data_state, browse_data_placeholder, toggle_browse_btn],
        )

        # Clear Button (must be after exclude_btn is defined)
        clear_outputs = [
            image_inp, image_url_inp, long_caption_inp, vlm_output, vlm_feedback, gen_button, vlm_model_dropdown, exampleid_btn,
            category_btn, concept_btn,
            category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2],
            category_concept_dropdowns[3], category_concept_dropdowns[4],
            category_state, concept_state, caption_state,
            image_url_state, quick_submit_state,
            category_1_concepts_state, category_2_concepts_state, category_3_concepts_state,
            category_4_concepts_state, category_5_concepts_state,
            timestamp_state, exampleid_state,
            vlm_caption_state, vlm_feedback_state, vlm_model_state, vlm_usage_state,
            exclude_btn, ori_img,
        ]
        for clear_button in (clear_btn, clear_top_btn):
            clear_button.click(
                fn=partial(clear_data, metadata_dict=metadata_dict),
                inputs=[local_storage],
                outputs=clear_outputs,
            ).then(
                fn=sync_preview,
                inputs=[image_inp],
                outputs=[image_preview, image_preview_step3],
            )

        # Step navigation
        step1_next_btn.click(
            fn=lambda show_browse, storage: (2, *set_step_from_storage(2, show_browse, storage), gr.update(interactive=True)),
            inputs=[show_browse_data_state, local_storage],
            outputs=[current_step_state, step1_section, step2_section, step3_section, browse_toggle_row, browse_data_placeholder, back_btn],
            scroll_to_output=False,
            show_progress_on=step2_title,
            js=scroll_to_top_js,
        ).then(
            fn=sync_preview,
            inputs=[image_inp],
            outputs=[image_preview, image_preview_step3],
            scroll_to_output=False,
        ).then(
            fn=partial(auto_generate_vlm_caption, metadata_dict=metadata_dict),
            inputs=[image_inp, vlm_caption_state, vlm_model_state, vlm_image_signature_state, exampleid_state, local_storage],
            outputs=[vlm_output, vlm_feedback, gen_button, vlm_model_dropdown, vlm_caption_state, vlm_model_state, vlm_image_signature_state, vlm_usage_state],
            concurrency_limit=3,
            concurrency_id="vlm_queue",
            scroll_to_output=False,
        )
        step2_back_btn.click(
            fn=lambda show_browse, storage: (1, *set_step_from_storage(1, show_browse, storage), gr.update(interactive=False)),
            inputs=[show_browse_data_state, local_storage],
            outputs=[current_step_state, step1_section, step2_section, step3_section, browse_toggle_row, browse_data_placeholder, back_btn],
            scroll_to_output=False,
            show_progress_on=step1_title,
            js=scroll_to_top_js,
        )
        step2_next_btn.click(
            fn=lambda show_browse, storage: (3, *set_step_from_storage(3, show_browse, storage), gr.update(interactive=True)),
            inputs=[show_browse_data_state, local_storage],
            outputs=[current_step_state, step1_section, step2_section, step3_section, browse_toggle_row, browse_data_placeholder, back_btn],
            scroll_to_output=False,
            show_progress_on=step3_title,
            js=scroll_to_top_js,
        )
        step3_back_btn.click(
            fn=lambda show_browse, storage: (2, *set_step_from_storage(2, show_browse, storage), gr.update(interactive=True)),
            inputs=[show_browse_data_state, local_storage],
            outputs=[current_step_state, step1_section, step2_section, step3_section, browse_toggle_row, browse_data_placeholder, back_btn],
            scroll_to_output=False,
            show_progress_on=step2_title,
            js=scroll_to_top_js,
        )

        def handle_global_back(current_step, show_browse, storage):
            if current_step == 2:
                step_outputs = set_step_from_storage(1, show_browse, storage)
                return (1, *step_outputs, gr.update(interactive=False))
            elif current_step == 3:
                step_outputs = set_step_from_storage(2, show_browse, storage)
                return (2, *step_outputs, gr.update(interactive=True))
            else:
                step_outputs = set_step_from_storage(1, show_browse, storage)
                return (1, *step_outputs, gr.update(interactive=False))

        back_btn.click(
            fn=handle_global_back,
            inputs=[current_step_state, show_browse_data_state, local_storage],
            outputs=[current_step_state, step1_section, step2_section, step3_section, browse_toggle_row, browse_data_placeholder, back_btn],
            scroll_to_output=False,
            js=scroll_to_top_js,
        )

        # ============================================ #
        # Submit Button Click events
        # login_btn.click(
        #     fn=login_user,
        #     inputs=[username, password],
        #     outputs=[supabase_user_client, persistent_session, login_status],
        # ).then(
        #     fn=restore_user_session,
        #     inputs=[persistent_session, login_status],
        #     outputs=[supabase_user_client, persistent_session, login_status, proceed_btn, login_btn, sign_up_btn, reset_password_btn, logout_btn, change_password_field, change_password_field_confirm, change_password_btn, change_password_status],
        # )

        # sign_up_btn.click(
        #     fn=sign_up,
        #     inputs=[username, password],
        #     outputs=[login_status],
        # )

        # logout_btn.click(
        #     fn=log_out,
        #     inputs=[supabase_user_client, persistent_session],
        #     outputs=[persistent_session]
        # ).then(
        #     fn=restore_user_session,
        #     inputs=[persistent_session],
        #     outputs=[supabase_user_client, persistent_session, login_status, proceed_btn, login_btn, sign_up_btn, reset_password_btn, logout_btn, change_password_field, change_password_field_confirm, change_password_btn, change_password_status],
        # )
        # change_password_btn.click(
        #     fn=change_password,
        #     inputs=[supabase_user_client, change_password_field, change_password_field_confirm],
        #     outputs=[change_password_status]
        # )
        # reset_password_btn.click(
        #     fn=reset_password,
        #     inputs=[username],
        #     outputs=[login_status]
        # )

        def _short_email_label(email_val, max_len=30):
            email_val = (email_val or "").strip()
            if len(email_val) <= max_len:
                return email_val
            if "@" not in email_val:
                return email_val[: max_len - 3] + "..."
            local, domain = email_val.split("@", 1)
            budget = max_len - len(domain) - 4
            if budget < 4:
                return email_val[: max_len - 3] + "..."
            return f"{local[:budget]}...@{domain}"

        def _update_user_page_btn(email_val, hf_username=""):
            email_val = (email_val or "").strip()
            hf_username = (hf_username or "").strip()
            if email_val and email_val != "anonymous":
                label = f"@{hf_username}" if hf_username else _short_email_label(email_val)
                return gr.update(value=f"👤 {label}", interactive=True)
            return gr.update(value="👤 Anonymous", interactive=False)

        language_outputs = [
            country_inp, language_inp, username_inp, category_btn, concept_btn, image_inp,
            image_url_inp, long_caption_inp, intro_text_inp, intro_text_inst_inp, back_btn, clear_top_btn, clear_btn,
            submit_btn, modal_saving_text, modal_data_saved_text, timestamp_btn, exit_btn,
            loading_msg, hide_all_faces_btn, hide_faces_btn, unhide_faces_btn, exclude_btn,
            category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2],
            category_concept_dropdowns[3], category_concept_dropdowns[4],
            step2_md_cmp,
            step1_md_cmp,
            step3_md_cmp,
            num_words_inp, additional_concepts_title, additional_concepts_instructions,
            exclude_modal_title, exclude_modal_text, cancel_exclude_btn, confirm_exclude_btn,
            gen_button, vlm_model_dropdown, vlm_feedback,
            step1_title, step2_title,
            step1_next_btn, step2_back_btn,
            step2_next_btn, step3_back_btn, step3_title,
            load_image_url_btn,
            rotate_left_btn, rotate_right_btn, reset_image_btn,
            quick_submit_checkbox,
            vlm_description,
            tab_your_data,
            vlm_accordion, toggle_browse_btn,
            vlm_output,
            image_accordion, image_accordion_content,
            description_accordion, description_accordion_content,
            user_examples,
            app_title,
        ]
        examples_modal_outputs = [examples_btn, examples_modal_title_md, examples_modal_content_md, examples_modal_close_btn]
        step_outputs = [step1_section, step2_section, step3_section, browse_toggle_row, browse_data_placeholder]
        user_data_outputs = [
            user_examples,
            loading_msg,
            vlm_captions,
            vlm_models,
            vlm_feedbacks,
            vlm_usages,
            additional_concepts_map,
        ]
        user_page_language_outputs = [
            user_page_back_btn, user_page_title,
            up_comparative_stats_title, up_comparative_stats_desc,
            up_contribution_metrics_title, up_missing_concepts_accordion,
            up_contributed_data_title,
            up_country_filter, up_language_filter, up_category_filter, up_concept_filter,
            up_clear_filters_btn, up_delete_btn, up_vlm_feedback_radio,
            up_delete_modal_title, up_delete_modal_text,
            up_cancel_delete_btn, up_confirm_delete_btn,
            up_data_table,
        ]

        def _has_saved_locale(country, language):
            return bool(country and language and metadata_dict.get(country, {}).get(language))

        def _skip_outputs(outputs):
            return tuple(gr.skip() for _ in outputs)



        def _auto_extract_country_language_username(local_storage_value):
            country, language, email = extract_country_language_username(local_storage_value)
            if not _has_saved_locale(country, language):
                return _skip_outputs([country_state, language_state, username_state])
            return country, language, email

        def _auto_update_language(local_storage_value):
            country, language, _ = extract_country_language_username(local_storage_value)
            if not _has_saved_locale(country, language):
                return _skip_outputs(language_outputs)
            return update_language(local_storage_value, metadata_dict, concepts_dict)

        def _auto_update_user_page_language(local_storage_value):
            country, language, _ = extract_country_language_username(local_storage_value)
            if not _has_saved_locale(country, language):
                return _skip_outputs(user_page_language_outputs)
            return update_user_page_language(local_storage_value, metadata_dict)

        def _auto_build_examples_modal(local_storage_value):
            country, language, _ = extract_country_language_username(local_storage_value)
            if not _has_saved_locale(country, language):
                return _skip_outputs(examples_modal_outputs)
            return build_examples_modal_updates(local_storage_value, metadata_dict)

        def _maybe_show_examples_for_first_visit(local_storage_value, hf_username, email):
            country, language, _ = extract_country_language_username(local_storage_value)
            if not _has_saved_locale(country, language):
                return gr.skip()
            hf_username = (hf_username or "").strip()
            email = (email or "").strip()
            if not hf_username and _is_anonymous_email(email):
                return gr.skip()
            first_visit = prefs_store.mark_main_ui_visit(hf_username, email) if prefs_store else False
            return gr.update(visible=True) if first_visit else gr.skip()

        def _auto_set_step(local_storage_value, show_browse):
            country, language, _ = extract_country_language_username(local_storage_value)
            if not _has_saved_locale(country, language):
                return _skip_outputs(step_outputs)
            return set_step_from_storage(1, show_browse, local_storage_value)

        def _auto_update_user_data(client, country, language, email):
            if not _has_saved_locale(country, language):
                return _skip_outputs(user_data_outputs)
            return update_user_data(client, country, language, email, HF_DATASET_NAME, LOCAL_DS_DIRECTORY_PATH)

        def _auto_update_user_page_btn(local_storage_value, email, hf_username):
            country, language, _ = extract_country_language_username(local_storage_value)
            if not _has_saved_locale(country, language):
                return gr.skip()
            return _update_user_page_btn(email, hf_username)

        def _auto_reset_step_and_back(local_storage_value):
            country, language, _ = extract_country_language_username(local_storage_value)
            if not _has_saved_locale(country, language):
                return gr.skip(), gr.skip()
            return 1, gr.update(interactive=False)

        def _persist_locale_session(email, country, language, hf_username, session):
            if not _has_saved_locale(country, language):
                return session
            return dict(
                session or {},
                email=(email or "").strip(),
                country=country,
                language=language,
                hf_username=(hf_username or "").strip(),
                force_language_selection=False,
            )

        def _mark_language_selection_requested(session):
            return dict(session or {}, force_language_selection=True)

        def _persist_locale_prefs(hf_username, email, country, language):
            if not _has_saved_locale(country, language):
                return None
            return prefs_store.upsert_prefs(hf_username, email, country, language) if prefs_store else None

        proceed_btn.click(
            fn=partial(switch_ui, flag=False, metadata_dict=metadata_dict),
            inputs=[country_choice, language_choice, email_inp],
            outputs=[local_storage, selection_page, main_ui_placeholder, browse_toggle_row, browse_data_placeholder, country_choice, language_choice, signin_title, selection_title, proceed_btn],
        ).then(
            fn=extract_country_language_username,
            inputs=[local_storage],
            outputs=[country_state, language_state, username_state],
        ).then(
            fn=partial(update_language, metadata_dict=metadata_dict, concepts_dict=concepts_dict),
            inputs=[local_storage],
            outputs=language_outputs,
        ).then(
            fn=partial(update_user_page_language, metadata_dict=metadata_dict),
            inputs=[local_storage],
            outputs=user_page_language_outputs,
        ).then(
            fn=partial(build_examples_modal_updates, metadata_dict=metadata_dict),
            inputs=[local_storage],
            outputs=examples_modal_outputs,
        ).then(
            fn=_maybe_show_examples_for_first_visit,
            inputs=[local_storage, hf_username_state, email_inp],
            outputs=[examples_modal],
        ).then(
            fn=lambda show_browse, storage: (1, *set_step_from_storage(1, show_browse, storage), gr.update(interactive=False)),
            inputs=[show_browse_data_state, local_storage],
            outputs=[current_step_state, step1_section, step2_section, step3_section, browse_toggle_row, browse_data_placeholder, back_btn],
        ).then(
            fn=partial(update_user_data, HF_DATASET_NAME=HF_DATASET_NAME, local_ds_directory_path = LOCAL_DS_DIRECTORY_PATH),
            inputs=[supabase_user_client, country_choice, language_choice, email_inp],
            outputs=user_data_outputs,
        ).then(
            fn=_update_user_page_btn,
            inputs=[email_inp, hf_username_state],
            outputs=[user_page_btn],
        ).then(
            fn=_persist_locale_session,
            inputs=[email_inp, country_choice, language_choice, hf_username_state, persistent_session],
            outputs=[persistent_session],
        ).then(
            fn=lambda hf_username, email, country, language: prefs_store.upsert_prefs(
                hf_username, email, country, language
            ) if prefs_store else None,
            inputs=[hf_username_state, email_inp, country_choice, language_choice],
        )

        def _chain_saved_locale_entry(event):
            return event.then(
                fn=partial(switch_ui, flag=False, metadata_dict=metadata_dict),
                inputs=[country_choice, language_choice, email_inp],
                outputs=[local_storage, selection_page, main_ui_placeholder, browse_toggle_row, browse_data_placeholder, country_choice, language_choice, signin_title, selection_title, proceed_btn],
            ).then(
                fn=_auto_extract_country_language_username,
                inputs=[local_storage],
                outputs=[country_state, language_state, username_state],
            ).then(
                fn=_auto_update_language,
                inputs=[local_storage],
                outputs=language_outputs,
            ).then(
                fn=_auto_update_user_page_language,
                inputs=[local_storage],
                outputs=user_page_language_outputs,
            ).then(
                fn=_auto_build_examples_modal,
                inputs=[local_storage],
                outputs=examples_modal_outputs,
            ).then(
                fn=_maybe_show_examples_for_first_visit,
                inputs=[local_storage, hf_username_state, email_inp],
                outputs=[examples_modal],
            ).then(
                fn=_auto_set_step,
                inputs=[local_storage, show_browse_data_state],
                outputs=step_outputs,
            ).then(
                fn=_auto_update_user_data,
                inputs=[supabase_user_client, country_choice, language_choice, email_inp],
                outputs=user_data_outputs,
            ).then(
                fn=_auto_update_user_page_btn,
                inputs=[local_storage, email_inp, hf_username_state],
                outputs=[user_page_btn],
            ).then(
                fn=_persist_locale_session,
                inputs=[email_inp, country_choice, language_choice, hf_username_state, persistent_session],
                outputs=[persistent_session],
            ).then(
                fn=_persist_locale_prefs,
                inputs=[hf_username_state, email_inp, country_choice, language_choice],
            ).then(
                fn=_auto_reset_step_and_back,
                inputs=[local_storage],
                outputs=[current_step_state, back_btn],
            )

        _chain_saved_locale_entry(profile_continue_event)
        _chain_saved_locale_entry(anon_continue_event)


        # Exit Button
        exit_btn.click(
            fn=exit_fn,
            outputs=[
                image_inp, image_url_inp, long_caption_inp, vlm_output, vlm_feedback, gen_button, vlm_model_dropdown, user_examples, loading_msg,
                username, local_storage, exampleid_btn, category_btn, concept_btn,
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2],
                category_concept_dropdowns[3], category_concept_dropdowns[4], ori_img
            ],
        ).success(
            fn=partial(switch_ui, flag=True, metadata_dict=metadata_dict),
            inputs=[country_choice, language_choice, username_inp],
            outputs=[local_storage, selection_page, main_ui_placeholder, browse_toggle_row, browse_data_placeholder, country_choice, language_choice, signin_title, selection_title, proceed_btn],
        ).success(
            fn=lambda: gr.update(value="👤 Anonymous", interactive=False),
            outputs=[user_page_btn],
        ).success(
            fn=_mark_language_selection_requested,
            inputs=[persistent_session],
            outputs=[persistent_session],
        )

        # Disable button while saving
        # Note: I think this is not longer needed as we clear all the inputs and disable the submit button if the data is saved correctly
        # def disable_submit():
        #     return gr.update(interactive=False)

        # def enable_submit():
        #     return gr.update(interactive=True)

        # STEP 1: show modal
        # submit_btn.click(lambda: Modal(visible=True), None, modal_saving)
        
        # STEP 2: disable button
        # submit_btn.click(disable_submit, None, [submit_btn], queue=False)

        #STEP 3: perform save_data
        gr.on(triggers=[quick_submit_checkbox.change], fn=lambda x: x if x is not None else False, inputs=[quick_submit_checkbox], outputs=[quick_submit_state])
        data_outputs = {
            "image": image_state,
            "image_url": image_url_state,
            "caption": caption_state,
            "quick_submit_mode": quick_submit_state,
            "country": country_state,
            "language": language_state,
            "category": category_state,
            "concept": concept_state,
            "category_1_concepts": category_1_concepts_state,
            "category_2_concepts": category_2_concepts_state,
            "category_3_concepts": category_3_concepts_state,
            "category_4_concepts": category_4_concepts_state,
            "category_5_concepts": category_5_concepts_state,
            "timestamp": timestamp_state,
            "username": username_state,
            "hf_username": hf_username_state,
            "id": exampleid_state,
            "excluded": gr.State(value=False),
            "concepts_dict": gr.State(value=concepts_dict),
            "country_lang_map": gr.State(value=lang2eng_mapping),
            "client": supabase_user_client,
            "vlm_caption": vlm_caption_state,
            "vlm_feedback": vlm_feedback_state,
            "vlm_model": vlm_model_state,
            "vlm_usage": vlm_usage_state,
        }
        # data_outputs = [image_inp, image_url_inp, long_caption_inp,
        #         country_inp, language_inp, category_btn, concept_btn,
        #         timestamp_btn, username_inp, password_inp, exampleid_btn]
        hf_writer.setup(list(data_outputs.keys()), local_ds_folder=LOCAL_DS_DIRECTORY_PATH, metadata_dict=metadata_dict)

        # STEP 4: Chain save_data, then update_user_data, then clear
        def _log_and_save(*values):
            return hf_writer.save(*values)

        # ── Rebuild user profile helper (used by submit + delete flows) ──
        def _rebuild_profile_after_save(email_val):
            email_val = (email_val or "").strip()
            if email_val:
                rebuild_user_profile(email_val, LOCAL_DS_DIRECTORY_PATH)

        def wire_submit_chain(button, js=None):
            e = button.click(
                _log_and_save,
                inputs=list(data_outputs.values()),
                outputs=None,
                js=js,
            ).success(
                fn=partial(clear_data, "submit", metadata_dict=metadata_dict),
                inputs=[local_storage],
                outputs=[
                    image_inp, image_url_inp, long_caption_inp, vlm_output, vlm_feedback, gen_button, vlm_model_dropdown, exampleid_btn,
                    category_btn, concept_btn,
                    category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2],
                    category_concept_dropdowns[3], category_concept_dropdowns[4],
                    category_state, concept_state, caption_state,
                    image_url_state, quick_submit_state,
                    category_1_concepts_state, category_2_concepts_state, category_3_concepts_state,
                    category_4_concepts_state, category_5_concepts_state,
                    timestamp_state, exampleid_state,
                    vlm_caption_state, vlm_feedback_state, vlm_model_state, vlm_usage_state,
                    exclude_btn, ori_img,
                ],
            ).success(
                # Navigate back to Step 1 immediately after save+clear succeed.
                # Performed early so any transient error in the data-reload steps
                # below can never leave the user stranded on Step 2.
                fn=lambda show_browse, storage: (1, *set_step_from_storage(1, show_browse, storage), gr.update(interactive=False)),
                inputs=[show_browse_data_state, local_storage],
                outputs=[current_step_state, step1_section, step2_section, step3_section, browse_toggle_row, browse_data_placeholder, back_btn],
            ).then(
                lambda: gr.update(value="**Loading your data, please wait ...**"),
                None, loading_msg
            ).then(
                fn=partial(update_user_data, HF_DATASET_NAME=HF_DATASET_NAME, local_ds_directory_path=LOCAL_DS_DIRECTORY_PATH, skip_snapshot=True),
                inputs=[supabase_user_client, country_choice, language_choice, username_inp],
                outputs=user_data_outputs
            ).then(
                fn=_rebuild_profile_after_save,
                inputs=[username_inp],
            ).then(
                fn=lambda hf_username, email, country, language: prefs_store.upsert_prefs(
                    hf_username, email, country, language
                ) if prefs_store else None,
                inputs=[hf_username_state, username_inp, country_choice, language_choice],
            )
            return e
        
        wire_submit_chain(submit_btn, js=scroll_to_top_js)


        # ============================================ #
        # "See Examples" button — opens a read-only reference modal
        # describing what/how to fill in each field. No click handlers on
        # the displayed examples (those caused the old Examples-tab error).
        examples_btn.click(lambda: gr.update(visible=True), None, examples_modal)
        examples_modal_close_btn.click(lambda: gr.update(visible=False), None, examples_modal)
        # ============================================ #

        # # Load saved values from local storage (browser storage)
        # @ui.load(inputs=[local_storage], outputs=[country_choice, language_choice, username, password])
        # def load_from_local_storage(saved_values):
        #     print("loading from local storage", saved_values)
        #     return saved_values[0], saved_values[1], saved_values[2], saved_values[3]
        
        # ============================================= #
        # Exclude button 
        # ============================================= #

        # Show confirmation modal when exclude button is clicked
        exclude_btn.click(
            fn=check_exclude_fn,
            inputs=[exampleid_state],
            outputs=[modal_exclude_confirm]
        )

        # Close modal when cancel button is clicked
        cancel_exclude_btn.click(
            fn=lambda: gr.update(visible=False),
            outputs=[modal_exclude_confirm]
        )

        def _do_exclude(example_id, country, language, username):
            mark_sample_excluded(example_id, country, language, username, LOCAL_DS_DIRECTORY_PATH)
            # Keep profiles.json in sync
            _rebuild_profile_after_save(username)

        confirm_exclude_btn.click(
            fn=lambda: gr.update(visible=False),
            outputs=[modal_exclude_confirm]
        ).success(
            fn=_do_exclude,
            inputs=[exampleid_state, country_state, language_state, username_state],
            outputs=None
        ).success(
            fn=partial(clear_data, "remove", metadata_dict=metadata_dict),
            inputs=[local_storage],
            outputs=[
                image_inp, image_url_inp, long_caption_inp, vlm_output, vlm_feedback, gen_button, vlm_model_dropdown, exampleid_btn,
                category_btn, concept_btn,
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2],
                category_concept_dropdowns[3], category_concept_dropdowns[4],
                category_state, concept_state, caption_state,
                image_url_state, quick_submit_state,
                category_1_concepts_state, category_2_concepts_state, category_3_concepts_state,
                category_4_concepts_state, category_5_concepts_state,
                timestamp_state, exampleid_state,
                vlm_caption_state, vlm_feedback_state, vlm_model_state, vlm_usage_state,
                exclude_btn, ori_img,
            ]
        ).success(
            fn=lambda: gr.update(value="**Refreshing your data, please wait...**"),
            outputs=loading_msg
        ).success(
            fn=partial(update_user_data, HF_DATASET_NAME=HF_DATASET_NAME, local_ds_directory_path=LOCAL_DS_DIRECTORY_PATH, skip_snapshot=True),
            inputs=[supabase_user_client, country_choice, language_choice, username_inp],
            outputs=user_data_outputs
        )
        # ============================================= #
        # VLM Gen button
        # ============================================= #
        gen_button.click(
            fn=partial(regenerate_vlm_caption, metadata_dict=metadata_dict),
            inputs=[image_inp, local_storage],
            outputs=[vlm_output, vlm_feedback, gen_button, vlm_model_dropdown, vlm_caption_state, vlm_model_state, vlm_image_signature_state, vlm_usage_state],
            # OpenAI API replaces the on-device VLM, so the GPU bottleneck is
            # gone. Keep a small queue (3) to stay polite to the API.
            concurrency_limit=3,
            concurrency_id="vlm_queue"
        )
        gr.on(triggers=[vlm_output.change], fn=_passthru, inputs=[vlm_output], outputs=[vlm_caption_state])
        gr.on(triggers=[vlm_feedback.change], fn=_passthru, inputs=[vlm_feedback], outputs=[vlm_feedback_state])
        # vlm_output.change(
        #     fn=lambda : gr.update(interactive=False) if vlm_output.value else gr.update(interactive=True),
        #     inputs=[],
        #     outputs=[gen_button]
        # )

        # ui.load(
        #     fn=login_user_recovery,
        #     inputs=gr.Textbox(visible=False, value=""), # hidden textbox to get the url tokens
        #     outputs=[supabase_user_client, persistent_session, login_status],
        #     js=js_code
        # ).then(
        #     fn=restore_user_session,
        #     inputs=[persistent_session],
        #     outputs=[supabase_user_client, persistent_session, login_status, proceed_btn, login_btn, sign_up_btn, reset_password_btn, logout_btn, change_password_field, change_password_field_confirm, change_password_btn, change_password_status],
        # )
        # ============================================= #
        # User Page wiring
        # ============================================= #

        def _open_user_page(email_val):
            """Navigate to user page — hide main UI, show user page."""
            email_val = (email_val or "").strip()
            if not email_val:
                gr.Warning("⚠️ Please sign in to view your profile.")
                return (
                    gr.skip(),  # main_ui_placeholder
                    gr.skip(),  # user_page_placeholder
                    gr.skip(),  # browse_toggle_row
                    gr.skip(),  # browse_data_placeholder
                )
            return (
                gr.update(visible=False),   # main_ui_placeholder
                gr.update(visible=True),    # user_page_placeholder
                gr.update(visible=False),   # browse_toggle_row
                gr.update(visible=False),   # browse_data_placeholder
            )

        def _load_user_page_data(email_val, local_storage_val):
            """Load all user data and stats for user page."""
            email_val = (email_val or "").strip()
            country, language, _ = extract_country_language_username(local_storage_val)
            meta = (metadata_dict or {}).get(country, {}).get(language, {}) if country and language else {}
            base_title = meta.get("UP_Title", "## 👤 Your Profile")
            if not email_val:
                return (
                    gr.Dataset(samples=[]),
                    gr.update(choices=[]),
                    gr.update(choices=[]),
                    gr.update(choices=[]),
                    gr.update(choices=[]),
                    [],  # up_all_samples_raw
                    "<p>No data.</p>",
                    "<p>No data.</p>",
                    "<p>No data.</p>",
                    f"{base_title} — Anonymous",
                )
            table_rows, countries, languages, categories, concepts, raw_samples = \
                load_all_user_data(email_val, HF_DATASET_NAME, LOCAL_DS_DIRECTORY_PATH, skip_snapshot=False)

            # Compute stats
            rebuild_user_profile(email_val, LOCAL_DS_DIRECTORY_PATH)
            comp_stats = compute_comparative_stats(email_val, LOCAL_DS_DIRECTORY_PATH)
            user_stats = compute_user_centric_stats(email_val, LOCAL_DS_DIRECTORY_PATH, concepts_dict)

            return (
                gr.Dataset(samples=table_rows),
                gr.update(choices=[None] + countries, value=None),
                gr.update(choices=[None] + languages, value=None),
                gr.update(choices=[None] + categories, value=None),
                gr.update(choices=[None] + concepts, value=None),
                raw_samples,
                render_comparative_stats(comp_stats, meta),
                render_user_centric_stats(user_stats, meta),
                render_missing_concepts(user_stats, meta),
                f"{base_title} — {email_val}",
            )

        user_page_btn.click(
            fn=_open_user_page,
            inputs=[username_inp],
            outputs=[main_ui_placeholder, user_page_placeholder, browse_toggle_row, browse_data_placeholder],
        ).then(
            fn=_load_user_page_data,
            inputs=[username_inp, local_storage],
            outputs=[
                up_data_table,
                up_country_filter, up_language_filter, up_category_filter, up_concept_filter,
                up_all_samples_raw,
                comparative_stats_html, user_centric_stats_html, missing_concepts_html,
                user_page_title,
            ],
        )

        # Back button — return to annotation page
        def _close_user_page():
            return (
                gr.update(visible=True),    # main_ui_placeholder
                gr.update(visible=False),   # user_page_placeholder
                gr.update(visible=True),    # browse_toggle_row
                gr.update(visible=True),    # browse_data_placeholder
            )

        user_page_back_btn.click(
            fn=_close_user_page,
            outputs=[main_ui_placeholder, user_page_placeholder, browse_toggle_row, browse_data_placeholder],
        )

        # ── Filters ──
        def _apply_filters(raw, country_f, lang_f, cat_f, concept_f):
            return filter_user_page_data(raw, country_f, lang_f, cat_f, concept_f)

        for _trigger in [up_country_filter, up_language_filter, up_category_filter, up_concept_filter]:
            _trigger.change(
                fn=_apply_filters,
                inputs=[up_all_samples_raw, up_country_filter, up_language_filter, up_category_filter, up_concept_filter],
                outputs=[up_data_table],
            )

        up_clear_filters_btn.click(
            fn=lambda raw: (
                gr.update(value=None), gr.update(value=None),
                gr.update(value=None), gr.update(value=None),
                gr.Dataset(samples=[[s["image"], s["caption"], s["country"], s["language"],
                    s["category"], s["concept"], s["vlm_caption"], s["id"]] for s in raw] if raw else []),
            ),
            inputs=[up_all_samples_raw],
            outputs=[up_country_filter, up_language_filter, up_category_filter, up_concept_filter, up_data_table],
        )

        # ── Row selection (clicking on a row in the Dataset) ──
        def _on_user_page_row_click_load(clicked_row, all_samples_raw):
            row = list(clicked_row)
            img = row[0]
            caption = row[1]
            country = row[2]
            language = row[3]
            category = row[4]
            concept = row[5]
            categories = concepts_dict.get(country, {}).get(lang2eng_mapping.get(language, language), {})
            category = _canonical_category_value(category, language, categories)
            if category is None and concept:
                matches = [cat for cat, values in categories.items() if concept in (values or [])]
                category = matches[0] if len(matches) == 1 else None
            vlm_cap = row[6]
            example_id = row[7]

            saved = next(
                (sample for sample in (all_samples_raw or []) if sample.get("id") == example_id),
                {},
            )
            saved_model = saved.get("vlm_model") or None
            saved_feedback = saved.get("vlm_feedback") or None
            saved_usage = saved.get("vlm_usage") if vlm_cap else None
            needs_feedback = bool(vlm_cap and not saved_feedback)
            ts = datetime.datetime.now().timestamp()
            vlm_sig = image_signature(img) if vlm_cap else None

            return (
                gr.update(visible=True),     # main_ui_placeholder
                gr.update(visible=False),    # user_page_placeholder
                gr.update(visible=True),     # browse_toggle_row
                gr.update(visible=True),     # browse_data_placeholder
                img,                         # image_inp
                caption,                     # long_caption_inp
                category,                    # category_btn
                concept,                     # concept_btn
                example_id,                  # exampleid_btn
                vlm_cap,                     # vlm_output
                gr.update(interactive=True), # exclude_btn
                (
                    gr.update(value=None, visible=True)
                    if needs_feedback
                    else gr.update(visible=False)
                ),                            # vlm_feedback
                # States
                img,                         # image_state
                caption,                     # caption_state
                category,                    # category_state
                concept,                     # concept_state
                example_id,                  # exampleid_state
                vlm_cap,                     # vlm_caption_state
                saved_model,                 # vlm_model_state
                saved_feedback,              # vlm_feedback_state
                saved_usage,                 # vlm_usage_state
                ts,                          # timestamp_state
                vlm_sig,                     # vlm_image_signature_state
            )

        up_data_table.click(
            fn=_on_user_page_row_click_load,
            inputs=[up_data_table, up_all_samples_raw],
            outputs=[
                main_ui_placeholder, user_page_placeholder, browse_toggle_row, browse_data_placeholder,
                image_inp, long_caption_inp, category_btn, concept_btn, exampleid_btn, vlm_output,
                exclude_btn, vlm_feedback,
                image_state, caption_state, category_state, concept_state, exampleid_state, vlm_caption_state,
                vlm_model_state, vlm_feedback_state, vlm_usage_state, timestamp_state,
                vlm_image_signature_state,
            ],
        ).then(
            fn=_copy_image_for_original_state,
            inputs=[image_inp],
            outputs=[ori_img],
        )

        def _on_hidden_concept_trigger(concept):
            return (
                gr.update(visible=True),    # main_ui_placeholder
                gr.update(visible=False),   # user_page_placeholder
                gr.update(visible=True),    # browse_toggle_row
                gr.update(visible=True),    # browse_data_placeholder
                gr.update(value=concept),   # concept_btn
            )

        up_hidden_concept_trigger.change(
            fn=_on_hidden_concept_trigger,
            inputs=[up_hidden_concept_trigger],
            outputs=[main_ui_placeholder, user_page_placeholder, browse_toggle_row, browse_data_placeholder, concept_btn],
        )

        # ── Delete from user page ──
        up_delete_btn.click(
            fn=lambda eid: gr.update(visible=True) if eid else gr.update(visible=False),
            inputs=[up_selected_example_id],
            outputs=[up_modal_delete_confirm],
        )
        up_cancel_delete_btn.click(
            fn=lambda: gr.update(visible=False),
            outputs=[up_modal_delete_confirm],
        )

        def _do_user_page_delete(example_id, country, language, email_val):
            email_val = (email_val or "").strip()
            mark_sample_excluded(example_id, country, language, email_val, LOCAL_DS_DIRECTORY_PATH)
            # Rebuild profile after deletion
            rebuild_user_profile(email_val, LOCAL_DS_DIRECTORY_PATH)

        up_confirm_delete_btn.click(
            fn=lambda: gr.update(visible=False),
            outputs=[up_modal_delete_confirm],
        ).success(
            fn=_do_user_page_delete,
            inputs=[up_selected_example_id, up_selected_country, up_selected_language, username_inp],
        ).success(
            fn=_load_user_page_data,
            inputs=[username_inp, local_storage],
            outputs=[
                up_data_table,
                up_country_filter, up_language_filter, up_category_filter, up_concept_filter,
                up_all_samples_raw,
                comparative_stats_html, user_centric_stats_html, missing_concepts_html,
                user_page_title,
            ],
        ).success(
            fn=lambda: gr.update(interactive=False),
            outputs=[up_delete_btn],
        )

        # ── VLM feedback from user page ──
        def _save_vlm_feedback_from_user_page(example_id, feedback_val, email_val):
            email_val = (email_val or "").strip()
            if example_id and feedback_val and email_val:
                update_vlm_feedback_in_json(example_id, feedback_val, email_val, LOCAL_DS_DIRECTORY_PATH)
                rebuild_user_profile(email_val, LOCAL_DS_DIRECTORY_PATH)
                gr.Info("✅ VLM feedback saved!", duration=3)

        up_vlm_feedback_radio.change(
            fn=_save_vlm_feedback_from_user_page,
            inputs=[up_selected_example_id, up_vlm_feedback_radio, username_inp],
        )

    return ui