app.py

import gradio as gr
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import base64
import os
import json
import random
import urllib.parse
import time

# Initialize the Google Generative AI client with the API key from environment variables
try:
    api_key = os.environ['GEMINI_API_KEY']
except KeyError:
    raise ValueError("Please set the GEMINI_API_KEY environment variable.")
client = genai.Client(api_key=api_key)

# Define safety settings to disable all filters for content generation
SAFETY_SETTINGS = [
    types.SafetySetting(
        category=types.HarmCategory.HARM_CATEGORY_HARASSMENT,
        threshold=types.HarmBlockThreshold.BLOCK_NONE,
    ),
    types.SafetySetting(
        category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
        threshold=types.HarmBlockThreshold.BLOCK_NONE,
    ),
    types.SafetySetting(
        category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
        threshold=types.HarmBlockThreshold.BLOCK_NONE,
    ),
    types.SafetySetting(
        category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
        threshold=types.HarmBlockThreshold.BLOCK_NONE,
    ),
    types.SafetySetting(
        category=types.HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY,
        threshold=types.HarmBlockThreshold.BLOCK_NONE,
    ),
]

# Define progress stages and corresponding witty messages
PROGRESS_STAGES = {
    "initializing": [
        "Warming up the creativity engine... 🚀",
        "Getting ready to spark some magic... ✨",
        "Loading the idea generator... ⚡"
    ],
    "generating_ideas": [
        "Brainstorming brilliant ideas... 💡",
        "Cooking up some viral concepts... 🍳",
        "Mixing creativity with a pinch of flair... 🎨"
    ],
    "creating_image": [
        "Painting a stunning visual... 🖌️",
        "Snapping the perfect shot... 📸",
        "Crafting a masterpiece pixel by pixel... 🖼️"
    ],
    "retrying_image": [
        "Tweaking the image for perfection... 🔧",
        "Giving it another go for that wow factor... 🔄",
        "Adjusting the lens for a better shot... 🎥"
    ],
    "generating_video": [
        "Rolling the cameras for a viral video... 🎬",
        "Animating the scene with cinematic vibes... 🌟",
        "Producing a TikTok-worthy clip... 🎞️"
    ],
    "finalizing": [
        "Putting the final touches... 🎀",
        "Polishing your content to shine... 💎",
        "Almost ready to blow up the feed... 🔥"
    ]
}

def clean_response_text(response_text):
    """
    Clean the API response by removing Markdown code block markers.
    """
    cleaned_text = response_text.strip()
    if cleaned_text.startswith("```json"):
        cleaned_text = cleaned_text[len("```json"):].strip()
    if cleaned_text.endswith("```"):
        cleaned_text = cleaned_text[:-len("```")].strip()
    return cleaned_text

def generate_ideas(user_input):
    """
    Generate a diverse set of ideas based on the user's input concept using the LLM.
    """
    prompt = f"""
    The user has provided the concept: "{user_input}". You must generate 5 diverse and creative ideas for a TikTok video that are directly and explicitly related to "{user_input}".
    Each idea must clearly incorporate and focus on the core theme of "{user_input}" without deviating into unrelated topics.
    Each idea should be a short sentence describing a specific scene or concept.
    Return the response as a JSON object with a single key 'ideas' containing a list of 5 ideas.
    Ensure the response is strictly in JSON format.
    """
    try:
        response = client.models.generate_content(
            model='gemini-2.0-flash-lite',
            contents=[prompt],
            config=types.GenerateContentConfig(
                temperature=1.2,
                safety_settings=SAFETY_SETTINGS
            )
        )
        if not response.text or response.text.isspace():
            raise ValueError("Empty response from API")
        cleaned_text = clean_response_text(response.text)
        response_json = json.loads(cleaned_text)
        if 'ideas' not in response_json or not isinstance(response_json['ideas'], list) or len(response_json['ideas']) != 5:
            raise ValueError("Invalid JSON format: 'ideas' key missing, not a list, or incorrect length")
        return response_json['ideas']
    except Exception as e:
        print(f"Error generating ideas: {e}")
        return [
            f"A dramatic {user_input} scene with cinematic lighting",
            f"A close-up of {user_input} in a futuristic setting",
            f"A high-energy {user_input} moment with vibrant colors",
            f"A serene {user_input} scene with soft focus",
            f"An action-packed {user_input} challenge with dynamic angles"
        ]

