import gradio as gr
import sqlite3
import hashlib
import uuid
import os
from datetime import datetime
from PIL import Image
import numpy as np
import shutil

# ===============================
# DATABASE LAYER
# ===============================

conn = sqlite3.connect("vibes.db", check_same_thread=False)
cursor = conn.cursor()

def init_db():
    # Users table
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS users(
    username TEXT PRIMARY KEY,
    password TEXT,
    display_name TEXT,
    bio TEXT,
    avatar TEXT,
    created_at TEXT
    )
    """)

    # Posts table (supports both images and videos)
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS posts(
    id TEXT PRIMARY KEY,
    username TEXT,
    caption TEXT,
    media_path TEXT,
    media_type TEXT,
    timestamp TEXT,
    views INTEGER DEFAULT 0
    )
    """)

    # Likes table
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS likes(
    post_id TEXT,
    username TEXT,
    PRIMARY KEY (post_id, username)
    )
    """)

    # Comments table
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS comments(
    id TEXT PRIMARY KEY,
    post_id TEXT,
    username TEXT,
    comment TEXT,
    timestamp TEXT
    )
    """)

    # Followers table
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS followers(
    user TEXT,
    follower TEXT,
    PRIMARY KEY (user, follower)
    )
    """)

    # Notifications table
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS notifications(
    id TEXT PRIMARY KEY,
    username TEXT,
    message TEXT,
    timestamp TEXT,
    read INTEGER DEFAULT 0
    )
    """)

    conn.commit()

init_db()

# Create directories
os.makedirs("uploads", exist_ok=True)
os.makedirs("avatars", exist_ok=True)

# ===============================
# AUTH LAYER
# ===============================

def hash_password(p):
    return hashlib.sha256(p.encode()).hexdigest()

def signup(username, password, display_name):
    if not username or not password:
        return "❌ Username and password required", "", "", ""
    
    cursor.execute("SELECT * FROM users WHERE username=?", (username,))
    if cursor.fetchone():
        return "❌ Username already exists", "", "", ""

    cursor.execute(
        "INSERT INTO users VALUES (?,?,?,?,?,?)",
        (username, hash_password(password), display_name or username, "✨ New to Vibes ✨", None, datetime.now().isoformat())
    )
    conn.commit()
    return "✅ Account created! Please login.", "", "", ""

def login(username, password):
    if not username or not password:
        return "❌ Username and password required", None
    
    cursor.execute("SELECT password, display_name FROM users WHERE username=?", (username,))
    result = cursor.fetchone()

    if not result:
        return "❌ User not found", None

    if result[0] != hash_password(password):
        return "❌ Wrong password", None

    return f"✅ Welcome back, {result[1]}!", username

# ===============================
# POST LAYER (SUPPORTS IMAGES & VIDEOS)
# ===============================

def upload_media(user, media, caption, media_type):
    if user is None or user == "guest":
        return "❌ Login to post", None

    if media is None:
        return f"❌ Select a {media_type}", None

    post_id = str(uuid.uuid4())
    
    # Determine file extension
    if media_type == "image":
        ext = "jpg"
        # Process image
        img = Image.fromarray(media)
        img = img.convert("RGB")
        
        # Resize if too large
        max_size = 1200
        if max(img.size) > max_size:
            ratio = max_size / max(img.size)
            new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
            img = img.resize(new_size, Image.Resampling.LANCZOS)
        
        path = f"uploads/{post_id}.{ext}"
        img.save(path, quality=85)
    else:
        ext = "mp4"
        path = f"uploads/{post_id}.{ext}"
        # Handle video file
        if isinstance(media, str):
            shutil.copy(media, path)
        else:
            return "❌ Invalid video format", None

    cursor.execute(
        "INSERT INTO posts VALUES (?,?,?,?,?,?,?)",
        (post_id, user, caption, path, media_type, datetime.now().isoformat(), 0)
    )
    conn.commit()
    
    # Notify followers
    cursor.execute("SELECT follower FROM followers WHERE user=?", (user,))
    followers = cursor.fetchall()
    for f in followers:
        notif_id = str(uuid.uuid4())
        cursor.execute(
            "INSERT INTO notifications VALUES (?,?,?,?,?)",
            (notif_id, f[0], f"@{user} posted new content!", datetime.now().isoformat(), 0)
        )
    conn.commit()
    
    return f"✅ {media_type.capitalize()} posted!", None

