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
import gradio

# Check Gradio version
required_version = "4.44.0"
current_version = gradio.__version__
if current_version < required_version:
    raise ValueError(f"Gradio version {current_version} is outdated. Please upgrade to {required_version} or later using 'pip install gradio=={required_version}'.")

# 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,
    ),
]

def clean_response_text(response_text):
    """
    Clean the API response by removing Markdown code block markers.
    
    Args:
        response_text (str): The raw response text from the API.
    
    Returns:
        str: The cleaned response text.
    """
    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.
    Yields progress updates for the loading UI.
    
    Args:
        user_input (str): The user's input concept or idea.
    
    Yields:
        tuple: (progress_percentage, message) for progress_html updates.
        list: Final list of ideas as strings.
    """
    yield (10, f"Brainstorming epic ideas for {user_input}... 🌟")
    
    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.
    Example for "blindfolded Rubik's Cube challenge":
    {{"ideas": [
        "A blindfolded speedcubing competition with dramatic music",
        "A close-up of a person solving a Rubik's Cube blindfolded under a spotlight",
        "A time-lapse of a blindfolded Rubik's Cube solve with colorful lighting",
        "A blindfolded Rubik's Cube challenge in a futuristic setting",
        "A split-screen of two people racing to solve a Rubik's Cube blindfolded"
    ]}}
    """
    try:
        response = client.models.generate_content(
            model='gemini-2.0-flash-lite',
            contents=[prompt],
            config=types.GenerateContentConfig(
                temperature=1.2,
                safety_settings=SAFETY_SETTINGS
            )
        )
        print(f"Raw response for ideas: {response.text}")  # Debugging
        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")
        
        ideas = response_json['ideas']
        yield (20, f"Ideas locked in for {user_input}! 🚀")
        return ideas
    
    except Exception as e:
        print(f"Error generating ideas: {e}")
        yield (20, f"Oops, tweaking the plan for {user_input}... 🔧")
        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 (image and optionally one video) using one of the ideas.
    Yields progress updates for the loading UI.
    
    Args:
        user_input (str): The user's input concept or idea.
        ideas (list): List of ideas to choose from.
        generate_video (bool): Whether to generate a video from the image.
        max_retries (int): Maximum number of retries for image generation per cycle.
    
    Yields:
        tuple: (progress_percentage, message) for progress_html updates.
        dict: Final dictionary with 'text', 'image_base64', 'video_base64', and 'ideas'.
    """
    video_base64 = None
    max_total_attempts = 3  # Maximum total attempts for combined image and video generation cycles
    
    total_attempts = 0
    while total_attempts < max_total_attempts:
        total_attempts += 1
        yield (20 + total_attempts * 10, f"Attempt {total_attempts} to craft your {user_input} masterpiece... 🎨")
        
        # Step 1: Generate an image (retry up to max_retries times)
        generated_image = None
        text = None
        img_str = None
        image_prompt = None
        
        for image_attempt in range(max_retries):
            yield (20 + total_attempts * 10 + image_attempt * 5, f"Crafting a stunning image for {user_input}... 📸")
            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.
            Example: {{"caption": "Blindfolded Rubik's Cube MAGIC! 🤯 #rubiks", "image_prompt": "A close-up view of a person solving a Rubik's Cube blindfolded, in a dramatic style, no text or letters"}}
            """
            try:
                response = client.models.generate_content(
                    model='gemini-2.0-flash-lite',
                    contents=[prompt],
                    config=types.GenerateContentConfig(
                        temperature=1.2,
                        safety_settings=SAFETY_SETTINGS
                    )
                )
                print(f"Raw response for item (image attempt {image_attempt + 1}, total attempt {total_attempts}): {response.text}")  # Debugging
                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 'caption' not in response_json or 'image_prompt' not in response_json:
                    raise ValueError("Invalid JSON format: 'caption' or 'image_prompt' key missing")
                text = response_json['caption']
                image_prompt = response_json['image_prompt']
            except Exception as e:
                print(f"Error generating item (image attempt {image_attempt + 1}, total attempt {total_attempts}): {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"

            # Attempt to generate the image
            try:
                yield (40 + image_attempt * 5, f"Rendering your {user_input} vision... ✨")
                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))
                    # Ensure the image matches the desired aspect ratio (9:16 = 0.5625)
                    target_width = 360
                    target_height = int(target_width / 9 * 16)  # 9:16 aspect ratio
                    image = image.resize((target_width, target_height), Image.LANCZOS)
                    # Convert image to base64
                    buffered = BytesIO()
                    image.save(buffered, format="PNG")
                    img_str = base64.b64encode(buffered.getvalue()).decode()
                    yield (50, f"Image for {user_input} is ready! 🎉")
                    break  # Successfully generated image, exit image retry loop
                else:
                    print(f"Image generation failed (image attempt {image_attempt + 1}, total attempt {total_attempts}): No images returned")
                    if image_attempt == max_retries - 1:
                        yield (50, f"Tweaking the image for {user_input}... 🔄")
                        if total_attempts == max_total_attempts:
                            # Max total attempts reached, use a gray placeholder
                            image = Image.new('RGB', (360, 640), color='gray')
                            buffered = BytesIO()
                            image.save(buffered, format="PNG")
                            img_str = base64.b64encode(buffered.getvalue()).decode()
                            yield (60, f"Using a placeholder for {user_input}... 🖼️")
                            return {
                                'text': text,
                                'image_base64': img_str,
                                'video_base64': None,
                                'ideas': ideas
                            }
                        break  # Exit inner loop to retry with new idea
            except Exception as e:
                print(f"Error generating image (image attempt {image_attempt + 1}, total attempt {total_attempts}): {e}")
                if image_attempt == max_retries - 1:
                    yield (50, f"Retrying image for {user_input}... 🔄")
                    if total_attempts == max_total_attempts:
                        # Max total attempts reached, use a gray placeholder
                        image = Image.new('RGB', (360, 640), color='gray')
                        buffered = BytesIO()
                        image.save(buffered, format="PNG")
                        img_str = base64.b64encode(buffered.getvalue()).decode()
                        yield (60, f"Using a placeholder for {user_input}... 🖼️")
                        return {
                            'text': text,
                            'image_base64': img_str,
                            'video_base64': None,
                            'ideas': ideas
                        }
                    break  # Exit inner loop to retry with new idea

        # Step 2: Generate video if enabled (with fallback to text-to-video if image-to-video fails)
        if generate_video and generated_image is not None:
            max_video_retries_per_image = 2  # Try text-to-video generation twice if needed
            video_generated = False

            # First, try image-to-video generation (only once)
            try:
                yield (60, f"Filming a viral video for {user_input}... 🎥")
                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.
                """
                print(f"Attempting image-to-video generation (total attempt {total_attempts}): {video_prompt}")
                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"
                    )
                )
                # Wait for video to generate
                while not operation.done:
                    time.sleep(20)
                    operation = client.operations.get(operation)
                
                # Log detailed information about the operation
                print(f"Video generation operation completed: {operation}")
                print(f"Operation done: {operation.done}")
                print(f"Operation error: {operation.error}")
                if operation.error:
                    print(f"Operation error message: {operation.error.message}")
                    if hasattr(operation.error, 'code'):
                        print(f"Operation error code: {operation.error.code}")
                    if hasattr(operation.error, 'details'):
                        print(f"Operation error details: {operation.error.details}")
                print(f"Operation response: {operation.response}")
                if operation.response:
                    print(f"Operation response has generated_videos: {hasattr(operation.response, 'generated_videos')}")
                    if hasattr(operation.response, 'generated_videos'):
                        print(f"Generated videos: {operation.response.generated_videos}")
                    else:
                        print("No generated_videos attribute in response")

                # Enhanced error handling for video generation response
                if operation.error:
                    raise ValueError(f"Video generation operation failed with error: {operation.error.message}")
                if operation.response is None:
                    raise ValueError("Video generation operation failed: No response")
                if not hasattr(operation.response, 'generated_videos') or operation.response.generated_videos is None:
                    raise ValueError("Video generation operation failed: No generated_videos in response")
                
                # Process the single generated video
                if len(operation.response.generated_videos) > 0:
                    video = operation.response.generated_videos[0]
                    if video is None or not hasattr(video, 'video'):
                        raise ValueError("Video is invalid or missing video data")
                    fname = 'with_image_input.mp4'
                    print(f"Generated video: {fname}")
                    # Download the video and get the raw bytes
                    video_data = client.files.download(file=video.video)
                    # Ensure video_data is in bytes
                    if isinstance(video_data, bytes):
                        video_bytes = video_data
                    else:
                        # If video_data is a file-like object, read the bytes
                        video_buffer = BytesIO()
                        for chunk in video_data:
                            video_buffer.write(chunk)
                        video_bytes = video_buffer.getvalue()
                    # Encode the video bytes as base64
                    video_base64 = base64.b64encode(video_bytes).decode()
                    video_generated = True
                    yield (90, f"Video for {user_input} is a wrap! 🎬")
                    # Successfully generated video, return the result
                    return {
                        'text': text,
                        'image_base64': img_str,
                        'video_base64': video_base64,
                        'ideas': ideas
                    }
                else:
                    raise ValueError("No video was generated")
            except Exception as e:
                print(f"Error generating video (image-to-video, total attempt {total_attempts}): {e}")
                yield (70, f"Switching to a new video approach for {user_input}... 🎞️")
                print("Image-to-video generation failed. Falling back to text-to-video generation.")

            # If image-to-video generation failed, try text-to-video generation
            if not video_generated:
                for video_attempt in range(max_video_retries_per_image):
                    try:
                        yield (75 + video_attempt * 5, f"Trying a fresh video take for {user_input}... 📹")
                        # Use the same video prompt but without the image
                        video_prompt_base = 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.
                        """
                        if video_attempt == 0:
                            video_prompt = video_prompt_base
                        else:
                            video_prompt = f"""
                            The user concept is "{user_input}". Based on this and a simplified scene: {image_prompt}, create a video.
                            Use a static close-up shot of the subject in a realistic style.
                            """

                        print(f"Attempting text-to-video generation (video attempt {video_attempt + 1}, total attempt {total_attempts}): {video_prompt}")
                        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"
                            )
                        )
                        # Wait for video to generate
                        while not operation.done:
                            time.sleep(20)
                            operation = client.operations.get(operation)
                        
                        # Log detailed information about the operation
                        print(f"Video generation operation completed: {operation}")
                        print(f"Operation done: {operation.done}")
                        print(f"Operation error: {operation.error}")
                        if operation.error:
                            print(f"Operation error message: {operation.error.message}")
                            if hasattr(operation.error, 'code'):
                                print(f"Operation error code: {operation.error.code}")
                            if hasattr(operation.error, 'details'):
                                print(f"Operation error details: {operation.error.details}")
                        print(f"Operation response: {operation.response}")
                        if operation.response:
                            print(f"Operation response has generated_videos: {hasattr(operation.response, 'generated_videos')}")
                            if hasattr(operation.response, 'generated_videos'):
                                print(f"Generated videos: {operation.response.generated_videos}")
                            else:
                                print("No generated_videos attribute in response")

                        # Enhanced error handling for video generation response
                        if operation.error:
                            raise ValueError(f"Video generation operation failed with error: {operation.error.message}")
                        if operation.response is None:
                            raise ValueError("Video generation operation failed: No response")
                        if not hasattr(operation.response, 'generated_videos') or operation.response.generated_videos is None:
                            raise ValueError("Video generation operation failed: No generated_videos in response")
                        
                        # Process the single generated video
                        if len(operation.response.generated_videos) > 0:
                            video = operation.response.generated_videos[0]
                            if video is None or not hasattr(video, 'video'):
                                raise ValueError("Video is invalid or missing video data")
                            fname = 'text_to_video.mp4'
                            print(f"Generated video: {fname}")
                            # Download the video and get the raw bytes
                            video_data = client.files.download(file=video.video)
                            # Ensure video_data is in bytes
                            if isinstance(video_data, bytes):
                                video_bytes = video_data
                            else:
                                # If video_data is a file-like object, read the bytes
                                video_buffer = BytesIO()
                                for chunk in video_data:
                                    video_buffer.write(chunk)
                                video_bytes = video_buffer.getvalue()
                            # Encode the video bytes as base64
                            video_base64 = base64.b64encode(video_bytes).decode()
                            video_generated = True
                            yield (90, f"Video for {user_input} is a wrap! 🎬")
                            # Successfully generated video, return the result
                            return {
                                'text': text,
                                'image_base64': img_str,
                                'video_base64': video_base64,
                                'ideas': ideas
                            }
                        else:
                            raise ValueError("No video was generated")
                    except Exception as e:
                        print(f"Error generating video (text-to-video attempt {video_attempt + 1}, total attempt {total_attempts}): {e}")
                        if video_attempt == max_video_retries_per_image - 1:
                            yield (85, f"Finalizing without video for {user_input}... 📌")
                            if total_attempts == max_total_attempts:
                                print("Max total attempts reached. Proceeding without video.")
                                video_base64 = None
                                yield (95, f"Polishing your {user_input} masterpiece... ✨")
                                return {
                                    'text': text,
                                    'image_base64': img_str,
                                    'video_base64': video_base64,
                                    'ideas': ideas
                                }
                            # Text-to-video failed, break to outer loop to try a new image
                            print(f"Text-to-video generation failed after {max_video_retries_per_image} attempts. Selecting a new idea and generating a new image.")
                            break
                        continue  # Retry text-to-video generation with a modified prompt
        
        # If video generation is not enabled or image generation failed, return the result
        if img_str is not None:
            yield (95, f"Polishing your {user_input} masterpiece... ✨")
            return {
                'text': text,
                'image_base64': img_str,
                'video_base64': video_base64,
                'ideas': ideas
            }
        # If img_str is None, continue to next cycle or fall back if max attempts reached

    # If max total attempts reached without success, use a gray placeholder image
    print("Max total attempts reached without successful image generation. Using placeholder.")
    yield (95, f"Falling back to a placeholder for {user_input}... 🖼️")
    image = Image.new('RGB', (360, 640), color='gray')
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    yield (100, f"Ready to roll with {user_input}! 🚀")
    return {
        'text': f"Amazing {user_input}! 🔥 #{user_input.replace(' ', '')}",
        'image_base64': img_str,
        'video_base64': None,
        'ideas': ideas
    }

