pages/Search_Personalization.py

"""
Search Personalization — OpenSearch UBI Boutique
A personalized product search experience powered by OpenSearch agentic memory.

This page demonstrates how user behavior (clickstream) data stored in OpenSearch
can drive real-time search personalization through an agentic memory pipeline.
"""
import streamlit as st
from typing import Dict, List, Any, Optional
from search_personalization.data_loader import get_opensearch_client
from search_personalization.image_utils import get_image_path, get_thumbnail_path, preload_image_cache
from opensearchpy import OpenSearchException
import os
from dotenv import load_dotenv
import uuid
from datetime import datetime, timedelta

# Import UBI tracking functions
from search_personalization.ubi_tracker import (
    track_query, track_event, is_ubi_enabled, validate_ubi_config
)

# Import agentic memory pipeline
from search_personalization.agentic_memory.demo_runner import run_pipeline_for_persona

# Import user preferences
from search_personalization.user_preferences import get_user_preferences

# Load environment variables
load_dotenv()

# UBI Configuration
UBI_CONFIG = {
    'events_endpoint': os.getenv('UBI_EVENTS_ENDPOINT'),
    'queries_endpoint': os.getenv('UBI_QUERIES_ENDPOINT'),
    'aws_region': os.getenv('AWS_REGION', 'us-east-1'),
    'debug_mode': os.getenv('UBI_DEBUG_MODE', 'false').lower() == 'true',
    'enabled': os.getenv('UBI_ENABLED', 'true').lower() == 'true',
    'application_name': os.getenv('UBI_APPLICATION_NAME', 'opensearch-ubi')
}

# User Personas for clickstream tracking
USER_PERSONAS = [
    {
        "user_id": "anonymous",
        "name": "Anonymous User",
        "age": 30,
        "gender": "unspecified",
        "location": "General",
        "occupation": "General User",
        "interests": ["General Shopping"],
        "shopping_behavior": "Standard search without personalization",
        "key": "anonymous",
        "emoji": "👤"
    },
    {
        "name": "Sarah",
        "user_id": "user1",
        "age": 32,
        "gender": "female",
        "location": "San Francisco, CA",
        "occupation": "Marketing Manager",
        "interests": ["Technology", "Fitness", "Travel"],
        "shopping_behavior": "Focused on quality and trends.",
        "key": "business_professional",
        "emoji": "👩‍💼"
    },
    {
        "name": "Alex",
        "user_id": "user2",
        "age": 21,
        "gender": "male",
        "location": "Austin, TX",
        "occupation": "Computer Science Student",
        "interests": ["Gaming", "Programming", "Music"],
        "shopping_behavior": "Budget-conscious, looks for deals and discounts",
        "key": "student",
        "emoji": "🎓"
    }
]

# Validate UBI configuration on startup
_ubi_valid, _ubi_error = validate_ubi_config()
if not _ubi_valid:
    print(f"WARNING: {_ubi_error}")

SESSION_DURATION = timedelta(minutes=1)

# Base path for product images (relative to the search_personalization module)
_MODULE_DIR = os.path.dirname(os.path.abspath(__file__))
_PRODUCT_SAMPLES_DIR = os.path.join(
    os.path.dirname(_MODULE_DIR), "search_personalization", "product_samples"
)


def _ensure_user_session(user_id: str) -> str:
    """Ensure the given user has a valid (non-expired) session."""
    now = datetime.utcnow()
    sessions = st.session_state.get('sp_user_sessions', {})
    user_session = sessions.get(user_id)

    if user_session is None or (now - user_session['started_at']) >= SESSION_DURATION:
        sessions[user_id] = {
            'session_id': str(uuid.uuid4()),
            'started_at': now,
        }
        st.session_state.sp_user_sessions = sessions

    return sessions[user_id]['session_id']