def like_post(post_id, user):
    if user is None or user == "guest":
        return "❤️ Login to like", None
    
    cursor.execute("SELECT username FROM posts WHERE id=?", (post_id,))
    post_owner = cursor.fetchone()
    
    cursor.execute("SELECT * FROM likes WHERE post_id=? AND username=?", (post_id, user))
    if cursor.fetchone():
        cursor.execute("DELETE FROM likes WHERE post_id=? AND username=?", (post_id, user))
        conn.commit()
        
        # Update notification (simplified - just remove like)
        return "🤍", "unliked"
    else:
        cursor.execute("INSERT INTO likes VALUES (?,?)", (post_id, user))
        conn.commit()
        
        # Notify post owner
        if post_owner and post_owner[0] != user:
            notif_id = str(uuid.uuid4())
            cursor.execute(
                "INSERT INTO notifications VALUES (?,?,?,?,?)",
                (notif_id, post_owner[0], f"@{user} liked your post!", datetime.now().isoformat(), 0)
            )
            conn.commit()
        
        return "❤️", "liked"

def add_comment(user, post_id, comment):
    if user is None or user == "guest":
        return "❌ Login to comment"
    if not comment or not comment.strip():
        return "❌ Enter a comment"
    
    cursor.execute("SELECT username FROM posts WHERE id=?", (post_id,))
    post_owner = cursor.fetchone()
    
    cursor.execute(
        "INSERT INTO comments VALUES (?,?,?,?,?)",
        (str(uuid.uuid4()), post_id, user, comment.strip(), datetime.now().isoformat())
    )
    conn.commit()
    
    # Notify post owner
    if post_owner and post_owner[0] != user:
        notif_id = str(uuid.uuid4())
        cursor.execute(
            "INSERT INTO notifications VALUES (?,?,?,?,?)",
            (notif_id, post_owner[0], f"@{user} commented: {comment[:30]}", datetime.now().isoformat(), 0)
        )
        conn.commit()
    
    return "✅ Comment added!"

# ===============================
# FEED RENDERER WITH STICKY HEADER
# ===============================

def render_feed():
    cursor.execute("""
        SELECT posts.id, posts.media_path, posts.caption, posts.media_type,
               users.display_name, users.username, users.avatar,
               COUNT(DISTINCT likes.username) as like_count,
               COUNT(DISTINCT comments.id) as comment_count
        FROM posts 
        JOIN users ON posts.username = users.username
        LEFT JOIN likes ON posts.id = likes.post_id 
        LEFT JOIN comments ON posts.id = comments.video_id
        GROUP BY posts.id 
        ORDER BY posts.timestamp DESC
        LIMIT 50
    """)

    rows = cursor.fetchall()

    if not rows:
        return """
        <div class="empty-state">
            <div class="empty-icon">🎵</div>
            <h3>No content yet</h3>
            <p>Be the first to share something amazing!</p>
        </div>
        """

    html = '<div class="feed-container">'

    for row in rows:
        post_id, media_path, caption, media_type, display_name, username, avatar, like_count, comment_count = row
        
        if media_type == "image":
            media_html = f'<img src="{media_path}" class="feed-media" alt="Post image">'
        else:
            media_html = f'''
            <video class="feed-media" loop controls playsinline preload="metadata">
                <source src="{media_path}" type="video/mp4">
            </video>
            '''
        
        avatar_html = f'<img src="{avatar}" class="avatar-img">' if avatar else '<div class="avatar-placeholder">👤</div>'
        
        html += f"""
        <div class="feed-post" data-id="{post_id}">
            {media_html}
            <div class="post-overlay">
                <div class="post-info">
                    <div class="creator-section">
                        {avatar_html}
                        <div class="creator-details">
                            <div class="display-name">{display_name}</div>
                            <div class="username">@{username}</div>
                        </div>
                    </div>
                    <div class="post-caption">{caption if caption else "✨"}</div>
                </div>
                <div class="post-actions">
                    <button class="action-btn like-btn" data-id="{post_id}">
                        <div class="action-icon">❤️</div>
                        <div class="action-count">{like_count}</div>
                    </button>
                    <button class="action-btn comment-btn" data-id="{post_id}">
                        <div class="action-icon">💬</div>
                        <div class="action-count">{comment_count}</div>
                    </button>
                    <button class="action-btn share-btn" data-id="{post_id}">
                        <div class="action-icon">📤</div>
                    </button>
                </div>
            </div>
        </div>
        """

    html += "</div>"
    return html