def generate_item(user_input, ideas, generate_video=False, max_retries=3):
    """
    Generate a single feed item with progress updates.
    Yields progress stage and message for UI updates.
    """
    video_base64 = None
    max_total_attempts = 3
    
    total_attempts = 0
    while total_attempts < max_total_attempts:
        total_attempts += 1
        
        yield {"stage": "initializing", "message": random.choice(PROGRESS_STAGES["initializing"]), "progress": 10}
        
        generated_image = None
        text = None
        img_str = None
        image_prompt = None
        
        for image_attempt in range(max_retries):
            yield {"stage": "creating_image", "message": random.choice(PROGRESS_STAGES["creating_image"]), "progress": 30 + (image_attempt * 10)}
            
            selected_idea = random.choice(ideas)
            prompt = f"""
            The user has provided the concept: "{user_input}". Based on this concept and the specific idea "{selected_idea}", create content for a TikTok video.
            Return a JSON object with two keys:
            - 'caption': A short, viral TikTok-style caption with hashtags that reflects "{user_input}".
            - 'image_prompt': A detailed image prompt for generating a high-quality visual scene, ensuring the theme of "{user_input}" is central.
            The image prompt should describe the scene vividly, specify a perspective and style, and ensure no text or letters are included.
            Ensure the response is strictly in JSON format.
            """
            try:
                response = client.models.generate_content(
                    model='gemini-2.0-flash-lite',
                    contents=[prompt],
                    config=types.GenerateContentConfig(
                        temperature=1.2,
                        safety_settings=SAFETY_SETTINGS
                    )
                )
                cleaned_text = clean_response_text(response.text)
                response_json = json.loads(cleaned_text)
                text = response_json['caption']
                image_prompt = response_json['image_prompt']
            except Exception as e:
                print(f"Error generating item: {e}")
                text = f"Amazing {user_input}! 🔥 #{user_input.replace(' ', '')}"
                image_prompt = f"A vivid scene of {selected_idea} related to {user_input}, in a vibrant pop art style, no text or letters"

            try:
                imagen = client.models.generate_images(
                    model='imagen-3.0-generate-002',
                    prompt=image_prompt,
                    config=types.GenerateImagesConfig(
                        aspect_ratio="9:16",
                        number_of_images=1
                    )
                )
                if imagen.generated_images and len(imagen.generated_images) > 0:
                    generated_image = imagen.generated_images[0]
                    image = Image.open(BytesIO(generated_image.image.image_bytes))
                    target_width = 360
                    target_height = int(target_width / 9 * 16)
                    image = image.resize((target_width, target_height), Image.LANCZOS)
                    buffered = BytesIO()
                    image.save(buffered, format="PNG")
                    img_str = base64.b64encode(buffered.getvalue()).decode()
                    break
                else:
                    if image_attempt == max_retries - 1:
                        if total_attempts == max_total_attempts:
                            image = Image.new('RGB', (360, 640), color='gray')
                            buffered = BytesIO()
                            image.save(buffered, format="PNG")
                            img_str = base64.b64encode(buffered.getvalue()).decode()
                            yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
                            return {
                                'text': text,
                                'image_base64': img_str,
                                'video_base64': None,
                                'ideas': ideas
                            }
                        break
                    yield {"stage": "retrying_image", "message": random.choice(PROGRESS_STAGES["retrying_image"]), "progress": 40 + (image_attempt * 10)}
            except Exception as e:
                print(f"Error generating image: {e}")
                if image_attempt == max_retries - 1:
                    if total_attempts == max_total_attempts:
                        image = Image.new('RGB', (360, 640), color='gray')
                        buffered = BytesIO()
                        image.save(buffered, format="PNG")
                        img_str = base64.b64encode(buffered.getvalue()).decode()
                        yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
                        return {
                            'text': text,
                            'image_base64': img_str,
                            'video_base64': None,
                            'ideas': ideas
                        }
                    break
                yield {"stage": "retrying_image", "message": random.choice(PROGRESS_STAGES["retrying_image"]), "progress": 40 + (image_attempt * 10)}

        if generate_video and generated_image is not None:
            yield {"stage": "generating_video", "message": random.choice(PROGRESS_STAGES["generating_video"]), "progress": 70}
            try:
                video_prompt = f"""
                The user concept is "{user_input}". Based on this and the scene: {image_prompt}, create a video.
                Use a close-up shot with a slow dolly shot circling around the subject, 
                using shallow focus on the main subject to emphasize details, in a realistic style with cinematic lighting.
                """
                operation = client.models.generate_videos(
                    model="veo-2.0-generate-001",
                    prompt=video_prompt,
                    image=generated_image.image,
                    config=types.GenerateVideosConfig(
                        aspect_ratio="9:16",
                        number_of_videos=1,
                        duration_seconds=8,
                        negative_prompt="blurry, low quality, text, letters"
                    )
                )
                while not operation.done:
                    time.sleep(20)
                    operation = client.operations.get(operation)
                if operation.error:
                    raise ValueError(f"Video generation failed: {operation.error.message}")
                if operation.response and hasattr(operation.response, 'generated_videos') and operation.response.generated_videos:
                    if len(operation.response.generated_videos) > 0:
                        video = operation.response.generated_videos[0]
                        if video and hasattr(video, 'video'):
                            video_data = client.files.download(file=video.video)
                            video_bytes = video_data if isinstance(video_data, bytes) else BytesIO(video_data).getvalue()
                            video_base64 = base64.b64encode(video_bytes).decode()
                            yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
                            return {
                                'text': text,
                                'image_base64': img_str,
                                'video_base64': video_base64,
                                'ideas': ideas
                            }
                    raise ValueError("No valid video generated")
                else:
                    raise ValueError("Video generation operation failed: No generated_videos in response")
            except Exception as e:
                print(f"Error generating video: {e}")
                yield {"stage": "generating_video", "message": random.choice(PROGRESS_STAGES["generating_video"]), "progress": 80}
                try:
                    operation = client.models.generate_videos(
                        model="veo-2.0-generate-001",
                        prompt=video_prompt,
                        config=types.GenerateVideosConfig(
                            aspect_ratio="9:16",
                            number_of_videos=1,
                            duration_seconds=8,
                            negative_prompt="blurry, low quality, text, letters"
                        )
                    )
                    while not operation.done:
                        time.sleep(20)
                        operation = client.operations.get(operation)
                    if operation.error:
                        raise ValueError(f"Video generation failed: {operation.error.message}")
                    if operation.response and hasattr(operation.response, 'generated_videos') and operation.response.generated_videos:
                        if len(operation.response.generated_videos) > 0:
                            video = operation.response.generated_videos[0]
                            if video and hasattr(video, 'video'):
                                video_data = client.files.download(file=video.video)
                                video_bytes = video_data if isinstance(video_data, bytes) else BytesIO(video_data).getvalue()
                                video_base64 = base64.b64encode(video_bytes).decode()
                                yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
                                return {
                                    'text': text,
                                    'image_base64': img_str,
                                    'video_base64': video_base64,
                                    'ideas': ideas
                                }
                        raise ValueError("No valid video generated")
                    else:
                        raise ValueError("Video generation operation failed: No generated_videos in response")
                except Exception as e:
                    print(f"Error generating text-to-video: {e}")
                    if total_attempts == max_total_attempts:
                        yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
                        return {
                            'text': text,
                            'image_base64': img_str,
                            'video_base64': None,
                            'ideas': ideas
                        }
        
        if img_str is not None:
            yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
            return {
                'text': text,
                'image_base64': img_str,
                'video_base64': video_base64,
                'ideas': ideas
            }

    image = Image.new('RGB', (360, 640), color='gray')
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
    return {
        'text': f"Amazing {user_input}! 🔥 #{user_input.replace(' ', '')}",
        'image_base64': img_str,
        'video_base64': None,
        'ideas': ideas
    }