def _init_session_state() -> None:
    """Initialize session state variables prefixed with sp_ to avoid collision."""
    defaults = {
        'sp_cart': [],
        'sp_products': [],
        'sp_filtered_products': [],
        'sp_search_query': '',
        'sp_query_id': 'browsing',
        'sp_current_page': 1,
        'sp_selected_product': None,
        'sp_page_view': 'categories',
        'sp_selected_category': None,
        'sp_products_loaded': False,
        'sp_load_error': None,
        'sp_show_opensearch_query': False,
        'sp_show_agent_reasoning': False,
        'sp_last_opensearch_query': None,
        'sp_selected_user_id': 'anonymous',
        'sp_pending_query_tracking': False,
        'sp_last_search_product_ids': [],
        'sp_ubi_client_id': str(uuid.uuid4()),
        'sp_user_sessions': {},
        'sp_user_preferences': None,
        'sp_preferences_applied': False,
    }
    for key, default in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = default

    # Ensure current user has an active session
    _ensure_user_session(st.session_state.sp_selected_user_id)
    st.session_state.sp_ubi_session_id = (
        st.session_state.sp_user_sessions[st.session_state.sp_selected_user_id]['session_id']
    )

    # Load user preferences on first run
    if (st.session_state.sp_user_preferences is None
            and st.session_state.sp_selected_user_id != 'anonymous'):
        _load_and_apply_preferences(st.session_state.sp_selected_user_id)


def _load_and_apply_preferences(user_id: str) -> None:
    """Load preferences from OpenSearch for the user."""
    if user_id == 'anonymous':
        st.session_state.sp_user_preferences = None
        st.session_state.sp_preferences_applied = False
        return
    prefs = get_user_preferences(user_id)
    st.session_state.sp_user_preferences = prefs
    st.session_state.sp_preferences_applied = False


def _load_products() -> None:
    """Load all products from OpenSearch."""
    if st.session_state.sp_products_loaded:
        return
    try:
        client = get_opensearch_client()
        response = client.search(
            index='products',
            body={
                "query": {"match_all": {}},
                "size": 10000,
                "_source": {"excludes": ["image_binary", "product_description_vector"]}
            }
        )
        products = [hit['_source'] for hit in response['hits']['hits']]
        st.session_state.sp_products = products
        st.session_state.sp_filtered_products = products
        st.session_state.sp_products_loaded = True
        st.session_state.sp_load_error = None
        if products:
            preload_image_cache(products)
    except Exception as e:
        st.session_state.sp_load_error = str(e)
        st.session_state.sp_products_loaded = False
        st.session_state.sp_products = []
        st.session_state.sp_filtered_products = []


def _get_product_by_id(product_id: str) -> Optional[Dict[str, Any]]:
    """Get product by ID from session state."""
    for product in st.session_state.sp_products:
        if product.get('id') == product_id:
            return product
    return None


def _add_to_cart(product_id: str) -> None:
    """Add a product to the shopping cart and track with UBI."""
    st.session_state.sp_cart.append(product_id)
    if is_ubi_enabled():
        try:
            product = _get_product_by_id(product_id)
            if product:
                query_id = st.session_state.get('sp_query_id', 'browsing')
                event_attributes = {
                    'object': {
                        'object_id': product_id,
                        'object_id_field': 'id',
                        'description': product.get('description', '')
                    },
                    'position': {'ordinal': 0},
                    'product_name': product.get('name', ''),
                    'product_price': product.get('price', 0),
                    'product_category': product.get('category', '')
                }
                track_event(
                    action_name='add_to_cart',
                    query_id=query_id,
                    object_id=product_id,
                    event_attributes=event_attributes,
                    message=f"Added {product.get('name', 'product')} to cart",
                    message_type='CONVERSION',
                    user_id=st.session_state.sp_selected_user_id
                )
        except Exception:
            pass


def _get_cart_items() -> List[Dict[str, Any]]:
    """Get cart items with quantities."""
    from collections import Counter
    product_counts = Counter(st.session_state.sp_cart)
    cart_items = []
    for product_id, quantity in product_counts.items():
        product = _get_product_by_id(product_id)
        if product:
            cart_items.append({
                'product': product,
                'quantity': quantity,
                'product_id': product_id
            })
    return cart_items


def _calculate_cart_total() -> float:
    """Calculate the total price of all items in the cart."""
    cart_items = _get_cart_items()
    return sum(item['product'].get('price', 0) * item['quantity'] for item in cart_items)


# --- Mapping from user_id to persona name for agentic memory ---
_PERSONA_NAME_MAP = {"user1": "sarah", "user2": "alex"}