# ===============================
# PROFILE LAYER
# ===============================

def load_profile(user):
    if user is None or user == "guest":
        return "### 👤 Guest Mode\n\nLogin to see your profile", []

    cursor.execute("SELECT display_name, bio, avatar, created_at FROM users WHERE username=?", (user,))
    u = cursor.fetchone()
    
    if not u:
        return "### User not found", []

    cursor.execute("SELECT media_path, media_type FROM posts WHERE username=? ORDER BY timestamp DESC", (user,))
    posts = [(p[0], p[1]) for p in cursor.fetchall()]

    cursor.execute("SELECT COUNT(*) FROM followers WHERE user=?", (user,))
    followers = cursor.fetchone()[0]

    cursor.execute("SELECT COUNT(*) FROM followers WHERE follower=?", (user,))
    following = cursor.fetchone()[0]

    cursor.execute("SELECT COUNT(*) FROM posts WHERE username=?", (user,))
    posts_count = cursor.fetchone()[0]
    
    cursor.execute("SELECT COUNT(*) FROM likes WHERE post_id IN (SELECT id FROM posts WHERE username=?)", (user,))
    total_likes = cursor.fetchone()[0]

    avatar_html = f'<img src="{u[2]}" class="profile-avatar-img">' if u[2] else '<div class="profile-avatar-placeholder">👤</div>'

    profile = f"""
<div class="profile-header">
    {avatar_html}
    <h2>{u[0]}</h2>
    <p class="profile-username">@{user}</p>
    <p class="profile-bio">{u[1]}</p>
    <p class="profile-joined">Joined {u[3][:10]}</p>
    <div class="profile-stats">
        <div class="stat">
            <div class="stat-number">{posts_count}</div>
            <div class="stat-label">Posts</div>
        </div>
        <div class="stat">
            <div class="stat-number">{followers}</div>
            <div class="stat-label">Followers</div>
        </div>
        <div class="stat">
            <div class="stat-number">{following}</div>
            <div class="stat-label">Following</div>
        </div>
        <div class="stat">
            <div class="stat-number">{total_likes}</div>
            <div class="stat-label">Likes</div>
        </div>
    </div>
</div>
"""
    # Return posts as gallery items
    gallery_items = [p[0] for p in posts[:9]]
    return profile, gallery_items

def update_bio(user, new_bio):
    if user is None or user == "guest":
        return "❌ Login to update bio"
    if not new_bio or not new_bio.strip():
        return "❌ Please enter a bio"
    
    cursor.execute("UPDATE users SET bio=? WHERE username=?", (new_bio.strip(), user))
    conn.commit()
    return "✅ Bio updated!"

def follow_user(current_user, target_user):
    if current_user is None or current_user == "guest":
        return "❌ Login to follow"
    if current_user == target_user:
        return "❌ Cannot follow yourself"
    
    cursor.execute("SELECT * FROM users WHERE username=?", (target_user,))
    if not cursor.fetchone():
        return "❌ User not found"

    cursor.execute("SELECT * FROM followers WHERE user=? AND follower=?", (target_user, current_user))
    if cursor.fetchone():
        cursor.execute("DELETE FROM followers WHERE user=? AND follower=?", (target_user, current_user))
        conn.commit()
        
        # Create unfollow notification
        notif_id = str(uuid.uuid4())
        cursor.execute(
            "INSERT INTO notifications VALUES (?,?,?,?,?)",
            (notif_id, target_user, f"@{current_user} unfollowed you", datetime.now().isoformat(), 0)
        )
        conn.commit()
        
        return f"✅ Unfollowed {target_user}"
    else:
        cursor.execute("INSERT INTO followers VALUES (?,?)", (target_user, current_user))
        conn.commit()
        
        # Create follow notification
        notif_id = str(uuid.uuid4())
        cursor.execute(
            "INSERT INTO notifications VALUES (?,?,?,?,?)",
            (notif_id, target_user, f"@{current_user} started following you!", datetime.now().isoformat(), 0)
        )
        conn.commit()
        
        return f"✅ Following {target_user}"