def start_feed(user_input, generate_video, current_index, feed_items):
    """
    Start or update the feed based on the user input with progress updates.
    """
    if not user_input.strip():
        user_input = "trending"
    current_user_input = user_input
    is_loading = True
    share_links = ""
    
    html_content = generate_html([], False, 0, user_input, is_loading, "initializing", random.choice(PROGRESS_STAGES["initializing"]), 10)
    yield current_user_input, current_index, feed_items, html_content, share_links, is_loading

    try:
        html_content = generate_html([], False, 0, user_input, is_loading, "generating_ideas", random.choice(PROGRESS_STAGES["generating_ideas"]), 20)
        yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
        ideas = generate_ideas(user_input)
        
        item_generator = generate_item(user_input, ideas, generate_video=generate_video)
        for progress in item_generator:
            if isinstance(progress, dict) and "stage" in progress:
                html_content = generate_html([], False, 0, user_input, is_loading, progress["stage"], progress["message"], progress["progress"])
                yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
            else:
                item = progress
                feed_items = [item]
                current_index = 0
                share_links = generate_share_links(item['image_base64'], item['video_base64'], item['text'])
                is_loading = False
                html_content = generate_html(feed_items, False, current_index, user_input, is_loading)
                yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                return
    except Exception as e:
        print(f"Error in start_feed: {e}")
        html_content = """
        <div style="
            display: flex;
            flex-direction: column;
            align-items: center;
            justify-content: center;
            max-width: 360px;
            margin: 0 auto;
            background-color: #000;
            height: 640px;
            border: 1px solid #333;
            border-radius: 10px;
            color: white;
            font-family: Arial, sans-serif;
        ">
            <p>Error generating content. Please try again!</p>
        </div>
        """
        is_loading = False
        yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
        return