def _search_products(search_query: str) -> tuple:
    """Search products in OpenSearch using neural knn search."""
    try:
        client = get_opensearch_client()
        if not search_query or search_query.strip() == '':
            query_body = {"query": {"match_all": {}}, "size": 10}
        else:
            query_body = {
                "query": {
                    "neural": {
                        "product_description_vector": {
                            "query_text": search_query.strip(),
                            "model_id": os.getenv('OPENSEARCH_MODEL_ID', ''),
                            "k": 50
                        }
                    }
                },
                "size": 10
            }
        import json
        st.session_state.sp_last_opensearch_query = json.dumps(query_body, indent=2)
        response = client.search(index='products', body=query_body)
        products = []
        product_ids = []
        for hit in response['hits']['hits']:
            product = hit['_source']
            products.append(product)
            product_ids.append(product.get('id', hit.get('_id')))
        total_count = response['hits']['total']['value']
        return products, total_count, product_ids
    except Exception as e:
        st.error(f"Search error: {str(e)}")
        return [], 0, []


def _agentic_search(search_query: str, user_id: str) -> tuple:
    """Run the agentic memory pipeline for personalized search."""
    import json as _json

    persona_name = _PERSONA_NAME_MAP[user_id]
    session_id = st.session_state.get("sp_ubi_session_id", "web-session")

    trace = run_pipeline_for_persona(
        query=search_query,
        persona_name=persona_name,
        session_id=session_id,
    )

    # Extract explanation
    _explanation_data = trace.get("agents", {}).get("explanation", {})
    explanation = _explanation_data.get("output")
    if explanation is None and "_future" in _explanation_data:
        _future = _explanation_data["_future"]
        try:
            _future.result(timeout=3.0)
            explanation = _future.result() if _future.done() else None
        except Exception:
            explanation = None

    # Extract enriched query and inferred attributes
    query_agent_output = trace["agents"].get("query", {}).get("output", "")
    enriched_query_text = search_query
    inferred_attrs = {}
    try:
        _qstart = query_agent_output.find("{")
        _qend = query_agent_output.rfind("}") + 1
        if _qstart >= 0 and _qend > _qstart:
            _parsed_query = _json.loads(query_agent_output[_qstart:_qend])
            enriched_query_text = _parsed_query.get("enriched_query", search_query)
            inferred_attrs = _parsed_query.get("inferred_attributes", {})
    except (_json.JSONDecodeError, TypeError):
        pass

    # Store the OpenSearch query for display
    actual_os_query = trace.get("opensearch_query")
    if actual_os_query:
        st.session_state.sp_last_opensearch_query = _json.dumps(actual_os_query, indent=2)

    # Parse ranked product IDs from ranking agent
    ranking_output = trace["agents"].get("ranking", {}).get("output", "")
    ranked_product_ids = []
    personalization_summary = None
    try:
        start = ranking_output.find("{")
        end = ranking_output.rfind("}") + 1
        if start >= 0 and end > start:
            parsed = _json.loads(ranking_output[start:end])
            for item in parsed.get("results", []):
                pid = item.get("product_id") or item.get("id")
                if pid:
                    ranked_product_ids.append(pid)
            personalization_summary = parsed.get("personalization_summary")
    except (_json.JSONDecodeError, TypeError):
        pass

    st.session_state._sp_personalization = personalization_summary

    # Store reasoning trace
    _reranked_count = len(ranked_product_ids)
    st.session_state._sp_reasoning = {
        "query_enrichment": {
            "original_query": search_query,
            "enriched_query": enriched_query_text,
            "inferred_attributes": inferred_attrs,
        },
        "reranking": {
            "model": "Cohere Rerank 3.5 (via Bedrock)",
            "results_reranked": _reranked_count,
        },
        "aversion_filter": {
            "aversions_from_profile": inferred_attrs.get("aversions", []),
        },
        "personalization_summary": personalization_summary,
        "pipeline_duration_ms": trace.get("total_duration_ms"),
    }

    # Resolve product IDs to full product dicts
    products = []
    product_ids = []
    for pid in ranked_product_ids:
        product = _get_product_by_id(pid)
        if product:
            products.append(product)
            product_ids.append(pid)

    # Fallback to regular search if pipeline returned nothing
    if not products:
        products, _, product_ids = _search_products(enriched_query_text)
        return products, len(products), product_ids, explanation

    return products, len(products), product_ids, explanation