def generate_progress_html(progress, message, user_input):
    """
    Generate HTML for the progress bar and witty text.
    
    Args:
        progress (float): Current progress percentage (0–100).
        message (str): Current loading message to display.
        user_input (str): The user's input concept or idea for context.
    
    Returns:
        str: HTML string for the progress bar.
    """
    return f"""
    <div id="progress-container" style="
        display: flex;
        flex-direction: column;
        align-items: center;
        justify-content: center;
        max-width: 360px;
        margin: 0 auto;
        background-color: #000;
        height: 200px;
        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);
        ">
            {message}
        </div>
        <div style="
            width: 80%;
            height: 10px;
            background-color: #333;
            border-radius: 5px;
            overflow: hidden;
        ">
            <div id="progress-bar" style="
                width: {progress}%;
                height: 100%;
                background: linear-gradient(to right, #ff2d55, #ff5e78);
                transition: width 0.5s ease-in-out;
            "></div>
        </div>
        <div style="
            margin-top: 10px;
            font-size: 14px;
            color: #ccc;
        ">
            {int(progress)}% Complete
        </div>
        <style>
            @keyframes pulse {{
                0% {{ opacity: 1; }}
                50% {{ opacity: 0.5; }}
                100% {{ opacity: 1; }}
            }}
            #loading-message {{
                animation: pulse 2s infinite;
            }}
        </style>
    </div>
    """