# ===============================
# NOTIFICATIONS
# ===============================

def get_notifications(user):
    if user is None or user == "guest":
        return []
    
    cursor.execute(
        "SELECT message, timestamp, read FROM notifications WHERE username=? ORDER BY timestamp DESC LIMIT 20",
        (user,)
    )
    notifs = cursor.fetchall()
    
    # Mark as read
    cursor.execute("UPDATE notifications SET read=1 WHERE username=?", (user,))
    conn.commit()
    
    return notifs

def render_notifications(user):
    notifs = get_notifications(user)
    if not notifs:
        return '<div class="empty-state">🔔 No notifications yet</div>'
    
    html = '<div class="notifications-list">'
    for n in notifs:
        unread_class = "" if n[2] else "unread"
        html += f'''
        <div class="notification-item {unread_class}">
            <div class="notification-message">{n[0]}</div>
            <div class="notification-time">{n[1][:16]}</div>
        </div>
        '''
    html += '</div>'
    return html

# ===============================
# SEARCH
# ===============================

def search_users(q):
    if len(q) < 2:
        return []
    cursor.execute(
        "SELECT username, display_name, bio FROM users WHERE username LIKE ? OR display_name LIKE ? LIMIT 10",
        ('%' + q + '%', '%' + q + '%')
    )
    return cursor.fetchall()

# ===============================
# BEAUTIFUL CSS - PURPLE/ORANGE THEME WITH STICKY HEADER
# ===============================

