visualization/app.py

"""
AI Messaging System v3 — Visualization Tool
Entry point: Login + data pre-fetch on authentication.
"""

import streamlit as st

st.set_page_config(
    page_title="Messaging Viewer",
    page_icon="💬",
    layout="wide",
    initial_sidebar_state="expanded",
)

from utils.auth import check_authentication, verify_login, logout, get_current_user
from utils.theme import inject_global_css
from utils import snowflake_client as sf
from utils import feedback_manager as fm

inject_global_css()


# ---------------------------------------------------------------------------
# Login gate
# ---------------------------------------------------------------------------

def show_login():
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.markdown("## 💬 Messaging Viewer")
        st.markdown("Sign in to review AI-generated push notification messages.")
        st.markdown("---")

        with st.form("login_form"):
            email = st.text_input("Email", placeholder="you@musora.com")
            token = st.text_input("Access Token", type="password")
            submitted = st.form_submit_button("Sign In", use_container_width=True)

        if submitted:
            if verify_login(email, token):
                st.session_state["authenticated"] = True
                st.session_state["user_email"] = email.lower().strip()
                st.rerun()
            else:
                st.error("Invalid email or token. Please try again.")


# ---------------------------------------------------------------------------
# Data pre-fetch (runs once after login, cached in session_state)
# ---------------------------------------------------------------------------

def prefetch_messages():
    """
    Load top 150 messages per campaign for all campaigns in a single query.
    Stores the result in session_state["all_messages"] so every page can
    filter client-side without additional Snowflake round-trips.
    Skips fetch if data is already cached in this session.
    """
    if "all_messages" in st.session_state:
        return  # already loaded

    with st.spinner("Loading messages…"):
        try:
            sf.ensure_feedback_table()
            df = sf.fetch_all_messages()
            st.session_state["all_messages"] = df
        except Exception as e:
            st.error(f"Failed to load messages: {e}")


# ---------------------------------------------------------------------------
# Main app (authenticated)
# ---------------------------------------------------------------------------

def show_home():
    user = get_current_user()

    with st.sidebar:
        st.markdown(f"**Signed in as:** {user}")
        if st.button("Sign Out", use_container_width=True):
            logout()
            fm.reset_session()
            st.session_state.pop("all_messages", None)
            st.rerun()
        if st.button("Reload Data", use_container_width=True, help="Re-fetch messages from Snowflake"):
            st.session_state.pop("all_messages", None)
            fm.reset_session()
            st.rerun()

    # Pre-fetch data as soon as the user lands on the home page
    prefetch_messages()

    df = st.session_state.get("all_messages")
    total = len(df) if df is not None else 0

    st.markdown("## 💬 Messaging Viewer")
    st.markdown(
        "Browse AI-generated push notifications, inspect user context, "
        "and flag messages that need improvement."
    )
    st.markdown("---")

    if df is not None and not df.empty:
        campaigns = df["CAMPAIGN_NAME"].value_counts().to_dict()
        cols = st.columns(len(campaigns) + 1)
        cols[0].metric("Total loaded", total)
        for i, (camp, count) in enumerate(campaigns.items(), 1):
            cols[i].metric(camp.replace("dailyPush_", ""), count)
        st.markdown("---")

    st.markdown("### How to use")
    c1, c2, c3 = st.columns(3)
    with c1:
        st.markdown("**1. Open Message Viewer**")
        st.markdown("Navigate to the Message Viewer page using the sidebar.")
    with c2:
        st.markdown("**2. Filter instantly**")
        st.markdown(
            "Use the sidebar filters (brand, campaign, scenario) to narrow down "
            "messages. Changes apply immediately — no Apply button needed."
        )
    with c3:
        st.markdown("**3. Review & reject**")
        st.markdown(
            "Read each message with its user context. Click **Reject** and pick "
            "a reason for any message that has issues. No action = message is good."
        )

    st.markdown("---")
    st.page_link("pages/1_Message_Viewer.py", label="Open Message Viewer →", icon="💬")


# ---------------------------------------------------------------------------
# Router
# ---------------------------------------------------------------------------

if not check_authentication():
    show_login()
else:
    show_home()