visualization/app.py

"""
AI Messaging System - Visualization Tool
Main entry point with authentication and brand selection.
"""

import streamlit as st
import sys
from pathlib import Path
from dotenv import load_dotenv

# Load environment variables from .env file FIRST
env_path = Path(__file__).parent / '.env'
if env_path.exists():
    load_dotenv(env_path)
else:
    # Try parent directory .env as fallback
    parent_env_path = Path(__file__).parent.parent / '.env'
    if parent_env_path.exists():
        load_dotenv(parent_env_path)

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from utils.auth import verify_login, check_authentication, get_current_user, logout
from utils.theme import apply_theme, get_brand_emoji, get_brand_theme
from utils.data_loader import DataLoader
from utils.config_manager import ConfigManager
from snowflake.snowpark import Session
import os

# Page configuration
st.set_page_config(
    page_title="AI Messaging System - Visualization",
    page_icon="🎵",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Helper function to create Snowflake session
def create_snowflake_session() -> Session:
    """Create a Snowflake session using environment variables."""
    conn_params = {
        "user": os.getenv("SNOWFLAKE_USER"),
        "password": os.getenv("SNOWFLAKE_PASSWORD"),
        "account": os.getenv("SNOWFLAKE_ACCOUNT"),
        "role": os.getenv("SNOWFLAKE_ROLE"),
        "database": os.getenv("SNOWFLAKE_DATABASE"),
        "warehouse": os.getenv("SNOWFLAKE_WAREHOUSE"),
        "schema": os.getenv("SNOWFLAKE_SCHEMA"),
    }
    return Session.builder.configs(conn_params).create()


# Initialize session state
def init_session_state():
    """Initialize session state variables."""
    defaults = {
        # Authentication
        "authenticated": False,
        "user_email": "",

        # Brand and configs
        "selected_brand": None,
        "cached_configs": {},  # Cache configs per brand: {brand: {config_name: config_data}}
        "configs_loaded": False,

        # Current experiment data (in-memory)
        "current_experiment_id": None,
        "ui_log_data": None,  # Accumulated dataframe for multi-stage messages
        "current_experiment_metadata": [],  # List of metadata dicts per stage
        "current_feedbacks": [],  # List of feedback dicts

        # AB Testing data
        "ab_testing_mode": False,
        "ui_log_data_a": None,
        "ui_log_data_b": None,
        "experiment_a_id": None,
        "experiment_b_id": None,
        "experiment_a_metadata": [],
        "experiment_b_metadata": [],
        "feedbacks_a": [],
        "feedbacks_b": [],
    }
    for key, value in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = value


init_session_state()

# Authentication check
if not check_authentication():
    # Show login page
    st.title("🔐 AI Messaging System - Login")

    st.markdown("""
    Welcome to the **AI Messaging System Visualization Tool**!

    This tool enables you to:
    - 🏗️ Build and configure message campaigns
    - 👀 Visualize generated messages across all stages
    - 📊 Analyze performance and provide feedback
    - 🧪 Run A/B tests to compare different approaches

    ---

    **Access is restricted to authorized team members only.**
    Please enter your credentials below.
    """)

    with st.form("login_form"):
        email = st.text_input("📧 Email Address", placeholder="your.name@musora.com")
        token = st.text_input("🔑 Access Token", type="password", placeholder="Enter your access token")
        submit = st.form_submit_button("🚀 Login", use_container_width=True)

    if submit:
        if verify_login(email, token):
            st.session_state.authenticated = True
            st.session_state.user_email = email
            st.success("✅ Login successful! Redirecting...")
            st.rerun()
        else:
            st.error("❌ Invalid email or access token. Please try again.")

    st.stop()

# User is authenticated - show main app
# Apply base theme initially
apply_theme("base")

# Sidebar - Brand Selection
with st.sidebar:
    st.title("🎵 AI Messaging System")

    st.markdown(f"**Logged in as:** {get_current_user()}")

    if st.button("🚪 Logout", use_container_width=True):
        logout()
        st.rerun()

    st.markdown("---")

    # Brand Selection
    st.subheader("🎨 Select Brand")

    brands = ["drumeo", "pianote", "guitareo", "singeo"]
    brand_labels = {
        "drumeo": "🥁 Drumeo",
        "pianote": "🎹 Pianote",
        "guitareo": "🎸 Guitareo",
        "singeo": "🎤 Singeo"
    }

    # Get current brand from session state if exists
    current_brand = st.session_state.get("selected_brand", brands[0])

    selected_brand = st.selectbox(
        "Brand",
        brands,
        index=brands.index(current_brand) if current_brand in brands else 0,
        format_func=lambda x: brand_labels[x],
        key="brand_selector",
        label_visibility="collapsed"
    )

    # Check if brand changed
    brand_changed = (st.session_state.selected_brand != selected_brand)

    # Update session state with selected brand
    st.session_state.selected_brand = selected_brand

    # Load configs from Snowflake if brand changed or not loaded yet
    if brand_changed or not st.session_state.get("configs_loaded", False):
        if selected_brand:
            with st.spinner(f"Loading configurations for {selected_brand}..."):
                try:
                    # Create Snowflake session for config loading
                    session = create_snowflake_session()

                    # Initialize config manager with session
                    config_manager = ConfigManager(session)

                    # Load all configs for this brand
                    configs = config_manager.get_all_configs(selected_brand)

                    # Cache in session state
                    if selected_brand not in st.session_state.cached_configs:
                        st.session_state.cached_configs[selected_brand] = {}
                    st.session_state.cached_configs[selected_brand] = configs
                    st.session_state.configs_loaded = True

                    # Close session after loading
                    session.close()

                    st.success(f"✅ Loaded {len(configs)} configurations from Snowflake")

                except Exception as e:
                    st.error(f"❌ Error loading configs from Snowflake: {e}")
                    st.info("💡 Make sure Snowflake credentials are set in .env file")

    # Apply brand theme
    if selected_brand:
        apply_theme(selected_brand)

    st.markdown("---")

    # Navigation info
    st.subheader("📖 Navigation")
    st.markdown("""
    Use the pages in the sidebar to navigate:
    - **Campaign Builder**: Create campaigns & A/B tests
    - **Message Viewer**: Review messages
    - **Analytics**: View current performance
    - **Historical Analytics**: Track improvements
    """)

# Main Content Area
if not selected_brand:
    st.title("🎵 Welcome to AI Messaging System")

    st.markdown("""
    ### Get Started

    **Please select a brand from the sidebar to begin.**

    Once you've selected a brand, you can:

    1. **Build Campaigns** - Configure and generate personalized messages
    2. **View Messages** - Browse and provide feedback on generated messages
    3. **Run A/B Tests** - Compare different configurations side-by-side
    4. **Analyze Results** - Review performance metrics and insights
    """)

    # Show feature cards
    col1, col2 = st.columns(2)

    with col1:
        st.markdown("""
        #### 🏗️ Campaign Builder
        - Configure multi-stage campaigns
        - Built-in A/B testing mode
        - Parallel experiment processing
        - Automatic experiment archiving
        - Save custom configurations
        """)

        st.markdown("""
        #### 👀 Message Viewer
        - View all generated messages
        - Side-by-side A/B comparison
        - User-centric or stage-centric views
        - Search and filter messages
        - Provide detailed feedback
        """)

    with col2:
        st.markdown("""
        #### 📊 Analytics Dashboard
        - Real-time performance metrics
        - A/B test winner determination
        - Rejection reason analysis
        - Stage-by-stage performance
        - Export capabilities
        """)

        st.markdown("""
        #### 📚 Historical Analytics
        - Track all past experiments
        - Rejection rate trends over time
        - Compare historical A/B tests
        - Export historical data
        """)

else:
    # Brand is selected - show brand-specific homepage
    emoji = get_brand_emoji(selected_brand)
    theme = get_brand_theme(selected_brand)

    st.title(f"{emoji} {selected_brand.title()} - AI Messaging System")

    st.markdown(f"""
    ### Welcome to {selected_brand.title()} Message Generation!

    You're all set to create personalized messages for {selected_brand.title()} users.
    """)

    # Quick stats
    data_loader = DataLoader()

    # Check if we have brand users
    has_users = data_loader.has_brand_users(selected_brand)

    # Check if we have generated messages
    messages_df = data_loader.load_generated_messages()
    has_messages = messages_df is not None and len(messages_df) > 0

    st.markdown("---")

    # Status cards
    col1, col2, col3 = st.columns(3)

    with col1:
        st.markdown("### 👥 Available Users")
        if has_users:
            user_count = data_loader.get_brand_user_count(selected_brand)
            st.metric("Users Available", user_count)
            st.success(f"✅ {user_count} users ready")
        else:
            st.metric("Users Available", 0)
            st.info("ℹ️ No users available")

    with col2:
        st.markdown("### 📨 Generated Messages")
        if has_messages:
            stats = data_loader.get_message_stats()
            st.metric("Total Messages", stats['total_messages'])
            st.success(f"✅ {stats['total_stages']} stages")
        else:
            st.metric("Total Messages", 0)
            st.info("ℹ️ No messages generated yet")

    with col3:
        st.markdown("### 🎯 Quick Actions")
        st.markdown("")  # spacing
        if st.button("🏗️ Build Campaign", use_container_width=True):
            st.switch_page("pages/1_Campaign_Builder.py")
        if has_messages:
            if st.button("👀 View Messages", use_container_width=True):
                st.switch_page("pages/2_Message_Viewer.py")

    st.markdown("---")

    # Getting started guide
    st.markdown("### 🚀 Quick Start Guide")

    st.markdown("""
    #### Step 1: Build a Campaign
    Navigate to **Campaign Builder** to:
    - Toggle A/B testing mode on/off as needed
    - Select number of users to experiment with (1-25)
    - Configure campaign stages (or two experiments side-by-side)
    - Set LLM models and instructions
    - Generate personalized messages with automatic archiving

    #### Step 2: Review Messages
    Go to **Message Viewer** to:
    - Browse all generated messages
    - View A/B tests side-by-side automatically
    - Switch between user-centric and stage-centric views
    - Provide detailed feedback with rejection reasons
    - Track message headers and content

    #### Step 3: Analyze Performance
    Check **Analytics Dashboard** for:
    - Real-time approval and rejection rates
    - A/B test comparisons with winner determination
    - Common rejection reasons with visual charts
    - Stage-by-stage performance insights
    - Export analytics data

    #### Step 4: Track Improvements
    View **Historical Analytics** to:
    - See all past experiments over time
    - Identify rejection rate trends
    - Compare historical A/B tests
    - Export comprehensive historical data
    """)

    st.markdown("---")

    # Tips
    with st.expander("💡 Tips & Best Practices"):
        st.markdown("""
        - **User Selection**: Select 1-25 users for quick experimentation
        - **Configuration Templates**: Start with default configurations and customize as needed
        - **A/B Testing**: Enable A/B mode in Campaign Builder to compare two configurations in parallel
        - **Feedback**: Reject poor messages with specific categories including the new 'Similar To Previous' option
        - **Historical Tracking**: Check Historical Analytics regularly to track improvements over time
        - **Stage Design**: Each stage should build upon previous stages with varied messaging approaches
        - **Automatic Archiving**: Previous experiments are automatically archived when you start a new one
        """)

st.markdown("---")
st.markdown("**Built with ❤️ for the Musora team**")