def load_next(user_input, generate_video, current_index, feed_items):
    """
    Load the next item in the feed with progress updates.
    """
    current_user_input = user_input if user_input.strip() else "trending"
    user_input = current_user_input
    is_loading = True
    share_links = ""
    
    html_content = generate_html(feed_items, False, current_index, user_input, is_loading, "initializing", random.choice(PROGRESS_STAGES["initializing"]), 10)
    yield current_user_input, current_index, feed_items, html_content, share_links, is_loading

    try:
        if current_index + 1 < len(feed_items):
            current_index += 1
            is_loading = False
            html_content = generate_html(feed_items, False, current_index, user_input, is_loading)
            share_links = generate_share_links(
                feed_items[current_index]['image_base64'],
                feed_items[current_index]['video_base64'],
                feed_items[current_index]['text']
            )
            yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
            return
        else:
            ideas = feed_items[-1]['ideas'] if feed_items else generate_ideas(user_input)
            html_content = generate_html(feed_items, False, current_index, user_input, is_loading, "generating_ideas", random.choice(PROGRESS_STAGES["generating_ideas"]), 20)
            yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
            
            item_generator = generate_item(user_input, ideas, generate_video=generate_video)
            for progress in item_generator:
                if isinstance(progress, dict) and "stage" in progress:
                    html_content = generate_html(feed_items, False, current_index, user_input, is_loading, progress["stage"], progress["message"], progress["progress"])
                    yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                else:
                    new_item = progress
                    feed_items.append(new_item)
                    current_index = len(feed_items) - 1
                    is_loading = False
                    html_content = generate_html(feed_items, False, current_index, user_input, is_loading)
                    share_links = generate_share_links(
                        feed_items[current_index]['image_base64'],
                        feed_items[current_index]['video_base64'],
                        feed_items[current_index]['text']
                    )
                    yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                    return
    except Exception as e:
        print(f"Error in load_next: {e}")
        html_content = """
        <div style="
            display: flex;
            flex-direction: column;
            align-items: center;
            justify-content: center;
            max-width: 360px;
            margin: 0 auto;
            background-color: #000;
            height: 640px;
            border: 1px solid #333;
            border-radius: 10px;
            color: white;
            font-family: Arial, sans-serif;
        ">
            <p>Error generating content. Please try again!</p>
        </div>
        """
        is_loading = False
        yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
        return