def start_feed(user_input, generate_video, current_index, feed_items):
    """
    Start or update the feed based on the user input.
    Yields progress updates and final feed content.
    
    Args:
        user_input (str): The user's input concept or idea.
        generate_video (bool): Whether to generate a video.
        current_index (int): The current item index.
        feed_items (list): The current list of feed items.
    
    Yields:
        tuple: (progress_html) for progress updates.
        tuple: (current_user_input, current_index, feed_items, feed_html, share_html, is_loading) for final result.
    """
    if not user_input.strip():
        user_input = "trending"
    
    # Update current_user_input with the new user_input
    current_user_input = user_input
    
    # Set initial loading state
    is_loading = True
    progress_html = generate_progress_html(0, f"Getting started with {user_input}... 🚀", user_input)
    yield (progress_html,)
    share_links = ""

    try:
        # Generate ideas with progress updates
        ideas_gen = generate_ideas(user_input)
        ideas = None
        for update in ideas_gen:
            if isinstance(update, tuple):
                progress, message = update
                progress_html = generate_progress_html(progress, message, user_input)
                yield (progress_html,)
            else:
                ideas = update
        
        # Generate item with progress updates
        item_gen = generate_item(user_input, ideas, generate_video=generate_video)
        item = None
        for update in item_gen:
            if isinstance(update, tuple):
                progress, message = update
                progress_html = generate_progress_html(progress, message, user_input)
                yield (progress_html,)
            else:
                item = update
        
        feed_items = [item]
        current_index = 0
        feed_html = generate_html(feed_items, False, current_index, user_input, is_loading=False)
        share_links = generate_share_links(
            item['image_base64'],
            item['video_base64'],
            item['text']
        )
    except Exception as e:
        print(f"Error in start_feed: {e}")
        feed_items = []
        current_index = 0
        feed_html = """
        <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>
        """
        progress_html = generate_progress_html(100, "Oops, something went wrong! 😅", user_input)
        yield (progress_html,)
        is_loading = False
        return current_user_input, current_index, feed_items, feed_html, share_links, is_loading

    # Set loading state to False and update UI
    is_loading = False
    progress_html = ""
    return current_user_input, current_index, feed_items, feed_html, share_links, is_loading