custom_css = """
:root {
    --primary: #8B5CF6;
    --primary-dark: #7C3AED;
    --secondary: #F59E0B;
    --accent: #EC4899;
    --bg: #0F0F0F;
    --card: #1A1A1A;
    --border: #2A2A2A;
    --text: #FFFFFF;
    --text-secondary: #A1A1AA;
}

* {
    margin: 0;
    padding: 0;
    box-sizing: border-box;
}

body, .gradio-container {
    background: var(--bg) !important;
    color: var(--text) !important;
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
}

/* Hide default Gradio header */
header, .gradio-header, .gradio-header-bar {
    display: none !important;
}

/* ===== STICKY HEADER ===== */
.app-header {
    position: sticky;
    top: 0;
    left: 0;
    right: 0;
    background: linear-gradient(135deg, var(--primary), var(--secondary));
    backdrop-filter: blur(0px);
    padding: 12px 20px;
    z-index: 100;
    box-shadow: 0 4px 20px rgba(0,0,0,0.3);
}

.app-header h1 {
    font-size: 24px;
    font-weight: 700;
    background: linear-gradient(135deg, #fff, #FFE4B5);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin: 0;
}

/* ===== BOTTOM NAVIGATION ===== */
.bottom-nav {
    position: fixed;
    bottom: 0;
    left: 0;
    right: 0;
    background: rgba(18, 18, 18, 0.95);
    backdrop-filter: blur(20px);
    display: flex;
    justify-content: space-around;
    padding: 10px 20px 20px;
    z-index: 100;
    border-top: 1px solid var(--border);
}

.nav-item {
    text-align: center;
    cursor: pointer;
    transition: all 0.2s;
    padding: 5px 10px;
    border-radius: 40px;
}

.nav-item:hover {
    transform: translateY(-2px);
}

.nav-icon {
    font-size: 24px;
    margin-bottom: 4px;
}

.nav-label {
    font-size: 11px;
    color: var(--text-secondary);
}

.nav-item.active .nav-icon,
.nav-item.active .nav-label {
    color: var(--secondary);
    font-weight: 600;
}

/* ===== FEED CONTAINER ===== */
.feed-container {
    height: calc(100vh - 60px);
    overflow-y: scroll;
    scroll-snap-type: y mandatory;
    scrollbar-width: none;
    padding-bottom: 70px;
}

.feed-container::-webkit-scrollbar {
    display: none;
}

/* ===== FEED POST ===== */
.feed-post {
    scroll-snap-align: start;
    height: calc(100vh - 60px);
    position: relative;
    background: var(--bg);
}

.feed-media {
    width: 100%;
    height: 100%;
    object-fit: contain;
}

/* Post Overlay */
.post-overlay {
    position: absolute;
    bottom: 0;
    left: 0;
    right: 0;
    background: linear-gradient(to top, rgba(0,0,0,0.9) 0%, transparent 100%);
    padding: 80px 20px 30px;
}

/* Creator Section */
.creator-section {
    display: flex;
    align-items: center;
    gap: 12px;
    margin-bottom: 12px;
}

.avatar-img, .avatar-placeholder {
    width: 48px;
    height: 48px;
    border-radius: 50%;
    object-fit: cover;
}

.avatar-placeholder {
    background: linear-gradient(135deg, var(--primary), var(--secondary));
    display: flex;
    align-items: center;
    justify-content: center;
    font-size: 24px;
}

.creator-details {
    flex: 1;
}

.display-name {
    font-weight: 600;
    font-size: 16px;
    color: var(--text);
}

.username {
    font-size: 13px;
    color: var(--text-secondary);
}

.post-caption {
    font-size: 14px;
    line-height: 1.4;
    color: var(--text);
    margin-top: 8px;
}

/* Action Buttons */
.post-actions {
    position: absolute;
    right: 16px;
    bottom: 100px;
    display: flex;
    flex-direction: column;
    gap: 24px;
}

.action-btn {
    background: rgba(0,0,0,0.6);
    border: none;
    border-radius: 50%;
    width: 48px;
    height: 48px;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    gap: 4px;
    cursor: pointer;
    transition: transform 0.2s;
    color: white;
}

.action-btn:hover {
    transform: scale(1.1);
    background: var(--primary);
}

.action-icon {
    font-size: 24px;
}

.action-count {
    font-size: 11px;
}

/* Tabs Container */
.tabs {
    margin-bottom: 70px !important;
    padding: 0 !important;
}

.tab-nav {
    display: none !important;
}

/* Forms */
.gr-box, .gr-form {
    background: var(--card) !important;
    border: 1px solid var(--border) !important;
    border-radius: 16px !important;
}

input, textarea {
    border-radius: 40px !important;
    background: #1F1F1F !important;
    border: 1px solid var(--border) !important;
    color: var(--text) !important;
    padding: 12px 20px !important;
}

input:focus, textarea:focus {
    border-color: var(--secondary) !important;
    outline: none !important;
}

button, .gr-button {
    border-radius: 40px !important;
    font-weight: 600 !important;
    padding: 12px 24px !important;
    transition: all 0.2s ease !important;
    border: none !important;
}

button:hover, .gr-button:hover {
    transform: translateY(-2px) !important;
    filter: brightness(0.95) !important;
}

.gr-button-primary {
    background: linear-gradient(135deg, var(--primary), var(--secondary)) !important;
    color: white !important;
}

/* Profile Section */
.profile-header {
    background: linear-gradient(135deg, var(--primary), var(--secondary));
    border-radius: 24px;
    padding: 30px;
    text-align: center;
    margin-bottom: 20px;
}

.profile-avatar-img, .profile-avatar-placeholder {
    width: 100px;
    height: 100px;
    border-radius: 50%;
    margin: 0 auto 16px;
    object-fit: cover;
    border: 3px solid white;
}

.profile-avatar-placeholder {
    background: rgba(255,255,255,0.2);
    display: flex;
    align-items: center;
    justify-content: center;
    font-size: 48px;
    margin: 0 auto 16px;
}

.profile-username {
    color: rgba(255,255,255,0.9);
    font-size: 14px;
    margin: 5px 0;
}

.profile-bio {
    margin: 15px 0;
    font-size: 14px;
}

.profile-joined {
    font-size: 12px;
    opacity: 0.8;
}

.profile-stats {
    display: flex;
    justify-content: space-around;
    margin-top: 20px;
    padding-top: 20px;
    border-top: 1px solid rgba(255,255,255,0.2);
}

.stat-number {
    font-size: 22px;
    font-weight: 700;
}

.stat-label {
    font-size: 11px;
    opacity: 0.8;
}

/* Gallery */
.gallery {
    border-radius: 16px !important;
    gap: 8px !important;
}

/* Notifications */
.notifications-list {
    padding: 20px;
}

.notification-item {
    background: var(--card);
    border-radius: 12px;
    padding: 15px;
    margin-bottom: 10px;
    border-left: 3px solid var(--secondary);
}

.notification-item.unread {
    background: #2A1F1F;
}

.notification-time {
    font-size: 11px;
    color: var(--text-secondary);
    margin-top: 5px;
}

/* Empty State */
.empty-state {
    text-align: center;
    padding: 80px 20px;
    color: var(--text-secondary);
}

.empty-icon {
    font-size: 64px;
    margin-bottom: 20px;
}
"""