def load_previous(user_input, generate_video, current_index, feed_items):
    """
    Load the previous item in the feed.
    """
    current_user_input = user_input if user_input.strip() else "trending"
    
    if current_index > 0:
        current_index -= 1
    html_content = generate_html(feed_items, False, current_index, user_input, False)
    share_links = generate_share_links(
        feed_items[current_index]['image_base64'],
        feed_items[current_index]['video_base64'],
        feed_items[current_index]['text']
    )
    return current_user_input, current_index, feed_items, html_content, share_links, False

def generate_share_links(image_base64, video_base64, caption):
    """
    Generate share links for social media platforms with download links for image and video.
    """
    image_data_url = f"data:image/png;base64,{image_base64}"
    encoded_caption = urllib.parse.quote(caption)

    download_links = f"""
        <p style="text-align: center; margin-bottom: 10px;">Download the media to share:</p>
        <div style="display: flex; flex-wrap: wrap; justify-content: center; gap: 8px; margin-bottom: 15px;">
            <a href="{image_data_url}" download="feed_item.png" style="
                background-color: #4CAF50;
                color: white;
                padding: 8px 16px;
                border-radius: 5px;
                text-decoration: none;
                font-size: 14px;
                font-weight: bold;
                transition: background-color 0.3s;
            " onmouseover="this.style.backgroundColor='#45a049'" onmouseout="this.style.backgroundColor='#4CAF50'">Download Image</a>
    """
    if video_base64:
        video_data_url = f"data:video/mp4;base64,{video_base64}"
        download_links += f"""
            <a href="{video_data_url}" download="feed_video.mp4" style="
                background-color: #4CAF50;
                color: white;
                padding: 8px 16px;
                border-radius: 5px;
                text-decoration: none;
                font-size: 14px;
                font-weight: bold;
                transition: background-color 0.3s;
            " onmouseover="this.style.backgroundColor='#45a049'" onmouseout="this.style.backgroundColor='#4CAF50'">Download Video</a>
        """
    download_links += "</div>"

    instruction = """
        <p style="text-align: center; margin-bottom: 10px;">
            Click a share button below to start a post with the caption, then manually upload the downloaded image or video.
        </p>
    """

    share_links = """
        <div style="
            display: flex;
            flex-wrap: wrap;
            justify-content: center;
            gap: 8px;
            margin-bottom: 15px;
        ">
            <a href="https://www.tiktok.com/upload?caption={caption}" target="_blank" style="
                background-color: #00f2ea;
                color: #000;
                padding: 8px 16px;
                border-radius: 5px;
                text-decoration: none;
                font-size: 14px;
                font-weight: bold;
                transition: background-color 0.3s;
            " onmouseover="this.style.backgroundColor='#00d9d1'" onmouseout="this.style.backgroundColor='#00f2ea'">Share on TikTok</a>
            <a href="https://www.instagram.com/?caption={caption}" target="_blank" style="
                background-color: #e1306c;
                color: white;
                padding: 8px 16px;
                border-radius: 5px;
                text-decoration: none;
                font-size: 14px;
                font-weight: bold;
                transition: background-color 0.3s;
            " onmouseover="this.style.backgroundColor='#c72b5e'" onmouseout="this.style.backgroundColor='#e1306c'">Share on Instagram</a>
            <a href="https://www.facebook.com/sharer/sharer.php?quote={caption}" target="_blank" style="
                background-color: #4267b2;
                color: white;
                padding: 8px 16px;
                border-radius: 5px;
                text-decoration: none;
                font-size: 14px;
                font-weight: bold;
                transition: background-color 0.3s;
            " onmouseover="this.style.backgroundColor='#395a9d'" onmouseout="this.style.backgroundColor='#4267b2'">Share on Facebook</a>
            <a href="https://twitter.com/intent/tweet?text={caption}" target="_blank" style="
                background-color: #1da1f2;
                color: white;
                padding: 8px 16px;
                border-radius: 5px;
                text-decoration: none;
                font-size: 14px;
                font-weight: bold;
                transition: background-color 0.3s;
            " onmouseover="this.style.backgroundColor='#1a91da'" onmouseout="this.style.backgroundColor='#1da1f2'">Share on X</a>
            <a href="https://pinterest.com/pin/create/button/?description={caption}" target="_blank" style="
                background-color: #bd081c;
                color: white;
                padding: 8px 16px;
                border-radius: 5px;
                text-decoration: none;
                font-size: 14px;
                font-weight: bold;
                transition: background-color 0.3s;
            " onmouseover="this.style.backgroundColor='#a30718'" onmouseout="this.style.backgroundColor='#bd081c'">Share on Pinterest</a>
        </div>
    """

    youtube_share = ""
    if video_base64:
        youtube_share = f"""
            <div style="
                display: flex;
                justify-content: center;
                margin-top: 10px;
            ">
                <a href="https://www.youtube.com/upload?description={caption}" target="_blank" style="
                    background-color: #ff0000;
                    color: white;
                    padding: 8px 16px;
                    border-radius: 5px;
                    text-decoration: none;
                    font-size: 14px;
                    font-weight: bold;
                    transition: background-color 0.3s;
                " onmouseover="this.style.backgroundColor='#e60000'" onmouseout="this.style.backgroundColor='#ff0000'">Share to YouTube as a Short</a>
            </div>
        """

    return f"""
    <div style="
        display: flex;
        flex-direction: column;
        align-items: center;
        gap: 10px;
        margin-top: 10px;
        color: white;
        font-family: Arial, sans-serif;
    ">
        {download_links}
        {instruction}
        {share_links}
        {youtube_share}
    </div>
    """.format(caption=encoded_caption)