# ─── UI COMPONENTS ───────────────────────────────────────────────────────────


def _render_header() -> None:
    """Render the header with title on top, then search bar + persona + cart below."""
    # Home button
    st.page_link("app.py", label=":orange[Home]", icon="🏠")
    # Title on its own row
    st.markdown(
        '<h1 style="font-size:1.8rem;font-weight:bold;color:#ffffff;margin:0 0 0.5rem 0;">'
        'Search personalization with agentic memory and UBI data</h1>',
        unsafe_allow_html=True
    )

    # Search bar, persona selector, cart on second row
    header_col2, header_col3, header_col4 = st.columns([4, 1.5, 0.5])

    with header_col2:
        search_query = st.text_input(
            "Search", placeholder="What are you looking for?",
            label_visibility="collapsed", key="sp_header_search"
        )
        if search_query != st.session_state.sp_search_query:
            st.session_state.sp_search_query = search_query
            if search_query and search_query.strip():
                st.session_state.sp_page_view = 'grid'
                st.session_state.sp_selected_category = None
            else:
                st.session_state.sp_page_view = 'categories'
                st.session_state.sp_selected_category = None
            new_query_id = str(uuid.uuid4())
            st.session_state.sp_query_id = new_query_id
            st.session_state.sp_pending_query_tracking = True
            if is_ubi_enabled():
                try:
                    query_text = search_query.strip() if search_query and search_query.strip() else "match_all"
                    track_event(
                        action_name="search", query_id=new_query_id,
                        event_attributes={"query_text": query_text},
                        message=f"User searched for: {query_text}",
                        user_id=st.session_state.sp_selected_user_id
                    )
                except Exception:
                    pass

    with header_col3:
        persona_options = [f"{p['emoji']} {p['name']}" for p in USER_PERSONAS]
        current_idx = next(
            (i for i, p in enumerate(USER_PERSONAS)
             if p['user_id'] == st.session_state.sp_selected_user_id), 0
        )
        selected_persona = st.selectbox(
            "User", options=persona_options, index=current_idx,
            label_visibility="collapsed", key="sp_persona_selector",
            help="Select a user persona for clickstream tracking"
        )
        selected_index = persona_options.index(selected_persona)
        new_user_id = USER_PERSONAS[selected_index]['user_id']
        if new_user_id != st.session_state.sp_selected_user_id:
            st.session_state.sp_selected_user_id = new_user_id
            st.session_state.sp_ubi_session_id = _ensure_user_session(new_user_id)
            st.session_state.sp_query_id = str(uuid.uuid4())
            # Clear agentic pipeline cache
            st.session_state.pop("_sp_cache_key", None)
            st.session_state.pop("_sp_cache", None)
            st.session_state.pop("_sp_reasoning", None)
            st.session_state.sp_last_opensearch_query = None
            _load_and_apply_preferences(new_user_id)

    with header_col4:
        cart_count = len(st.session_state.sp_cart)
        if st.button(f"🛒 {cart_count}", key="sp_cart_button", help="View Cart"):
            st.session_state.sp_page_view = 'cart'
            st.rerun()