# ===============================
# MAIN UI WITH STICKY HEADER
# ===============================

def get_header_html(username=None):
    user_section = f'<span style="color: #F59E0B;">👤 {username}</span>' if username and username != "guest" else '<span style="color: #888;">👻 Guest Mode</span>'
    
    return f"""
    <div class="app-header">
        <div style="display: flex; justify-content: space-between; align-items: center;">
            <h1>✨ Vibes</h1>
            <div style="font-size: 14px;">
                {user_section}
            </div>
        </div>
    </div>
    """

with gr.Blocks(title="Vibes", css=custom_css, theme=gr.themes.Soft()) as app:
    
    user_state = gr.State("guest")
    
    # Login Screen
    with gr.Column(visible=True, elem_id="login_screen") as login_screen:
        gr.HTML("""
        <div style="text-align: center; padding: 60px 20px;">
            <div style="font-size: 80px; margin-bottom: 20px;">✨</div>
            <h1 style="font-size: 48px; background: linear-gradient(135deg, #8B5CF6, #F59E0B); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">Vibes</h1>
            <p style="color: #888; margin-bottom: 40px;">Share moments. Connect with creators.</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                pass
            with gr.Column(scale=2):
                with gr.Group():
                    with gr.Tabs():
                        with gr.TabItem("🔐 Login"):
                            login_user = gr.Textbox(label="Username", placeholder="Enter username")
                            login_pass = gr.Textbox(label="Password", type="password", placeholder="Enter password")
                            login_btn = gr.Button("Login", variant="primary")
                            login_msg = gr.Markdown("")
                        
                        with gr.TabItem("📝 Sign Up"):
                            su_user = gr.Textbox(label="Username", placeholder="Choose username")
                            su_pass = gr.Textbox(label="Password", type="password", placeholder="Choose password")
                            su_display = gr.Textbox(label="Display Name", placeholder="What should we call you?")
                            su_btn = gr.Button("Create Account", variant="primary")
                            su_msg = gr.Markdown("")
                    
                    gr.Markdown("---")
                    guest_btn = gr.Button("👻 Browse without login", variant="secondary", size="lg")
            with gr.Column(scale=1):
                pass
    
    # Main App
    with gr.Column(visible=False, elem_id="main_app") as main_app:
        
        # Dynamic header
        header_html = gr.HTML(get_header_html())
        
        with gr.Tabs():
            with gr.TabItem("🎵 Home", id="feed_tab"):
                feed_html = gr.HTML(value=render_feed())
                refresh_btn = gr.Button("🔄 Refresh", variant="secondary", size="sm")
            
            with gr.TabItem("📸 Create", id="create_tab"):
                with gr.Column():
                    gr.Markdown("### Create New Post")
                    media_type = gr.Radio(choices=["Image", "Video"], label="Media Type", value="Image")
                    image_input = gr.Image(type="numpy", label="Upload Image", visible=True, height=350)
                    video_input = gr.Video(label="Upload Video", visible=False, height=350)
                    caption_input = gr.Textbox(label="Caption", placeholder="What's on your mind?", lines=2)
                    upload_btn = gr.Button("🚀 Post", variant="primary")
                    upload_msg = gr.Markdown("")
            
            with gr.TabItem("👤 Profile", id="profile_tab"):
                profile_display = gr.Markdown("")
                bio_input = gr.Textbox(label="Update Bio", placeholder="Write something about yourself...", lines=2)
                bio_btn = gr.Button("Update Bio", variant="secondary")
                bio_msg = gr.Markdown("")
                gr.Markdown("### My Posts")
                profile_gallery = gr.Gallery(columns=3, height=400, object_fit="cover")
                with gr.Row():
                    refresh_profile_btn = gr.Button("🔄 Refresh", variant="secondary")
                    logout_btn = gr.Button("🚪 Logout", variant="stop")
            
            with gr.TabItem("🔍 Discover", id="discover_tab"):
                search_input = gr.Textbox(label="Search creators", placeholder="Enter username...")
                search_results = gr.Dataframe(
                    headers=["Username", "Display Name", "Bio"],
                    datatype=["str", "str", "str"],
                    interactive=False,
                    height=300
                )
                gr.Markdown("---")
                follow_input = gr.Textbox(label="Follow user", placeholder="Enter username to follow")
                follow_btn = gr.Button("➕ Follow/Unfollow", variant="secondary")
                follow_msg = gr.Markdown("")
            
            with gr.TabItem("🔔 Notifications", id="notifications_tab"):
                notifications_html = gr.HTML()
                refresh_notifs_btn = gr.Button("🔄 Refresh", variant="secondary")
    
    # ===============================
    # UI UPDATE FUNCTIONS
    # ===============================
    
    def update_header(user):
        return get_header_html(user)
    
    def toggle_media_type(media_type):
        return gr.update(visible=(media_type == "Image")), gr.update(visible=(media_type == "Video"))
    
    def do_login(user, pwd):
        msg, logged_user = login(user, pwd)
        if logged_user:
            return (msg, gr.update(visible=False), gr.update(visible=True), logged_user, 
                    render_feed(), render_notifications(logged_user), get_header_html(logged_user))
        return msg, gr.update(), gr.update(), "guest", render_feed(), "", get_header_html("guest")
    
    def do_signup(user, pwd, disp):
        msg, _, _, _ = signup(user, pwd, disp)
        return msg, "", "", ""
    
    def do_guest():
        return (gr.update(visible=False), gr.update(visible=True), "guest", 
                render_feed(), "", get_header_html("guest"))
    
    def do_logout():
        return (gr.update(visible=True), gr.update(visible=False), "guest", 
                render_feed(), "", get_header_html("guest"))
    
    def do_upload(user, media_type, img, vid, caption):
        if user == "guest":
            return "❌ Login to post"
        if media_type == "Image":
            return upload_media(user, img, caption, "image")
        else:
            return upload_media(user, vid, caption, "video")
    
    def do_load_profile(user):
        return load_profile(user)
    
    def do_update_bio(user, bio):
        return update_bio(user, bio)
    
    def do_follow(user, target):
        return follow_user(user, target)
    
    def do_refresh_notifications(user):
        return render_notifications(user)
    
    def do_search(query):
        return search_users(query)
    
    # Wire Events
    login_btn.click(do_login, [login_user, login_pass], 
                   [login_msg, login_screen, main_app, user_state, feed_html, notifications_html, header_html])
    su_btn.click(do_signup, [su_user, su_pass, su_display], [su_msg, su_user, su_pass, su_display])
    guest_btn.click(do_guest, None, [login_screen, main_app, user_state, feed_html, notifications_html, header_html])
    
    refresh_btn.click(lambda: render_feed(), None, feed_html)
    
    media_type.change(toggle_media_type, [media_type], [image_input, video_input])
    
    upload_btn.click(do_upload, [user_state, media_type, image_input, video_input, caption_input], upload_msg)
    upload_btn.click(lambda: render_feed(), None, feed_html)
    upload_btn.click(lambda: (None, None, "", "", None, None), None, [image_input, video_input, caption_input, upload_msg, media_type])
    
    refresh_profile_btn.click(lambda u: load_profile(u), [user_state], [profile_display, profile_gallery])
    bio_btn.click(do_update_bio, [user_state, bio_input], bio_msg)
    bio_btn.click(lambda u: load_profile(u), [user_state], [profile_display, profile_gallery])
    logout_btn.click(do_logout, None, [login_screen, main_app, user_state, feed_html, notifications_html, header_html])
    
    follow_btn.click(do_follow, [user_state, follow_input], follow_msg)
    search_input.change(do_search, [search_input], [search_results])
    
    refresh_notifs_btn.click(lambda u: render_notifications(u), [user_state], notifications_html)

app.launch(server_name="0.0.0.0", server_port=7860)