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 logging

# Set up logging
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

# 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.
    """
    logging.debug("Cleaning 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.
    """
    logging.debug(f"Generating ideas for input: {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.
    """
    try:
        response = client.models.generate_content(
            model='gemini-2.0-flash-lite',
            contents=[prompt],
            config=types.GenerateContentConfig(
                temperature=1.2,
                safety_settings=SAFETY_SETTINGS
            )
        )
        logging.debug(f"Generate ideas response: {response.text}")
        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")
        logging.debug(f"Generated ideas: {response_json['ideas']}")
        return response_json['ideas']
    except Exception as e:
        logging.error(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.
    """
    logging.debug("Starting generate_item")
    video_base64 = None
    max_total_attempts = 3
    timeout_seconds = 60  # Timeout for API calls

    total_attempts = 0
    while total_attempts < max_total_attempts:
        total_attempts += 1
        logging.debug(f"Total attempt {total_attempts}")
        
        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):
            logging.debug(f"Image attempt {image_attempt + 1}")
            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:
                start_time = time.time()
                response = client.models.generate_content(
                    model='gemini-2.0-flash-lite',
                    contents=[prompt],
                    config=types.GenerateContentConfig(
                        temperature=1.2,
                        safety_settings=SAFETY_SETTINGS
                    )
                )
                if time.time() - start_time > timeout_seconds:
                    raise TimeoutError("Image caption generation timed out")
                logging.debug(f"Generate content response: {response.text}")
                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:
                logging.error(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:
                start_time = time.time()
                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 time.time() - start_time > timeout_seconds:
                    raise TimeoutError("Image generation timed out")
                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()
                    logging.debug("Image generated successfully")
                    break
                else:
                    logging.warning("No images generated")
                    if image_attempt == max_retries - 1 and 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}
                        logging.debug("Returning with placeholder image")
                        return {
                            'text': text,
                            'image_base64': img_str,
                            'video_base64': None,
                            'ideas': ideas
                        }
                    yield {"stage": "retrying_image", "message": random.choice(PROGRESS_STAGES["retrying_image"]), "progress": 40 + (image_attempt * 10)}
            except Exception as e:
                logging.error(f"Error generating image: {e}")
                if image_attempt == max_retries - 1 and 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}
                    logging.debug("Returning with placeholder image")
                    return {
                        'text': text,
                        'image_base64': img_str,
                        'video_base64': None,
                        'ideas': ideas
                    }
                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:
            logging.debug("Attempting video generation")
            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.
                """
                start_time = time.time()
                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=5,
                        negative_prompt="blurry, low quality, text, letters"
                    )
                )
                while not operation.done:
                    if time.time() - start_time > timeout_seconds:
                        raise TimeoutError("Video generation timed out")
                    time.sleep(5)
                    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}
                            logging.debug("Video generated successfully")
                            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:
                logging.error(f"Error generating video: {e}")
                logging.debug("Falling back to text-to-video")
                yield {"stage": "generating_video", "message": random.choice(PROGRESS_STAGES["generating_video"]), "progress": 80}
                try:
                    start_time = time.time()
                    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=5,
                            negative_prompt="blurry, low quality, text, letters"
                        )
                    )
                    while not operation.done:
                        if time.time() - start_time > timeout_seconds:
                            raise TimeoutError("Text-to-video generation timed out")
                        time.sleep(5)
                        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}
                                logging.debug("Text-to-video generated successfully")
                                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:
                    logging.error(f"Error generating text-to-video: {e}")
                    if total_attempts == max_total_attempts:
                        yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
                        logging.debug("Returning without video")
                        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}
            logging.debug("Returning with image only")
            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}
    logging.debug("Returning with placeholder image")
    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.
    """
    logging.debug("Starting start_feed")
    if not user_input.strip():
        user_input = "trending"
    current_user_input = user_input
    is_loading = True
    share_links = ""
    timeout_seconds = 120  # Timeout for entire generation process
    
    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:
        start_time = time.time()
        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)
        try:
            while True:
                if time.time() - start_time > timeout_seconds:
                    logging.error("Generation timed out")
                    raise TimeoutError("Feed generation timed out")
                progress = next(item_generator)
                logging.debug(f"Progress update: {progress}")
                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:
                    logging.debug(f"Received final item: {progress}")
                    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)
                    logging.debug("Feed generation complete")
                    yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                    return
        except StopIteration as e:
            # Handle the generator's return value
            if e.value is not None:
                logging.debug(f"Generator returned final item: {e.value}")
                item = e.value
                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)
                logging.debug("Feed generation complete via StopIteration")
                yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                return
            else:
                logging.warning("Generator returned without a final item")
                raise ValueError("Generator did not return a valid item")
    except Exception as e:
        logging.error(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: {str(e)}. 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.
    """
    logging.debug("Starting load_next")
    current_user_input = user_input if user_input.strip() else "trending"
    user_input = current_user_input
    is_loading = True
    share_links = ""
    timeout_seconds = 120  # Timeout for entire generation process
    
    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']
            )
            logging.debug("Loaded existing feed item")
            yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
            return
        else:
            start_time = time.time()
            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)
            try:
                while True:
                    if time.time() - start_time > timeout_seconds:
                        logging.error("Generation timed out")
                        raise TimeoutError("Feed generation timed out")
                    progress = next(item_generator)
                    logging.debug(f"Progress update: {progress}")
                    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:
                        logging.debug(f"Received final item: {progress}")
                        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']
                        )
                        logging.debug("New feed item generated")
                        yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                        return
            except StopIteration as e:
                # Handle the generator's return value
                if e.value is not None:
                    logging.debug(f"Generator returned final item: {e.value}")
                    new_item = e.value
                    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']
                    )
                    logging.debug("New feed item generated via StopIteration")
                    yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                    return
                else:
                    logging.warning("Generator returned without a final item")
                    raise ValueError("Generator did not return a valid item")
    except Exception as e:
        logging.error(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: {str(e)}. 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.
    """
    logging.debug("Loading previous item")
    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']
    )
    logging.debug("Previous item loaded")
    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.
    """
    logging.debug("Generating share links")
    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={encoded_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.
    """
    logging.debug(f"Generating HTML, is_loading: {is_loading}, progress_stage: {progress_stage}")
    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):
        logging.debug("No feed items to display")
        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>
    """
    logging.debug("Feed item HTML generated")
    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()