def _render_product_card(product: Dict[str, Any], position: int = 0) -> None:
    """Render a single product card with actions."""
    product_id = product.get('id', '')
    product_name = product.get('name', 'Unknown Product')
    price = product.get('price', 0)
    current_stock = product.get('current_stock', 0)
    category = product.get('category', 'Unknown')
    image_path = get_thumbnail_path(product)
    query_id = st.session_state.get('sp_query_id', 'browsing')
    user_id = st.session_state.get('sp_selected_user_id', 'anonymous')

    col_image, col_content = st.columns([1, 4])

    with col_image:
        try:
            st.image(image_path, use_container_width=True)
        except Exception:
            st.image(
                os.path.join(_PRODUCT_SAMPLES_DIR, "product_image_coming_soon.png"),
                use_container_width=True
            )

    with col_content:
        st.markdown(f"**{product_name}**")
        price_col, stock_col = st.columns([1, 2])
        with price_col:
            st.markdown(f"${price:,.2f}")
        with stock_col:
            if current_stock < 10:
                st.markdown(f"⚠️ {current_stock} items left")
            else:
                st.caption(f"{current_stock} items left")

        btn1, btn2, btn3, btn4 = st.columns([1, 1, 1, 2])
        with btn1:
            if st.button("🛒 Add", key=f"sp_cart_{product_id}", type="primary"):
                _add_to_cart(product_id)
                st.rerun()
        with btn2:
            if st.button("❤️", key=f"sp_like_{product_id}"):
                if is_ubi_enabled():
                    try:
                        track_event(
                            action_name='user_feedback', query_id=query_id,
                            object_id=product_id,
                            event_attributes={
                                'object': {'object_id': product_id, 'object_id_field': 'id'},
                                'position': {'ordinal': position},
                                'feedback': 'thumbs_up'
                            },
                            message=f"Thumbs up: {product_name}",
                            message_type='LIKE', user_id=user_id
                        )
                    except Exception:
                        pass
                st.toast(f"You liked {product_name}!")

        with btn3:
            if st.button("👎", key=f"sp_dislike_{product_id}"):
                if is_ubi_enabled():
                    try:
                        track_event(
                            action_name='user_feedback', query_id=query_id,
                            object_id=product_id,
                            event_attributes={
                                'object': {'object_id': product_id, 'object_id_field': 'id'},
                                'position': {'ordinal': position},
                                'feedback': 'thumbs_down'
                            },
                            message=f"Thumbs down: {product_name}",
                            message_type='REJECT', user_id=user_id
                        )
                    except Exception:
                        pass
                st.toast(f"You disliked {product_name}!")
        with btn4:
            if st.button("View Details", key=f"sp_details_{product_id}", type="secondary"):
                if is_ubi_enabled():
                    try:
                        track_event(
                            action_name='view_product_details', query_id=query_id,
                            object_id=product_id,
                            event_attributes={
                                'object': {'object_id': product_id, 'object_id_field': 'id'},
                                'position': {'ordinal': position},
                            },
                            message=f"Viewed: {product_name}",
                            user_id=user_id
                        )
                    except Exception:
                        pass
                st.session_state.sp_selected_product = product_id
                st.session_state.sp_page_view = 'details'
                st.rerun()

    st.divider()


def _render_categories() -> None:
    """Render the category landing page."""
    categories = sorted(set(
        p.get('category', '').lower()
        for p in st.session_state.sp_products if p.get('category')
    ))
    if not categories:
        st.info("No categories available.")
        return

    category_meta = {
        'accessories': {'emoji': '👜', 'desc': 'Bags, watches, belts and more'},
        'apparel': {'emoji': '👕', 'desc': 'Shirts, jackets, pants and dresses'},
        'electronics': {'emoji': '💻', 'desc': 'Gadgets, devices and tech gear'},
        'footwear': {'emoji': '👟', 'desc': 'Shoes, boots, sneakers and sandals'},
        'jewelry': {'emoji': '💍', 'desc': 'Rings, necklaces, bracelets and earrings'},
    }

    st.markdown("### Shop by Category")
    cols_per_row = 3
    for row_start in range(0, len(categories), cols_per_row):
        row_cats = categories[row_start:row_start + cols_per_row]
        cols = st.columns(cols_per_row)
        for idx, cat in enumerate(row_cats):
            meta = category_meta.get(cat, {'emoji': '📦', 'desc': ''})
            with cols[idx]:
                # Show sample thumbnail
                thumb_dir = os.path.join(_PRODUCT_SAMPLES_DIR, ".thumbnails", cat)
                sample_image = None
                if os.path.isdir(thumb_dir):
                    images = [f for f in os.listdir(thumb_dir) if f.endswith('.jpg')]
                    if images:
                        sample_image = os.path.join(thumb_dir, images[0])
                if sample_image:
                    st.image(sample_image, use_container_width=True)
                else:
                    st.markdown(f"<div style='font-size:4rem;text-align:center;'>{meta['emoji']}</div>",
                                unsafe_allow_html=True)
                st.markdown(f"**{meta['emoji']} {cat.title()}**")
                st.caption(meta['desc'])
                if st.button(f"Browse {cat.title()}", key=f"sp_cat_{cat}", use_container_width=True):
                    st.session_state.sp_selected_category = cat
                    st.session_state.sp_page_view = 'grid'
                    st.rerun()