def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_input="", is_loading=False, progress_stage="initializing", progress_message="", progress_percent=10):
    """
    Generate an HTML string to display the current feed item or a loading screen with a progress bar.
    """
    if is_loading:
        progress_percent = max(0, min(100, progress_percent))
        return f"""
        <div id="feed-container" style="
            display: flex;
            flex-direction: column;
            align-items: center;
            justify-content: center;
            max-width: 360px;
            margin: 0 auto;
            background-color: #000;
            height: 640px;
            border: 1px solid #333;
            border-radius: 10px;
            color: white;
            font-family: Arial, sans-serif;
            position: relative;
        ">
            <div id="loading-message" style="
                font-size: 18px;
                font-weight: bold;
                text-align: center;
                margin-bottom: 20px;
                text-shadow: 1px 1px 2px rgba(0,0,0,0.5);
                max-width: 90%;
            ">
                {progress_message}
            </div>
            <div style="
                width: 80%;
                height: 10px;
                background-color: #333;
                border-radius: 5px;
                overflow: hidden;
            ">
                <div style="
                    width: {progress_percent}%;
                    height: 100%;
                    background: linear-gradient(to right, #ff2d55, #ff5e78);
                    transition: width 0.5s ease-in-out;
                "></div>
            </div>
            <style>
                @keyframes pulse {{
                    0% {{ opacity: 1; }}
                    50% {{ opacity: 0.7; }}
                    100% {{ opacity: 1; }}
                }}
                #loading-message {{
                    animation: pulse 2s infinite;
                }}
            </style>
            <script>
                const stages = {json.dumps(PROGRESS_STAGES)};
                const currentStage = "{progress_stage}";
                let currentMessageIndex = 0;
                const messageElement = document.getElementById('loading-message');
                function rotateMessages() {{
                    if (stages[currentStage] && stages[currentStage].length > 0) {{
                        currentMessageIndex = (currentMessageIndex + 1) % stages[currentStage].length;
                        messageElement.textContent = stages[currentStage][currentMessageIndex];
                    }}
                }}
                setInterval(rotateMessages, 2000);
            </script>
        </div>
        """

    if not feed_items or current_index >= len(feed_items):
        return """
        <div style="
            display: flex;
            flex-direction: column;
            align-items: center;
            justify-content: center;
            max-width: 360px;
            margin: 0 auto;
            background-color: #000;
            height: 640px;
            border: 1px solid #333;
            border-radius: 10px;
            color: white;
            font-family: Arial, sans-serif;
        ">
            <p>Enter a concept or idea to start your feed!</p>
        </div>
        """

    item = feed_items[current_index]
    media_element = f"""
        <img id="feed-image" src="data:image/png;base64,{item['image_base64']}" style="
            width: 100%;
            height: 100%;
            object-fit: cover;
            position: absolute;
            top: 0;
            left: 0;
            z-index: 1;
        ">
    """ if not item['video_base64'] else f"""
        <video id="feed-video" controls style="
            width: 100%;
            height: 100%;
            object-fit: cover;
            position: absolute;
            top: 0;
            left: 0;
            z-index: 1;
        ">
            <source src="data:video/mp4;base64,{item['video_base64']}" type="video/mp4">
            Your browser does not support the video tag.
        </video>
    """

    html_str = f"""
    <div id="feed-container" style="
        display: flex;
        flex-direction: column;
        align-items: center;
        max-width: 360px;
        margin: 0 auto;
        background-color: #000;
        height: 640px;
        border: 1px solid #333;
        border-radius: 10px;
        position: relative;
    ">
        <div class="feed-item" style="
            width: 100%;
            height: 100%;
            position: relative;
            display: flex;
            flex-direction: column;
            justify-content: flex-end;
            overflow: hidden;
            cursor: pointer;
        " onclick="handleClick(event)">
            {media_element}
            <div style="
                position: relative;
                z-index: 2;
                background: linear-gradient(to top, rgba(0,0,0,0.7), transparent);
                padding: 20px;
                color: white;
                font-family: Arial, sans-serif;
                font-size: 18px;
                font-weight: bold;
                text-shadow: 1px 1px 2px rgba(0,0,0,0.5);
            ">
                {item['text']}
            </div>
        </div>
    </div>
    <script>
        function handleClick(event) {{
            const media = document.getElementById('feed-video') || document.getElementById('feed-image');
            const rect = media.getBoundingClientRect();
            const clickX = event.clientX - rect.left;
            const width = rect.width;
            if (clickX > width * 0.75) {{
                document.getElementById('previous-button').click();
            }}
        }}
    </script>
    """
    return html_str