def load_next(user_input, generate_video, current_index, feed_items):
    """
    Load the next item in the feed.
    Yields progress updates and final feed content.
    
    Args:
        user_input (str): The user's input concept or idea (updated from the textbox).
        generate_video (bool): Whether to generate a video.
        current_index (int): The current item index.
        feed_items (list): The current list of feed_items.
    
    Yields:
        tuple: (progress_html) for progress updates.
        tuple: (current_user_input, current_index, feed_items, feed_html, share_html, is_loading) for final result.
    """
    # Update current_user_input with the latest user_input from the textbox
    current_user_input = user_input if user_input.strip() else "trending"
    
    # Use current_user_input for generating the next item
    user_input = current_user_input
    
    is_loading = True
    progress_html = generate_progress_html(0, f"Loading next {user_input} vibe... 🚀", user_input)
    yield (progress_html,)
    share_links = ""

    try:
        if current_index + 1 < len(feed_items):
            current_index += 1
            progress_html = generate_progress_html(50, f"Switching to the next {user_input} moment... 🔄", user_input)
            yield (progress_html,)
        else:
            ideas = feed_items[-1]['ideas'] if feed_items else None
            if not ideas:
                # Generate new ideas if none exist
                ideas_gen = generate_ideas(user_input)
                for update in ideas_gen:
                    if isinstance(update, tuple):
                        progress, message = update
                        progress_html = generate_progress_html(progress, message, user_input)
                        yield (progress_html,)
                    else:
                        ideas = update
            
            # Generate new item
            new_item_gen = generate_item(user_input, ideas, generate_video=generate_video)
            new_item = None
            for update in new_item_gen:
                if isinstance(update, tuple):
                    progress, message = update
                    progress_html = generate_progress_html(progress, message, user_input)
                    yield (progress_html,)
                else:
                    new_item = update
            
            feed_items.append(new_item)
            current_index = len(feed_items) - 1
        
        feed_html = generate_html(feed_items, False, current_index, user_input, is_loading=False)
        share_links = generate_share_links(
            feed_items[current_index]['image_base64'],
            feed_items[current_index]['video_base64'],
            feed_items[current_index]['text']
        )
    except Exception as e:
        print(f"Error in load_next: {e}")
        feed_html = """
        <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>
        """
        progress_html = generate_progress_html(100, "Oops, something went wrong! 😅", user_input)
        yield (progress_html,)
        is_loading = False
        return current_user_input, current_index, feed_items, feed_html, share_links, is_loading

    is_loading = False
    progress_html = ""
    return current_user_input, current_index, feed_items, feed_html, share_links, is_loading