def _render_product_grid(products: List[Dict[str, Any]]) -> None:
    """Render the product list."""
    if not products:
        st.markdown("### 🔍 No results found")
        st.caption("Try adjusting your search query to see more products.")
        return
    st.markdown(f"**Showing {len(products)} products**")
    for i, product in enumerate(products):
        _render_product_card(product, position=i + 1)


def _render_product_details(product_id: str) -> None:
    """Render the product details page."""
    product = _get_product_by_id(product_id)
    if not product:
        st.error("Product not found")
        if st.button("← Back to Products", type="primary"):
            st.session_state.sp_page_view = 'grid'
            st.session_state.sp_selected_product = None
            st.rerun()
        return

    if st.button("← Back to Products", key="sp_back_details", type="secondary"):
        st.session_state.sp_page_view = 'grid'
        st.session_state.sp_selected_product = None
        st.rerun()

    st.divider()
    product_name = product.get('name', 'Unknown')
    price = product.get('price', 0)
    description = product.get('description', '')
    category = product.get('category', '')
    style = product.get('style', '')
    current_stock = product.get('current_stock', 0)
    image_path = get_thumbnail_path(product)

    col_img, col_info = st.columns([1, 1], gap="large")
    with col_img:
        try:
            st.image(image_path, use_container_width=True)
        except Exception:
            st.image(os.path.join(_PRODUCT_SAMPLES_DIR, "product_image_coming_soon.png"),
                     use_container_width=True)
    with col_info:
        st.markdown(f"## {product_name}")
        st.markdown(description)
        st.divider()
        st.markdown(f"### ${price:,.2f}")
        if current_stock > 10:
            st.success(f"✓ In Stock ({current_stock} available)")
        elif current_stock > 0:
            st.warning(f"⚠️ Low Stock ({current_stock} left)")
        else:
            st.error("✗ Out of Stock")
        if current_stock > 0:
            if st.button("🛒 Add to Cart", key="sp_detail_add", type="primary", use_container_width=True):
                _add_to_cart(product_id)
                st.success(f"✓ Added to cart!")
                st.rerun()
        st.divider()
        c1, c2, c3 = st.columns(3)
        c1.metric("Category", category.title())
        c2.metric("Style", style.title())
        gender = product.get('gender_affinity', '')
        c3.metric("Gender", "Women" if gender == "F" else "Men" if gender == "M" else "Unisex")


