src/streamlit_app.py

import streamlit as st
import os
import json
from typing import List, Dict, Optional
from dotenv import load_dotenv
from groq import Groq
import spotipy
from spotipy.oauth2 import SpotifyOAuth
import time

# Load environment variables
load_dotenv()

# Page configuration
st.set_page_config(
    page_title="VibeSync - AI Playlist Generator",
    page_icon="🎵",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Custom CSS for premium styling
st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;600;700&family=Inter:wght@300;400;500;600&family=Space+Grotesk:wght@300;400;500;600;700&display=swap');
    
    :root {
        --primary: #6366f1;
        --primary-glow: rgba(99, 102, 241, 0.5);
        --secondary: #a855f7;
        --accent: #f472b6;
        --bg-dark: #0f172a;
        --card-bg: rgba(255, 255, 255, 0.05);
        --text-main: #f8fafc;
        --text-muted: #94a3b8;
    }

    * {
        font-family: 'Inter', sans-serif;
    }
    
    h1, h2, h3, .hero-title {
        font-family: 'Space Grotesk', sans-serif;
    }
    
    .stApp {
        background: radial-gradient(circle at 0% 0%, #1e1b4b 0%, #0f172a 100%);
        background-attachment: fixed;
    }

    /* Mesh Gradient Overlay */
    .stApp::before {
        content: "";
        position: fixed;
        top: 0;
        left: 0;
        width: 100%;
        height: 100%;
        background: 
            radial-gradient(at 0% 0%, rgba(99, 102, 241, 0.15) 0px, transparent 50%),
            radial-gradient(at 100% 0%, rgba(168, 85, 247, 0.15) 0px, transparent 50%),
            radial-gradient(at 100% 100%, rgba(244, 114, 182, 0.15) 0px, transparent 50%),
            radial-gradient(at 0% 100%, rgba(99, 102, 241, 0.15) 0px, transparent 50%);
        pointer-events: none;
        z-index: 0;
    }
    
    .main-container {
        background: rgba(15, 23, 42, 0.6);
        backdrop-filter: blur(20px) saturate(180%);
        -webkit-backdrop-filter: blur(20px) saturate(180%);
        border: 1px solid rgba(255, 255, 255, 0.1);
        border-radius: 32px;
        padding: 4rem;
        box-shadow: 0 25px 50px -12px rgba(0, 0, 0, 0.5);
        margin: 2rem auto;
        max-width: 1000px;
        position: relative;
        z-index: 1;
    }
    
    .hero-title {
        font-size: 4.5rem;
        font-weight: 800;
        background: linear-gradient(to right, #818cf8, #c084fc, #f472b6);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        text-align: center;
        margin-bottom: 0.5rem;
        letter-spacing: -0.05em;
        animation: titleReveal 1.2s cubic-bezier(0.16, 1, 0.3, 1);
    }
    
    .hero-subtitle {
        font-size: 1.4rem;
        color: var(--text-muted);
        text-align: center;
        margin-bottom: 3rem;
        font-weight: 400;
        letter-spacing: 0.02em;
        animation: fadeIn 1.5s ease-out;
    }
    
    .song-card {
        background: rgba(30, 41, 59, 0.4);
        backdrop-filter: blur(10px);
        border-radius: 20px;
        padding: 1.5rem;
        margin: 1.2rem 0;
        border: 1px solid rgba(255, 255, 255, 0.05);
        transition: all 0.4s cubic-bezier(0.4, 0, 0.2, 1);
        display: flex;
        align-items: center;
        gap: 1.5rem;
        position: relative;
        overflow: hidden;
    }
    
    .song-card:hover {
        transform: scale(1.02) translateY(-5px);
        background: rgba(30, 41, 59, 0.6);
        border-color: rgba(129, 140, 248, 0.3);
        box-shadow: 0 20px 40px -15px rgba(0, 0, 0, 0.5), 0 0 20px rgba(99, 102, 241, 0.1);
    }
    
    .song-card::after {
        content: '';
        position: absolute;
        top: 0;
        right: 0;
        width: 100%;
        height: 100%;
        background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.03), transparent);
        transform: translateX(-100%);
        transition: 0.6s;
    }
    
    .song-card:hover::after {
        transform: translateX(100%);
    }
    
    .song-index {
        font-family: 'Space Grotesk', sans-serif;
        font-size: 2.5rem;
        font-weight: 700;
        color: rgba(129, 140, 248, 0.2);
        min-width: 60px;
        text-align: center;
    }
    
    .song-title {
        font-size: 1.4rem;
        font-weight: 700;
        color: var(--text-main);
        margin-bottom: 0.2rem;
    }
    
    .song-artist {
        font-size: 1.1rem;
        color: #818cf8;
        font-weight: 500;
    }
    
    .tag {
        background: rgba(129, 140, 248, 0.1);
        color: #818cf8;
        padding: 0.4rem 1rem;
        border-radius: 100px;
        font-size: 0.75rem;
        font-weight: 600;
        text-transform: uppercase;
        letter-spacing: 0.05em;
        border: 1px solid rgba(129, 140, 248, 0.2);
    }
    
    .question-card {
        background: rgba(30, 41, 59, 0.3);
        border-radius: 24px;
        padding: 3rem;
        margin: 2rem 0;
        border: 1px solid rgba(255, 255, 255, 0.05);
        animation: slideUp 0.6s cubic-bezier(0.16, 1, 0.3, 1);
    }
    
    .question-text {
        font-size: 1.8rem;
        font-weight: 700;
        color: var(--text-main);
        margin-bottom: 2rem;
        line-height: 1.3;
    }
    
    /* Custom Radio Styling */
    [data-testid="stRadio"] > div {
        background: transparent !important;
        padding: 0 !important;
        gap: 0.8rem;
    }
    
    [data-testid="stRadio"] label {
        background: rgba(255, 255, 255, 0.03) !important;
        border: 1px solid rgba(255, 255, 255, 0.05) !important;
        padding: 1.2rem 1.5rem !important;
        border-radius: 16px !important;
        transition: all 0.3s ease !important;
        color: var(--text-main) !important;
        margin-bottom: 0.5rem !important;
        width: 100%;
    }
    
    [data-testid="stRadio"] label:hover {
        background: rgba(255, 255, 255, 0.08) !important;
        border-color: rgba(129, 140, 248, 0.3) !important;
        transform: translateX(10px);
    }

    [data-testid="stRadio"] label p {
        color: var(--text-main) !important;
        font-size: 1.1rem !important;
        font-weight: 500 !important;
    }

    /* Multiselect Styling */
    .stMultiSelect div[data-baseweb="select"] {
        background: rgba(255, 255, 255, 0.03) !important;
        border: 1px solid rgba(255, 255, 255, 0.1) !important;
        border-radius: 16px !important;
        color: white !important;
    }

    .stMultiSelect span[data-baseweb="tag"] {
        background: var(--primary) !important;
        border-radius: 8px !important;
    }
    
    /* Button Styling */
    div.stButton > button {
        background: linear-gradient(135deg, #6366f1 0%, #a855f7 100%) !important;
        color: white !important;
        border: none !important;
        padding: 1rem 2rem !important;
        border-radius: 16px !important;
        font-size: 1.2rem !important;
        font-weight: 700 !important;
        letter-spacing: 0.02em !important;
        transition: all 0.4s cubic-bezier(0.16, 1, 0.3, 1) !important;
        box-shadow: 0 10px 20px -5px rgba(99, 102, 241, 0.4) !important;
        text-transform: uppercase;
    }
    
    div.stButton > button:hover {
        transform: translateY(-3px) scale(1.02) !important;
        box-shadow: 0 20px 30px -10px rgba(99, 102, 241, 0.6) !important;
        filter: brightness(1.1);
    }

    div.stButton > button:active {
        transform: translateY(0) scale(0.98) !important;
    }
    
    .ai-reasoning {
        background: rgba(129, 140, 248, 0.05);
        border-left: 3px solid #818cf8;
        padding: 1rem;
        margin-top: 1rem;
        border-radius: 0 12px 12px 0;
        font-size: 0.95rem;
        color: var(--text-muted);
        line-height: 1.5;
    }
    
    .success-box {
        background: linear-gradient(135deg, rgba(34, 197, 94, 0.1) 0%, rgba(20, 184, 166, 0.1) 100%);
        border: 1px solid rgba(34, 197, 94, 0.2);
        padding: 2rem;
        border-radius: 24px;
        text-align: center;
        color: #4ade80;
        font-weight: 700;
        font-size: 1.2rem;
    }

    /* Animations */
    @keyframes titleReveal {
        from { opacity: 0; transform: translateY(40px) scale(0.95); filter: blur(10px); }
        to { opacity: 1; transform: translateY(0) scale(1); filter: blur(0); }
    }
    
    @keyframes slideUp {
        from { opacity: 0; transform: translateY(30px); }
        to { opacity: 1; transform: translateY(0); }
    }
    
    @keyframes fadeIn {
        from { opacity: 0; }
        to { opacity: 1; }
    }

    /* Progress Bar */
    .stProgress > div > div > div > div {
        background: linear-gradient(to right, #6366f1, #a855f7) !important;
    }

    /* Scrollbar */
    ::-webkit-scrollbar {
        width: 8px;
    }
    ::-webkit-scrollbar-track {
        background: var(--bg-dark);
    }
    ::-webkit-scrollbar-thumb {
        background: #334155;
        border-radius: 10px;
    }
    ::-webkit-scrollbar-thumb:hover {
        background: #475569;
    }
</style>

""", unsafe_allow_html=True)

# Quiz Questions
QUIZ_QUESTIONS = [
    {
        "question": "How would you describe your current mood?",
        "options": ["Happy & Upbeat 😄", "Calm & Relaxed 😌", "Reflective & Thoughtful 🤔", "Energized & Motivated 💪", "Romantic & Dreamy 💕"],
        "key": "mood"
    },
    {
        "question": "What's your ideal Friday night?",
        "options": ["Dancing at a party 🎉", "Netflix and chill 📺", "Deep conversations with friends 💬", "Working out or sports 🏋️", "Candlelit dinner 🕯️"],
        "key": "friday_night"
    },
    {
        "question": "Pick a weather that matches your vibe:",
        "options": ["Sunny & warm ☀️", "Cloudy & cool ☁️", "Rainy & cozy 🌧️", "Storm & intense ⛈️", "Sunset & peaceful 🌅"],
        "key": "weather"
    },
    {
        "question": "What's your energy level right now?",
        "options": ["Super high! ⚡", "Relaxed & steady 🌊", "Low & contemplative 🌙", "Ready to conquer! 🏆", "Soft & gentle 🦋"],
        "key": "energy"
    },
    {
        "question": "Choose a color that speaks to you:",
        "options": ["Bright Yellow 💛", "Ocean Blue 💙", "Deep Purple 💜", "Fiery Red ❤️", "Soft Pink 💗"],
        "key": "color"
    },
    {
        "question": "What kind of lyrics do you prefer?",
        "options": ["Uplifting & positive ✨", "Mellow & smooth 🎵", "Deep & meaningful 📖", "Powerful & inspiring 💥", "Sweet & emotional 💌"],
        "key": "lyrics"
    },
    {
        "question": "If your life was a movie, what genre would it be?",
        "options": ["Comedy 😂", "Indie Drama 🎬", "Psychological Thriller 🧠", "Action Adventure 🎯", "Romance 💑"],
        "key": "movie"
    },
    {
        "question": "What time of day do you feel most alive?",
        "options": ["Morning & fresh ☀️", "Afternoon & steady 🌤️", "Late night & introspective 🌙", "Peak hours & busy 📈", "Golden hour & magical ✨"],
        "key": "time_of_day"
    },
    {
        "question": "Which languages are you familiar with? (Select all that apply)",
        "options": ["English", "Spanish", "French", "German", "Italian", "Portuguese", "Japanese", "Korean", "Chinese","Punjabi", "Hindi","Bengali", "Arabic", "Russian", "Turkish", "Indonesian"],
        "key": "languages",
        "type": "multiselect"
    }
]

# Initialize session state
if 'stage' not in st.session_state:
    st.session_state.stage = 'landing'
if 'quiz_index' not in st.session_state:
    st.session_state.quiz_index = 0
if 'quiz_answers' not in st.session_state:
    st.session_state.quiz_answers = {}
if 'playlist' not in st.session_state:
    st.session_state.playlist = []
if 'spotify_connected' not in st.session_state:
    st.session_state.spotify_connected = False
if 'spotify_playlist_id' not in st.session_state:
    st.session_state.spotify_playlist_id = None

# API Configuration
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
SPOTIFY_CLIENT_ID = os.getenv("SPOTIFY_CLIENT_ID")
SPOTIFY_CLIENT_SECRET = os.getenv("SPOTIFY_CLIENT_SECRET")
SPOTIFY_REDIRECT_URI = os.getenv("SPOTIFY_REDIRECT_URI")

# --- MULTI-USER TOKEN MANAGEMENT ---
class StreamlitSessionCacheHandler(spotipy.cache_handler.CacheHandler):
    """
    Custom cache handler for Spotify tokens that stores them in Streamlit session state.
    This is essential for web deployment to isolate tokens between different users.
    """
    def __init__(self):
        if 'spotify_token' not in st.session_state:
            st.session_state.spotify_token = None

    def get_cached_token(self):
        return st.session_state.spotify_token

    def save_token_to_cache(self, token_info):
        st.session_state.spotify_token = token_info

def get_groq_client():
    """Initialize Groq client"""
    if GROQ_API_KEY:
        return Groq(api_key=GROQ_API_KEY)
    return None

def get_spotify_oauth():
    """Initialize Spotify OAuth object with session-based cache"""
    scope = "playlist-modify-public playlist-modify-private"
    return SpotifyOAuth(
        client_id=SPOTIFY_CLIENT_ID,
        client_secret=SPOTIFY_CLIENT_SECRET,
        redirect_uri=SPOTIFY_REDIRECT_URI,
        scope=scope,
        cache_handler=StreamlitSessionCacheHandler(),
        open_browser=False
    )

def get_spotify_client():
    """Initialize Spotify client using session-based cache"""
    if SPOTIFY_CLIENT_ID and SPOTIFY_CLIENT_SECRET:
        try:
            sp_oauth = get_spotify_oauth()
            
            # Check for cached token in session state
            token_info = sp_oauth.get_cached_token()
            
            if token_info:
                # Refresh token if expired
                if sp_oauth.is_token_expired(token_info):
                    token_info = sp_oauth.refresh_access_token(token_info['refresh_token'])
                
                return spotipy.Spotify(auth=token_info['access_token'])
            
            # Check if we got a code from the redirect
            query_params = st.query_params
            if 'code' in query_params:
                code = query_params['code']
                try:
                    token_info = sp_oauth.get_access_token(code, as_dict=True, check_cache=False)
                    if token_info:
                        # Clear the code from URL
                        st.query_params.clear()
                        return spotipy.Spotify(auth=token_info['access_token'])
                except Exception as e:
                    st.error(f"Error getting token: {str(e)}")
                    return None
            
            return None
            
        except Exception as e:
            st.error(f"Spotify authentication error: {str(e)}")
            return None
    return None

def generate_playlist_with_groq(quiz_answers: Dict, count: int = 20) -> List[Dict]:
    """Generate playlist using Groq AI based on quiz answers"""
    groq_client = get_groq_client()
    
    if not groq_client:
        st.warning("⚠️ Groq API not configured. Using demo mode.")
        return get_demo_playlist()
    
    # Create prompt from quiz answers
    languages = quiz_answers.get('languages', ['English'])
    if isinstance(languages, list):
        languages_str = ", ".join(languages)
    else:
        languages_str = str(languages)

    prompt = f"""Based on the following personality quiz responses, recommend {count} songs that match this person's vibe perfectly.

IMPORTANT: The user is familiar with these languages: {languages_str}. 
You MUST ONLY recommend songs in one of these languages. Do not include songs in other languages.

Quiz Responses:
{json.dumps(quiz_answers, indent=2)}

Please analyze the responses and recommend songs that match their mood, energy, and preferences. For each song, provide:
1. Song title
2. Artist name
3. Genre
4. A brief reason why this song matches their vibe (MAX 10 words)

Format your response as a JSON array with this structure:
[
  {{
    "title": "Song Title",
    "artist": "Artist Name",
    "genre": "Genre",
    "reasoning": "Why this song matches their vibe"
  }}
]

Only return the JSON array, no additional text."""

    try:
        with st.spinner("✨ AI is curating your perfect playlist..."):
            response = groq_client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=[
                    {"role": "system", "content": "You are a music expert who understands personality and creates perfect playlists. Always respond with valid JSON."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.5,
                max_tokens=6000
            )
            
            content = response.choices[0].message.content.strip()
            
            # Extract JSON if wrapped in code blocks
            if "```json" in content:
                content = content.split("```json")[1].split("```")[0].strip()
            elif "```" in content:
                content = content.split("```")[1].split("```")[0].strip()
            
            try:
                songs = json.loads(content)
            except json.JSONDecodeError as je:
                # If JSON is truncated, try to fix the last entry
                if "Expecting value" in str(je) or "Unterminated string" in str(je):
                    # Find the last complete song object
                    last_bracket = content.rfind('}')
                    if last_bracket != -1:
                        fixed_content = content[:last_bracket+1] + ']'
                        try:
                            songs = json.loads(fixed_content)
                            st.warning("⚠️ Some songs were omitted due to length limits.")
                        except:
                            raise je
                    else:
                        raise je
                else:
                    raise je
            
            return songs[:count]
            
    except Exception as e:
        print(f"Error generating playlist with AI: {str(e)}")
        st.error(f"Error generating playlist with AI: {str(e)}")
        return get_demo_playlist()

def get_demo_playlist() -> List[Dict]:
    """Fallback demo playlist"""
    return [
        {"title": "Blinding Lights", "artist": "The Weeknd", "genre": "Pop", "reasoning": "High energy and upbeat vibe"},
        {"title": "Levitating", "artist": "Dua Lipa", "genre": "Pop", "reasoning": "Perfect for happy moods"},
        {"title": "Good 4 U", "artist": "Olivia Rodrigo", "genre": "Pop", "reasoning": "Energetic and powerful"},
        {"title": "Shivers", "artist": "Ed Sheeran", "genre": "Pop", "reasoning": "Romantic and catchy"},
        {"title": "Heat Waves", "artist": "Glass Animals", "genre": "Indie", "reasoning": "Chill yet engaging"},
        {"title": "Stay", "artist": "The Kid LAROI & Justin Bieber", "genre": "Pop", "reasoning": "Emotional and melodic"},
        {"title": "Peaches", "artist": "Justin Bieber", "genre": "R&B", "reasoning": "Smooth and relaxed"},
        {"title": "Montero", "artist": "Lil Nas X", "genre": "Hip-Hop", "reasoning": "Bold and confident"},
        {"title": "drivers license", "artist": "Olivia Rodrigo", "genre": "Pop", "reasoning": "Deep emotional resonance"},
        {"title": "Save Your Tears", "artist": "The Weeknd", "genre": "Pop", "reasoning": "Uplifting with depth"},
        {"title": "Positions", "artist": "Ariana Grande", "genre": "R&B", "reasoning": "Sweet and romantic"},
        {"title": "Willow", "artist": "Taylor Swift", "genre": "Pop", "reasoning": "Dreamy and enchanting"}
    ]

def search_and_add_to_spotify_playlist(songs: List[Dict], playlist_id: str):
    """Search for songs on Spotify and add them to playlist"""
    sp = get_spotify_client()
    if not sp:
        return False
    
    track_uris = []
    
    with st.spinner("🔍 Finding songs on Spotify..."):
        for song in songs:
            try:
                query = f"{song['title']} {song['artist']}"
                results = sp.search(q=query, type='track', limit=1)
                
                if results['tracks']['items']:
                    track_uris.append(results['tracks']['items'][0]['uri'])
                    time.sleep(0.1)  # Rate limiting
                    
            except Exception as e:
                st.warning(f"Couldn't find '{song['title']}' on Spotify")
                continue
    
    # Add tracks to playlist in batches of 100
    try:
        for i in range(0, len(track_uris), 100):
            sp.playlist_add_items(playlist_id, track_uris[i:i+100])
        return True
    except Exception as e:
        st.error(f"Error adding songs to playlist: {str(e)}")
        return False

def create_spotify_playlist(playlist_name: str) -> Optional[str]:
    """Create a new Spotify playlist"""
    sp = get_spotify_client()
    if not sp:
        return None
    
    try:
        user_id = sp.current_user()['id']
        playlist = sp.user_playlist_create(
            user=user_id,
            name=playlist_name,
            public=True,
            description="Created by VibeSync - AI-powered playlist generator"
        )
        return playlist['id']
    except Exception as e:
        st.error(f"Error creating Spotify playlist: {str(e)}")
        return None

def display_song_card(song: Dict, index: int):
    """Display a song in a beautiful card format"""
    reasoning_html = f'<div class="ai-reasoning">“{song.get("reasoning", "Perfectly curated for your vibe.")}”</div>' if song.get("reasoning") else ""
    
    st.markdown(f"""
    <div class="song-card">
        <div class="song-index">{index:02d}</div>
        <div style="flex: 1;">
            <div class="song-title">{song['title']}</div>
            <div class="song-artist">{song['artist']}</div>
            <div style="margin-top: 0.8rem; display: flex; align-items: center; gap: 0.5rem;">
                <span class="tag">{song.get('genre', 'Music')}</span>
            </div>
            {reasoning_html}
        </div>
    </div>
    """, unsafe_allow_html=True)

# Check Spotify connection
def check_spotify_auth():
    """Check if Spotify is authenticated"""
    sp = get_spotify_client()
    if sp:
        try:
            user = sp.current_user()
            st.session_state.spotify_connected = True
            return user
        except:
            st.session_state.spotify_connected = False
            return None
    return None

# Landing Page
if st.session_state.stage == 'landing':
    st.markdown('<div class="main-container">', unsafe_allow_html=True)
    st.markdown('<h1 class="hero-title">VibeSync</h1>', unsafe_allow_html=True)
    st.markdown('<p class="hero-subtitle">Your personality, translated into sound.</p>', unsafe_allow_html=True)
    
    # Spotify Connection Status
    user = check_spotify_auth()
    if user:
        st.markdown(f'<div style="text-align: center; margin-bottom: 2rem;"><div class="spotify-badge">✓ Connected as {user["display_name"]}</div></div>', unsafe_allow_html=True)
    elif SPOTIFY_CLIENT_ID:
        st.markdown('<div style="text-align: center; margin: 1rem 0;"><p style="color: var(--text-muted);">Connect Spotify to export your magic.</p></div>', unsafe_allow_html=True)
        col1, col2, col3 = st.columns([1, 1, 1])
        with col2:
            if st.button("🎵 Connect Spotify", key="connect_spotify"):
                # Generate auth URL
                sp_oauth = get_spotify_oauth()
                auth_url = sp_oauth.get_authorize_url()
                
                st.markdown(f"""
                <div style="background: rgba(255,255,255,0.05); padding: 2rem; border-radius: 24px; margin: 1rem 0; border: 1px solid rgba(255,255,255,0.1); text-align: center;">
                    <h4 style="color: white; margin-bottom: 1rem;">Spotify Integration</h4>
                    <a href="{auth_url}" target="_blank" style="
                        display: inline-block;
                        background: #1DB954;
                        color: white;
                        padding: 1rem 2.5rem;
                        border-radius: 100px;
                        text-decoration: none;
                        font-weight: 700;
                        margin: 1rem 0;
                        transition: 0.3s;
                    ">Login with Spotify</a>
                </div>
                """, unsafe_allow_html=True)
    
    st.markdown("""
    <div style="margin: 3rem 0; padding: 2rem; background: rgba(255,255,255,0.02); border-radius: 24px; border: 1px solid rgba(255,255,255,0.05);">
        <div style="display: grid; grid-template-columns: repeat(3, 1fr); gap: 2rem; text-align: center;">
            <div>
                <div style="font-size: 2rem; margin-bottom: 0.5rem;">🧠</div>
                <div style="color: white; font-weight: 600;">AI Analysis</div>
                <div style="color: var(--text-muted); font-size: 0.9rem;">Deep mood mapping</div>
            </div>
            <div>
                <div style="font-size: 2rem; margin-bottom: 0.5rem;">✨</div>
                <div style="color: white; font-weight: 600;">Curated Vibes</div>
                <div style="color: var(--text-muted); font-size: 0.9rem;">20+ perfect tracks</div>
            </div>
            <div>
                <div style="font-size: 2rem; margin-bottom: 0.5rem;">🎧</div>
                <div style="color: white; font-weight: 600;">Instant Export</div>
                <div style="color: var(--text-muted); font-size: 0.9rem;">Direct to Spotify</div>
            </div>
        </div>
    </div>
    """, unsafe_allow_html=True)
    
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        if st.button("🚀 Start Your Journey", key="start_quiz"):
            st.session_state.stage = 'quiz'
            st.rerun()
    st.markdown('</div>', unsafe_allow_html=True)

# Quiz Stage
elif st.session_state.stage == 'quiz':
    current_q = QUIZ_QUESTIONS[st.session_state.quiz_index]
    
    st.markdown('<div class="main-container">', unsafe_allow_html=True)
    st.markdown('<h1 class="hero-title">VibeSync Quiz</h1>', unsafe_allow_html=True)
    st.markdown(f'<p style="text-align: center; color: var(--text-muted); margin-bottom: 2rem;">Question {st.session_state.quiz_index + 1} of {len(QUIZ_QUESTIONS)}</p>', unsafe_allow_html=True)
    
    # Progress bar
    progress = (st.session_state.quiz_index) / len(QUIZ_QUESTIONS)
    st.progress(progress)
    
    st.markdown('<div class="question-card">', unsafe_allow_html=True)
    st.markdown(f'<p class="question-text">{current_q["question"]}</p>', unsafe_allow_html=True)
    
    if current_q.get("type") == "multiselect":
        answer = st.multiselect("Choose languages:", current_q["options"], key=f"q_{st.session_state.quiz_index}")
    else:
        answer = st.radio("Choose one:", current_q["options"], key=f"q_{st.session_state.quiz_index}", label_visibility="collapsed")
    
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        if st.button("Next Question ➡️", key=f"next_{st.session_state.quiz_index}"):
            # Save answer
            st.session_state.quiz_answers[current_q["key"]] = answer
            
            if st.session_state.quiz_index < len(QUIZ_QUESTIONS) - 1:
                st.session_state.quiz_index += 1
                st.rerun()
            else:
                # Quiz completed, generate playlist with AI
                st.session_state.playlist = generate_playlist_with_groq(st.session_state.quiz_answers)
                st.session_state.stage = 'results'
                st.rerun()
    
    st.markdown('</div></div>', unsafe_allow_html=True)

# Results Stage
elif st.session_state.stage == 'results':
    st.markdown('<div class="main-container">', unsafe_allow_html=True)
    st.markdown('<h1 class="hero-title">Your Sonic Identity</h1>', unsafe_allow_html=True)
    st.markdown(f"<p style='text-align: center; color: var(--text-muted); margin-bottom: 3rem;'>A {len(st.session_state.playlist)}-track journey curated by AI.</p>", unsafe_allow_html=True)
    
    # Display playlist
    for idx, song in enumerate(st.session_state.playlist, 1):
        display_song_card(song, idx)
    
    # Add more songs with AI
    st.markdown("<br><br>", unsafe_allow_html=True)
    st.markdown("<h3 style='color: #818cf8; text-align: center; font-family: \"Space Grotesk\", sans-serif;'>Curious for more? 🎵</h3>", unsafe_allow_html=True)
    
    user_request = st.text_input("Describe what kind of songs you'd like to add:", placeholder="e.g., more upbeat songs, slower tempo, different genre...")
    
    if st.button("✨ Get AI Suggestions"):
        if user_request:
            groq_client = get_groq_client()
            if groq_client:
                prompt = f"""Based on the user's request and their current playlist, suggest 5 additional songs.

User's request: {user_request}

Current playlist context:
{json.dumps(st.session_state.playlist[:3], indent=2)}

Recommend 5 songs that match their request. Format as JSON array:
[
  {{
    "title": "Song Title",
    "artist": "Artist Name",
    "genre": "Genre",
    "reasoning": "Why this matches their request (MAX 10 words)"
  }}
]"""
                
                try:
                    with st.spinner("✨ AI is finding more gems..."):
                        response = groq_client.chat.completions.create(
                            model="llama-3.3-70b-versatile",
                            messages=[
                                {"role": "system", "content": "You are a music expert. Always respond with valid JSON."},
                                {"role": "user", "content": prompt}
                            ],
                            temperature=0.8,
                            max_tokens=1000
                        )
                        
                        content = response.choices[0].message.content.strip()
                        if "```json" in content:
                            content = content.split("```json")[1].split("```")[0].strip()
                        elif "```" in content:
                            content = content.split("```")[1].split("```")[0].strip()
                        
                        new_songs = json.loads(content)
                        st.session_state.temp_suggestions = new_songs
                        
                except Exception as e:
                    st.error(f"Error getting AI suggestions: {str(e)}")
            else:
                st.warning("⚠️ Groq API not configured")

    # Display suggestions if they exist
    if 'temp_suggestions' in st.session_state:
        st.markdown("<h4 style='color: #818cf8; margin-top: 2rem; font-family: \"Space Grotesk\", sans-serif;'>AI Suggestions:</h4>", unsafe_allow_html=True)
        for idx, song in enumerate(st.session_state.temp_suggestions, 1):
            col1, col2 = st.columns([5, 1])
            with col1:
                st.markdown(f"""
                <div class="song-card" style="margin: 0.5rem 0; padding: 1rem;">
                    <div style="flex: 1;">
                        <div class="song-title" style="font-size: 1.1rem;">{song['title']}</div>
                        <div class="song-artist" style="font-size: 0.9rem;">{song['artist']}</div>
                        <div class="ai-reasoning" style="font-size: 0.8rem; padding: 0.5rem;">“{song.get('reasoning', 'Matches your request.')}”</div>
                    </div>
                </div>
                """, unsafe_allow_html=True)
            with col2:
                st.markdown("<div style='height: 20px;'></div>", unsafe_allow_html=True)
                if st.button("➕", key=f"add_{idx}_{song['title'][:10]}"):
                    st.session_state.playlist.append(song)
                    
                    # Add to Spotify if connected
                    if st.session_state.spotify_playlist_id:
                        sp = get_spotify_client()
                        if sp:
                            try:
                                query = f"{song['title']} {song['artist']}"
                                results = sp.search(q=query, type='track', limit=1)
                                if results['tracks']['items']:
                                    sp.playlist_add_items(
                                        st.session_state.spotify_playlist_id, 
                                        [results['tracks']['items'][0]['uri']]
                                    )
                                    st.toast(f"✅ Added {song['title']} to Spotify!")
                            except:
                                pass
                    
                    # Remove from suggestions after adding
                    st.session_state.temp_suggestions.pop(idx-1)
                    if not st.session_state.temp_suggestions:
                        del st.session_state.temp_suggestions
                    st.rerun()

    # Export to Spotify
    st.markdown("<br><br>", unsafe_allow_html=True)
    
    if st.session_state.spotify_connected and not st.session_state.spotify_playlist_id:
        st.markdown("""
        <div class="success-box">
            <h3 style="margin: 0; color: #4ade80;">🎧 Ready to export to Spotify!</h3>
        </div>
        """, unsafe_allow_html=True)
        
        col1, col2, col3 = st.columns([1, 2, 1])
        with col2:
            playlist_name = st.text_input("Playlist Name:", value="My VibeSync Playlist")
            if st.button("📤 Create Spotify Playlist"):
                playlist_id = create_spotify_playlist(playlist_name)
                if playlist_id:
                    if search_and_add_to_spotify_playlist(st.session_state.playlist, playlist_id):
                        st.session_state.spotify_playlist_id = playlist_id
                        st.success("✅ Playlist created and songs added to your Spotify account!")
                        st.balloons()
                        st.rerun()
    
    elif st.session_state.spotify_playlist_id:
        st.markdown("""
        <div class="success-box">
            <h3 style="margin: 0; color: #4ade80;">✅ Playlist exported to Spotify!</h3>
            <p style="margin: 0.5rem 0 0 0; color: var(--text-muted);">Check your Spotify account</p>
        </div>
        """, unsafe_allow_html=True)
    
    # Action buttons
    st.markdown("<br><br>", unsafe_allow_html=True)
    col1, col2 = st.columns(2)
    with col1:
        if st.button("🔄 Start Over"):
            st.session_state.stage = 'landing'
            st.session_state.quiz_index = 0
            st.session_state.quiz_answers = {}
            st.session_state.playlist = []
            st.session_state.spotify_playlist_id = None
            if 'temp_suggestions' in st.session_state:
                del st.session_state.temp_suggestions
            st.rerun()
    with col2:
        if st.button("✅ I'm Happy with My Playlist"):
            st.balloons()
            st.success(f"🎉 Awesome! Your playlist of {len(st.session_state.playlist)} songs is ready!")
    st.markdown('</div>', unsafe_allow_html=True)