# Define the Gradio interface
with gr.Blocks(
    css="""
        body { background-color: #000; color: #fff; font-family: Arial, sans-serif; }
        .gradio-container { max-width: 400px; margin: 0 auto; padding: 10px; }
        input, select, button, .gr-checkbox { border-radius: 5px; background-color: #222; color: #fff; border: 1px solid #444; }
        button { background-color: #ff2d55; border: none; }
        button:hover { background-color: #e0264b; }
        .gr-button { width: 100%; margin-top: 10px; }
        .gr-form { background-color: #111; padding: 15px; border-radius: 10px; }
    """,
    title="Create Your Feed"
) as demo:
    current_user_input = gr.State(value="")
    current_index = gr.State(value=0)
    feed_items = gr.State(value=[])
    is_loading = gr.State(value=False)
    share_links = gr.State(value="")

    with gr.Column(elem_classes="gr-form"):
        gr.Markdown("### Create Your Feed")
        user_input = gr.Textbox(
            label="Enter Concept or Ideas",
            value="",
            placeholder="e.g., sushi adventure, neon tech",
            submit_btn=False
        )
        generate_video_checkbox = gr.Checkbox(
            label="Generate Video (may take longer)",
            value=False
        )
        magic_button = gr.Button("✨ Create ✨", elem_classes="gr-button")

    feed_html = gr.HTML()
    share_html = gr.HTML(label="Share this item:")

    user_input.submit(
        fn=start_feed,
        inputs=[user_input, generate_video_checkbox, current_index, feed_items],
        outputs=[current_user_input, current_index, feed_items, feed_html, share_html, is_loading]
    )

    magic_button.click(
        fn=load_next,
        inputs=[user_input, generate_video_checkbox, current_index, feed_items],
        outputs=[current_user_input, current_index, feed_items, feed_html, share_html, is_loading]
    )

    previous_button = gr.Button("Previous", elem_id="previous-button", visible=False)
    previous_button.click(
        fn=load_previous,
        inputs=[user_input, generate_video_checkbox, current_index, feed_items],
        outputs=[current_user_input, current_index, feed_items, feed_html, share_html, is_loading]
    )

demo.launch()