def _render_cart() -> None:
    """Render the shopping cart view."""
    if st.button("← Continue Shopping", key="sp_back_shop", type="secondary"):
        st.session_state.sp_page_view = 'categories'
        st.session_state.sp_selected_category = None
        st.rerun()

    st.divider()
    cart_count = len(st.session_state.sp_cart)
    st.markdown(f"### 🛒 Shopping Cart ({cart_count} items)")

    if cart_count == 0:
        st.info("Your cart is empty. Add some products to get started!")
        if st.button("Browse Products", type="primary", use_container_width=True):
            st.session_state.sp_page_view = 'categories'
            st.rerun()
        return

    cart_items = _get_cart_items()
    for idx, item in enumerate(cart_items):
        product = item['product']
        quantity = item['quantity']
        product_id = item['product_id']
        image_path = get_thumbnail_path(product)

        col_img, col_det, col_act = st.columns([1, 3, 1])
        with col_img:
            try:
                st.image(image_path, use_container_width=True)
            except Exception:
                st.image(os.path.join(_PRODUCT_SAMPLES_DIR, "product_image_coming_soon.png"),
                         use_container_width=True)
        with col_det:
            st.markdown(f"**{product.get('name', 'Unknown')}**")
            st.caption(f"Category: {product.get('category', '').title()}")
            st.markdown(f"${product.get('price', 0):,.2f} × {quantity} = "
                        f"**${product.get('price', 0) * quantity:,.2f}**")
        with col_act:
            if st.button("➖", key=f"sp_rm1_{idx}"):
                if product_id in st.session_state.sp_cart:
                    st.session_state.sp_cart.remove(product_id)
                st.rerun()
            if st.button("🗑️", key=f"sp_rmall_{idx}"):
                st.session_state.sp_cart = [
                    pid for pid in st.session_state.sp_cart if pid != product_id
                ]
                st.rerun()
        st.divider()

    # Cart summary
    total = _calculate_cart_total()
    st.markdown(f"### Total: ${total:,.2f}")

    col_clear, col_checkout = st.columns(2)
    with col_clear:
        if st.button("🗑️ Clear Cart", key="sp_clear_cart"):
            st.session_state.sp_cart = []
            st.rerun()
    with col_checkout:
        if st.button("💳 Checkout", key="sp_checkout", type="primary"):
            if is_ubi_enabled():
                try:
                    track_event(
                        action_name='checkout',
                        query_id=st.session_state.get('sp_query_id', 'browsing'),
                        event_attributes={
                            'cart_total': total,
                            'cart_item_count': cart_count,
                        },
                        message=f"Checkout: {cart_count} items (${total:.2f})",
                        message_type='CONVERSION',
                        user_id=st.session_state.sp_selected_user_id
                    )
                except Exception:
                    pass
            st.session_state.sp_cart = []
            st.success("✓ Thank you for your purchase!")
            st.rerun()


def _render_developer_sidebar() -> None:
    """Render developer tools in the sidebar."""
    with st.sidebar:
        st.subheader("🔧 Developer Tools")
        show_query = st.checkbox("Show OpenSearch Query", key="sp_show_query")
        if show_query and st.session_state.sp_last_opensearch_query:
            st.markdown("**Current Query:**")
            st.code(st.session_state.sp_last_opensearch_query, language="json")
        elif show_query:
            st.info("No query executed yet.")

        show_reasoning = st.checkbox("Show Agent Reasoning", key="sp_show_reasoning")
        if show_reasoning:
            reasoning = st.session_state.get("_sp_reasoning")
            if reasoning:
                st.markdown("---")
                st.markdown("#### 🧠 Agent Reasoning")
                enrichment = reasoning.get("query_enrichment", {})
                with st.expander("1. Query Enrichment", expanded=True):
                    st.markdown(f"**Original:** `{enrichment.get('original_query', '')}`")
                    st.markdown(f"**Enriched:** `{enrichment.get('enriched_query', '')}`")
                    attrs = enrichment.get("inferred_attributes", {})
                    if attrs:
                        for k, v in attrs.items():
                            if v:
                                st.markdown(f"- {k}: `{v}`")
                reranking = reasoning.get("reranking", {})
                with st.expander("2. Reranking", expanded=True):
                    st.markdown(f"**Model:** {reranking.get('model', 'N/A')}")
                    st.markdown(f"**Results reranked:** {reranking.get('results_reranked', 0)}")
                aversion = reasoning.get("aversion_filter", {})
                with st.expander("3. Aversions", expanded=True):
                    aversions = aversion.get("aversions_from_profile", [])
                    if aversions:
                        st.markdown(f"**Profile aversions:** {', '.join(f'`{a}`' for a in aversions)}")
                    else:
                        st.markdown("No aversions detected.")
                duration = reasoning.get("pipeline_duration_ms")
                if duration:
                    st.caption(f"⏱️ Pipeline: {duration}ms")
            else:
                st.info("Select Sarah or Alex and search to see reasoning.")


# ─── MAIN PAGE ────────────────────────────────────────────────────────────────