def load_previous(user_input, generate_video, current_index, feed_items):
    """
    Load the previous item in the feed.
    
    Args:
        user_input (str): The user's input concept or idea.
        generate_video (bool): Whether to generate a video (not used here).
        current_index (int): The current item index.
        feed_items (list): The current list of feed items.
    
    Returns:
        tuple: (current_user_input, current_index, feed_items, feed_html, share_html, is_loading)
    """
    # Update current_user_input with the latest user_input
    current_user_input = user_input if user_input.strip() else "trending"
    
    if current_index > 0:
        current_index -= 1
    feed_html = generate_html(feed_items, False, current_index, user_input, is_loading=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, feed_html, 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.
    Share links only include the caption, and users are instructed to manually upload the media.
    
    Args:
        image_base64 (str): The base64-encoded image data.
        video_base64 (str or None): The base64-enoded video data (single video).
        caption (str): The caption to share.
    
    Returns:
        str: HTML string with share links and download instructions.
    """
    image_data_url = f"data:image/png;base64,{image_base64}"
    encoded_caption = urllib.parse.quote(caption)

    # Generate download links for image and video (if available)
    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:  # Only include video download link if a video exists
        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 for users
    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>
    """

    # Generate share links for social media platforms (only passing the caption)
    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>
    """

    # Add YouTube Shorts share button if a video is available
    youtube_share = ""
    if video_base64:  # Only show YouTube Shorts share button if a video is generated
        youtube_share = f"""
            <div style="
                display: flex;
                justify-content: center;
                margin-top: 10px;
            ">
                <a href="https://studio.youtube.com/channel/UC/videos/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=0, message=""):
    """
    Generate an HTML string to display the current feed item or loading state.
    
    Args:
        feed_items (list): List of dictionaries containing 'text', 'image_base64', and 'video_base64'.
        scroll_to_latest (bool): Whether to auto-scroll to the latest item (not used here).
        current_index (int): The index of the item to display.
        user_input (str): The user's input concept or idea for context.
        is_loading (bool): Whether the feed is currently loading.
        progress (float): Current progress percentage (0–100).
        message (str): Current loading message to display.
    
    Returns:
        str: HTML string representing the feed or loading state.
    """
    if is_loading:
        # When loading, return an empty feed container (progress is handled by progress_html)
        return """
        <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;
            color: white;
            font-family: Arial, sans-serif;
        ">
        </div>
        """

    if not feed_items or current_index >= len(feed_items):
        return """
        <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;
        ">
            <p>Enter a concept or idea to start your feed!</p>
        </div>
        """

    item = feed_items[current_index]
    # Check if there is a video to display
    if item['video_base64']:
        media_element = 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>
        """
    else:
        # Fallback to image if no video is available
        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;
            ">
        """

    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>
    <button id="previous-button" style="display: none;" onclick="document.getElementById('previous-button').click()"></button>
    """
    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:
    # State variables
    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="")

    # Input section
    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("✨ Generate Next Item", elem_classes="gr-button")

    # Output display
    progress_html = gr.HTML(label="Progress", visible=True)
    feed_html = gr.HTML()
    share_html = gr.HTML(label="Share this item:")

    # Event handlers
    # Handle Enter keypress in the concept input
    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, progress_html]
    )

    # Handle magic button click to generate next item
    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, progress_html]
    )

    # Hidden button for previous item navigation
    previous_button = gr.Button("Previous", elem_id="previous-button", visible=False)

    # Handle click to go to previous item
    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, progress_html]
    )

# Launch the app
demo.launch()