def main() -> None:
    """Main entry point for the Search Personalization page."""
    _init_session_state()

    # Refresh the current user's session
    st.session_state.sp_ubi_session_id = _ensure_user_session(
        st.session_state.sp_selected_user_id
    )

    # Render header (search + persona + cart)
    _render_header()
    st.divider()

    # Route to the correct view
    if st.session_state.sp_page_view == 'details' and st.session_state.sp_selected_product:
        _render_product_details(st.session_state.sp_selected_product)
        _render_developer_sidebar()
        return

    if st.session_state.sp_page_view == 'cart':
        _render_cart()
        return

    # Load products
    if not st.session_state.sp_products_loaded and not st.session_state.sp_load_error:
        with st.spinner("🔄 Loading products..."):
            _load_products()
    else:
        _load_products()

    # Categories landing
    if st.session_state.sp_page_view == 'categories' and not st.session_state.sp_search_query:
        if st.session_state.sp_products_loaded and st.session_state.sp_products:
            _render_categories()
            _render_developer_sidebar()
            return

    # Error states
    if st.session_state.sp_load_error:
        st.error("❌ Unable to connect to the product catalog.")
        st.info("Check your AWS credentials and OpenSearch endpoint.")
        if st.button("🔄 Retry", key="sp_retry", type="primary"):
            st.session_state.sp_products_loaded = False
            st.session_state.sp_load_error = None
            st.rerun()
        return

    if not st.session_state.sp_products_loaded:
        st.info("⏳ Loading products...")
        return

    if not st.session_state.sp_products:
        st.warning("No products available in the catalog.")
        return

    # ─── Search + product grid ───
    search_query = st.session_state.sp_search_query
    page_size = 10

    # Back to categories button when browsing a category
    if st.session_state.sp_selected_category and not (search_query and search_query.strip()):
        if st.button("← Back to Categories", key="sp_back_cats", type="secondary"):
            st.session_state.sp_selected_category = None
            st.session_state.sp_page_view = 'categories'
            st.rerun()
        st.markdown(f"## {st.session_state.sp_selected_category.title()}")

    # Perform search
    if search_query and search_query.strip():
        persona_user_id = st.session_state.sp_selected_user_id
        if persona_user_id in _PERSONA_NAME_MAP:
            cache_key = f"{persona_user_id}:{search_query}"
            if st.session_state.get("_sp_cache_key") != cache_key:
                products_to_display, total_count, product_ids, explanation = (
                    _agentic_search(search_query, persona_user_id)
                )
                st.session_state._sp_cache_key = cache_key
                st.session_state._sp_cache = (products_to_display, total_count, product_ids, explanation)
            else:
                products_to_display, total_count, product_ids, explanation = st.session_state._sp_cache
        else:
            products_to_display, total_count, product_ids = _search_products(search_query)
            explanation = None
    else:
        all_products = st.session_state.sp_products
        if st.session_state.sp_selected_category:
            all_products = [
                p for p in all_products
                if p.get('category', '').lower() == st.session_state.sp_selected_category
            ]
        total_count = len(all_products)
        start = (st.session_state.sp_current_page - 1) * page_size
        products_to_display = all_products[start:start + page_size]
        product_ids = [p.get('id', '') for p in products_to_display]
        explanation = None

    # UBI query tracking (non-blocking)
    if is_ubi_enabled() and st.session_state.get('sp_pending_query_tracking', False):
        st.session_state.sp_pending_query_tracking = False
        import threading
        _q_text = search_query.strip() if search_query and search_query.strip() else "match_all"
        _q_id = st.session_state.get('sp_query_id', 'browsing')
        _q_uid = st.session_state.sp_selected_user_id
        _q_hits = list(product_ids)
        def _track():
            try:
                track_query(_q_text, _q_id, user_id=_q_uid, query_response_hit_ids=_q_hits)
            except Exception:
                pass
        threading.Thread(target=_track, daemon=True).start()

    # Show search info
    if search_query and search_query.strip():
        if not products_to_display:
            st.info(f"🔍 No results found for '{search_query}'")
        else:
            if explanation:
                st.info(f"💡 {explanation}")
            personalization = st.session_state.get("_sp_personalization")
            persona_user_id = st.session_state.sp_selected_user_id
            if persona_user_id in _PERSONA_NAME_MAP and personalization and "none" not in personalization.lower():
                persona_name = _PERSONA_NAME_MAP.get(persona_user_id, "").title()
                st.caption(f"📊 Personalized results for {persona_name}")

    # Render product grid
    _render_product_grid(products_to_display)

    # Developer tools sidebar
    _render_developer_sidebar()


# Run the page
main()