import os
import requests
import gradio as gr
import uuid
import datetime
from supabase import create_client, Client
from supabase.lib.client_options import ClientOptions
import dotenv
from google.cloud import storage
import json
from pathlib import Path
import mimetypes
from workflow_handler import WanVideoWorkflow
from video_config import MODEL_FRAME_RATES, calculate_frames
import asyncio
from openai import OpenAI
import base64
from google.cloud import vision
from google.oauth2 import service_account

dotenv.load_dotenv()

SCRIPT_DIR = Path(__file__).parent
CONFIG_PATH = SCRIPT_DIR / "config.json"
WORKFLOW_PATH = SCRIPT_DIR / "wani2v.json"

loras = [
    {
      "image": "https://storage.googleapis.com/remade-v2/huggingface_assets/69e1adfc-b99a-4559-b745-f193a1bca0e2.gif",
      "id": "c8972c6d-ab8a-4988-9a9d-38082264ef22",
      "title": "Jumpscare",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },
    {
      "image": "https://storage.googleapis.com/remade-v2/huggingface_assets/2dede2c5-38e0-4acb-9e99-027e804a0455.gif",
      "id": "d7cbf9b4-82cd-4a94-ba2f-040e809635fa",
      "title": "Angry",
      "example_prompt": "The video starts with a man looking at the camera with a neutral face. Then his facial expression changes to 4ngr23 angry face, and he begins to yell with clenched fists. "
    },
    
    {
      "image": "https://storage.googleapis.com/remade-v2/huggingface_assets/7ba11998-2421-4ae1-8cf8-47b076387c2e.gif",
      "id": "e17959c4-9fa5-4e5b-8f69-d1fb01bbe4fa",
      "title": "Cartoon Jaw Drop",
      "example_prompt": "The video shows Pluto the dog, wearing a red collar, who is smiling wide, then his mouth transforms into a dr0p_j88 comical jaw drop, extending down in a long, rectangular shape, and revealing his tongue and teeth."
    },
    {
      "image": "https://storage.googleapis.com/remade-v2/huggingface_assets/46b5c7a3-60d5-469e-9454-9d16bf20afe4.gif",
      "id": "687255bb-959e-4422-bdbb-5aba93c7c180",
      "title": "Kissing",
      "example_prompt": "A man with a beard is shown smiling. A woman comes into the scene and starts passionately k144ing kissing the man."
    },
    {
      "image": "https://storage.googleapis.com/remade-v2/huggingface_assets/9f11b800-25c1-42f9-a687-97d706fef06d.gif",
      "id": "4ac2fb4e-5ca2-4338-a59c-549167f5b6d0",
      "title": "Laughing",
      "example_prompt": "A [Subject] is smiling at the camera. He/she then begins l4a6ing laughing."
    },
    {
      "image": "https://storage.googleapis.com/remade-v2/huggingface_assets/7cc38d79-f09e-4a5d-9e17-65867f8fd3a7.gif",
      "id": "bcc4163d-ebda-4cdc-b153-7136cdbf563a",
      "title": "Crying",
      "example_prompt": "The video starts with a [Subject] with a solemn expression. Then a tear rolls down his/her cheek, as he/she is cr471ng crying."
    },
    {
      "image": "https://storage.googleapis.com/remade-v2/huggingface_assets/62256895-10b2-4bd1-8ffe-4ecec530c4ec.gif",
      "id": "13093298-652c-4df8-ba28-62d9d5924754",
      "title": "Take a selfie with your younger self",
      "example_prompt": "The video starts with the a man with a beard smiling at the camera, then s31lf13 taking a selfie with their younger self, and the younger self appears next to him with similar facial features and eye color. The younger self wears a white t-shirt and has a cream white jacket. The younger self is smiling slightly."
    },
    {
      "image": "https://huggingface.co/Remade-AI/Squish/resolve/main/example_gifs/person_squish.gif",
      "id": "06ce6840-f976-4963-9644-b6cf7f323f90",
      "title": "Squish",
      "example_prompt": "In the video, a miniature rodent is presented. The rodent is held in a person's hands. The person then presses on the rodent, causing a sq41sh squish effect. The person keeps pressing down on the rodent, further showing the sq41sh squish effect.",
    },  
    {
      "image": "https://huggingface.co/Remade-AI/Rotate/resolve/main/example_videos/chair-rotate.gif",
      "id": "4ac08cfa-841e-4aa9-9022-c3fc80fb6ef4",
      "title": "Rotate",
      "example_prompt": "The video shows an elderly Asian man's head and shoulders with blurred background, performing a r0t4tion 360 degrees rotation.",
    },
    {
      "image": "https://huggingface.co/Remade-AI/Cakeify/resolve/main/example_gifs/timberland_cakeify.gif",
      "id": "b05c1dc7-a71c-4d24-b512-4877a12dea7e",
      "title": "Cakeify",
      "example_prompt": "The video opens on a woman. A knife, held by a hand, is coming into frame and hovering over the woman. The knife then begins cutting into the woman to c4k3 cakeify it. As the knife slices the woman open, the inside of the woman is revealed to be cake with chocolate layers. The knife cuts through and the contents of the woman are revealed."
    },    
]

# Initialize Supabase client with async support
supabase: Client = create_client(
    os.getenv('SUPABASE_URL'),
    os.getenv('SUPABASE_KEY'),
   
)

# Initialize OpenAI client
openai_client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

def initialize_gcs():
    """Initialize Google Cloud Storage client with credentials from environment"""
    try:
        # Parse service account JSON from environment variable
        service_account_json = os.getenv('SERVICE_ACCOUNT_JSON')
        if not service_account_json:
            raise ValueError("SERVICE_ACCOUNT_JSON environment variable not found")
        
        credentials_info = json.loads(service_account_json)
        
        # Initialize storage client
        storage_client = storage.Client.from_service_account_info(credentials_info)
        print("Successfully initialized Google Cloud Storage client")
        return storage_client
    except Exception as e:
        print(f"Error initializing Google Cloud Storage: {e}")
        raise

def upload_to_gcs(file_path, content_type=None, folder='user_uploads'):
    """
    Uploads a file to Google Cloud Storage
    Args:
        file_path: Path to the file to upload
        content_type: MIME type of the file (optional)
        folder: Folder path in bucket (default: 'user_uploads')
    Returns:
        str: Public URL of the uploaded file
    """
    try:
        bucket_name = 'remade-v2'
        storage_client = initialize_gcs()
        bucket = storage_client.bucket(bucket_name)

        # Get file extension and generate unique filename
        file_extension = Path(file_path).suffix
        if not content_type:
            content_type = mimetypes.guess_type(file_path)[0] or 'application/octet-stream'
        
        # Validate file type
        valid_types = ['image/jpeg', 'image/png', 'image/gif']
        if content_type not in valid_types:
            raise ValueError("Invalid file type. Please upload a JPG, PNG or GIF image.")

        # Generate unique filename with proper path structure
        filename = f"{str(uuid.uuid4())}{file_extension}"
        file_path_in_gcs = f"{folder}/{filename}"
        
        # Create blob and set metadata
        blob = bucket.blob(file_path_in_gcs)
        blob.content_type = content_type
        blob.cache_control = 'public, max-age=31536000'
        
        print(f'Uploading file to GCS: {file_path_in_gcs}')
        
        # Upload the file
        blob.upload_from_filename(
            file_path,
            timeout=120  # 2 minute timeout
        )
        
        # Generate public URL with correct path format
        image_url = f"https://storage.googleapis.com/{bucket_name}/{file_path_in_gcs}"
        print(f"Successfully uploaded to GCS: {image_url}")
        return image_url

    except Exception as e:
        print(f"Error uploading to GCS: {e}")
        raise ValueError(f"Failed to upload image to storage: {str(e)}")

def build_lora_prompt(subject, lora_id):
    """
    Builds a standardized prompt based on the selected LoRA and subject
    """
    # Get LoRA config
    lora_config = next((lora for lora in loras if lora["id"] == lora_id), None)
    if not lora_config:
        raise ValueError(f"Invalid LoRA ID: {lora_id}")
        
    if lora_id == "c8972c6d-ab8a-4988-9a9d-38082264ef22":  # Jumpscare
        return (
            f"The video shows the {subject} with a slight smile, then the j432mpscare jumpscare occurs, "
            f"revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
        )
    elif lora_id == "d7cbf9b4-82cd-4a94-ba2f-040e809635fa":  # Angry
        return (
            f"The video starts with the {subject} looking at the camera with a neutral face. "
            f"Then the facial expression of the {subject} changes to 4ngr23 angry face, and begins to yell with clenched fists."

        )
    elif lora_id == "e17959c4-9fa5-4e5b-8f69-d1fb01bbe4fa":  # Cartoon Jaw Drop
        return (
            f"The video shows {subject} smiling wide, "
            f"then {subject} mouth transforms into a dr0p_j88 comical jaw drop, extending down in a long, rectangular shape, and revealing his tongue and teeth."
        )
    elif lora_id == "687255bb-959e-4422-bdbb-5aba93c7c180":  # Kissing
        return (
            f"A {subject} is shown smiling. A man/woman comes into the scene and starts passionately k144ing kissing the {subject}."
        )
    elif lora_id == "4ac2fb4e-5ca2-4338-a59c-549167f5b6d0":  # Laughing
        return (
            f"A {subject} is smiling at the camera. He/she then begins l4a6ing laughing."
        )
    elif lora_id == "bcc4163d-ebda-4cdc-b153-7136cdbf563a":  # Crying
        return (
            f"The video starts with a {ubject} with a solemn expression. Then a tear rolls down his/her cheek, as he/she is cr471ng crying."
        )
    elif lora_id == "13093298-652c-4df8-ba28-62d9d5924754":  # Take a selfie with your younger self
        return (
            f"The video starts with the {subject} smiling at the camera, then s31lf13 taking a selfie with their younger self, "
            f"and the younger self appears next to the {subject} with similar facial features and eye color. "
            f"The younger self wears a white t-shirt and has a cream white jacket. The younger self is smiling slightly."
        )
        
    elif lora_id == "06ce6840-f976-4963-9644-b6cf7f323f90":  # Squish
        return (
            f"In the video, a miniature {subject} is presented. "
            f"The {subject} is held in a person's hands. "
            f"The person then presses on the {subject}, causing a sq41sh squish effect. "
            f"The person keeps pressing down on the {subject}, further showing the sq41sh squish effect."
        )
    
    elif lora_id == "4ac08cfa-841e-4aa9-9022-c3fc80fb6ef4":  # Rotate
        return (
            f"The video shows a {subject} performing a r0t4tion 360 degrees rotation."
        )
    
    elif lora_id == "b05c1dc7-a71c-4d24-b512-4877a12dea7e":  # Cakeify
        return (
            f"The video opens on a {subject}. A knife, held by a hand, is coming into frame "
            f"and hovering over the {subject}. The knife then begins cutting into the {subject} "
            f"to c4k3 cakeify it. As the knife slices the {subject} open, the inside of the "
            f"{subject} is revealed to be cake with chocolate layers. The knife cuts through "
            f"and the contents of the {subject} are revealed."
        )
    else:
        # Fallback to using the example prompt from the LoRA config
        if "example_prompt" in lora_config:
            # Replace any specific subject in the example with the user's subject
            return lora_config["example_prompt"].replace("rodent", subject).replace("woman", subject).replace("man", subject)
        else:
            raise ValueError(f"Unknown LoRA ID: {lora_id} and no example prompt available")

def poll_generation_status(generation_id):
    """Poll generation status from database"""
    try:
        # Query the database for the current status
        response = supabase.table('generations') \
            .select('*') \
            .eq('generation_id', generation_id) \
            .execute()
        
        if not response.data:
            return None
        
        return response.data[0]
    except Exception as e:
        print(f"Error polling generation status: {e}")
        raise e

async def moderate_prompt(prompt: str) -> dict:
    """
    Check if a text prompt contains NSFW content with strict rules against inappropriate content
    """
    try:
        # First check with OpenAI moderation
        response = openai_client.moderations.create(input=prompt)
        result = response.results[0]
        
        if result.flagged:
            # Find which categories were flagged
            flagged_categories = [
                category for category, flagged in result.categories.model_dump().items() 
                if flagged
            ]
            
            return {
                "isNSFW": True,
                "reason": f"Content flagged for: {', '.join(flagged_categories)}"
            }
        
        # Additional checks for keywords related to minors or inappropriate content
        keywords = [
            "child", "kid", "minor", "teen", "young", "baby", "infant", "underage",
            "naked", "nude", "nsfw", "porn", "xxx", "sex", "explicit",
            "inappropriate", "adult content"
        ]
        
        lower_prompt = prompt.lower()
        found_keywords = [word for word in keywords if word in lower_prompt]
        
        if found_keywords:
            return {
                "isNSFW": True,
                "reason": f"Content contains inappropriate keywords: {', '.join(found_keywords)}"
            }
        
        return {"isNSFW": False, "reason": None}
    except Exception as e:
        print(f"Error during prompt moderation: {e}")
        # If there's an error, reject the prompt to be safe
        return {
            "isNSFW": True,
            "reason": "Failed to verify prompt safety - please try again"
        }

async def moderate_image(image_path: str) -> dict:
    """
    Check if an image contains NSFW content using both Google Cloud Vision API's SafeSearch detection
    and OpenAI's vision model for double verification
    """
    try:
        # Convert image to base64 for OpenAI
        with open(image_path, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
        
        # 1. Google Cloud Vision API Check using proper client library
        try:
            # Get service account info from environment
            service_account_info = json.loads(os.getenv('SERVICE_ACCOUNT_JSON'))
            
            # Initialize Vision client with credentials
            credentials = service_account.Credentials.from_service_account_info(service_account_info)
            vision_client = vision.ImageAnnotatorClient(credentials=credentials)
            
            # Load image content
            with open(image_path, "rb") as image_file:
                content = image_file.read()
            
            # Create image object
            image = vision.Image(content=content)
            
            # Perform safe search detection
            response = vision_client.safe_search_detection(image=image)
            safe_search = response.safe_search_annotation
            
            # Map likelihood values
            likelihood_values = {
                vision.Likelihood.VERY_LIKELY: 4,
                vision.Likelihood.LIKELY: 3,
                vision.Likelihood.POSSIBLE: 2,
                vision.Likelihood.UNLIKELY: 1,
                vision.Likelihood.VERY_UNLIKELY: 0,
                vision.Likelihood.UNKNOWN: 0
            }
            
            # Get likelihood scores
            adult_score = likelihood_values[safe_search.adult]
            racy_score = likelihood_values[safe_search.racy]
            violence_score = likelihood_values[safe_search.violence]
            medical_score = likelihood_values[safe_search.medical]
            
            # Determine if content is NSFW according to Vision API
            vision_reasons = []
            if adult_score >= 3:  # LIKELY or VERY_LIKELY
                vision_reasons.append("adult content")
            if racy_score >= 3:  # LIKELY or VERY_LIKELY
                vision_reasons.append("suggestive content")
            if violence_score >= 3:  # LIKELY or VERY_LIKELY
                vision_reasons.append("violent content")

                
            # Print Vision API results
            print("Google Cloud Vision API Results:")
            print(f"Adult: {vision.Likelihood(safe_search.adult).name}")
            print(f"Racy: {vision.Likelihood(safe_search.racy).name}")
            print(f"Violence: {vision.Likelihood(safe_search.violence).name}")
            print(f"Medical: {vision.Likelihood(safe_search.medical).name}")
            
        except Exception as vision_error:
            print(f"Error with Vision API: {vision_error}")
            vision_reasons = []  # Continue with OpenAI check if Vision API fails
        
        # 2. OpenAI Vision Check
        openai_response = openai_client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "system",
                    "content": "You are a content review system. Your task is to analyze the image and determine if it contains any inappropriate content that would make it unsuitable for general audiences. Especially don't allow any content involving minors (children). Please respond with 'SAFE' if the image is appropriate, or 'NSFW: <reason>' if you detect any concerning content."
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Is this image appropriate for general audiences?"},
                        {"type": "image_url", "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}"
                        }}
                    ]
                }
            ],
            max_tokens=150
        )
        
        openai_result = openai_response.choices[0].message.content.strip()
        print("\nOpenAI Vision Check Result:")
        print(openai_result)
        
        # Combine results from both services
        is_nsfw_vision = bool(vision_reasons)
        is_nsfw_openai = openai_result.startswith("NSFW:")
        
        # If either service flags the content, reject it
        if is_nsfw_vision or is_nsfw_openai:
            reasons = []
            if vision_reasons:
                reasons.append(f"Google Vision: {', '.join(vision_reasons)}")
            if is_nsfw_openai:
                reasons.append(f"OpenAI: {openai_result[5:].strip()}")
                
            return {
                "isNSFW": True,
                "reason": f"Image contains inappropriate content: {'; '.join(reasons)}"
            }
        
        return {"isNSFW": False, "reason": None}
        
    except Exception as e:
        print(f"Error during image moderation: {e}")
        # If there's an error, reject the image to be safe
        return {
            "isNSFW": True,
            "reason": "Failed to verify image safety - please try again"
        }

async def moderate_combined(prompt: str, image_path: str) -> dict:
    """
    Check if both a prompt and image contain NSFW content together with strict rules
    """
    try:
        # Convert image to base64
        with open(image_path, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
        
        response = openai_client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "system",
                    "content": """You are a strict content moderation system. Your task is to determine if an image and prompt combination would result in inappropriate content.

STRICT RULES:
1. NO content involving minors in any way
2. NO nudity or sexually suggestive content
3. NO extreme violence or gore
4. NO hate speech or discriminatory content
5. NO illegal activities

Respond with 'NSFW: <reason>' if ANY of these rules are violated, or 'SAFE' if appropriate.
Be extremely cautious - if there's any doubt, mark it as NSFW."""
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": f'Please moderate this image and prompt combination for an image-to-video generation:\n\nPrompt: "{prompt}"\n\nEnsure NO inappropriate content, especially involving minors.'
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ],
            max_tokens=150
        )
        
        result = response.choices[0].message.content.strip()
        if result.startswith("NSFW:"):
            return {
                "isNSFW": True,
                "reason": result[5:].strip()
            }
        return {
            "isNSFW": False,
            "reason": None
        }
    except Exception as e:
        print(f"Error during combined moderation: {e}")
        # If there's an error, reject to be safe
        return {
            "isNSFW": True,
            "reason": "Failed to verify content safety - please try again"
        }

async def generate_video(input_image, subject, duration, selected_index, progress=gr.Progress()):
    try:
        # Initialize workflow handler with explicit paths
        workflow_handler = WanVideoWorkflow(
            supabase,
            config_path=str(CONFIG_PATH),
            workflow_path=str(WORKFLOW_PATH)
        )
        
        # Check if the input is a URL (example image) or a file path (user upload)
        if input_image.startswith('http'):
            # It's already a URL, use it directly
            image_url = input_image
        else:
            # It's a file path, upload to GCS
            image_url = upload_to_gcs(input_image)
        
        # Map duration selection to actual seconds
        duration_mapping = {
            "Short (3s)": 3,
            "Long (5s)": 5
        }
        video_duration = duration_mapping[duration]
        
        # Get LoRA config
        lora_config = next((lora for lora in loras if lora["id"] == selected_index), None)
        if not lora_config:
            raise ValueError(f"Invalid LoRA ID: {selected_index}")

        # Generate unique ID
        generation_id = str(uuid.uuid4())

        # Update workflow
        prompt = build_lora_prompt(subject, selected_index)
        workflow_handler.update_prompt(prompt)
        workflow_handler.update_input_image(image_url)
        await workflow_handler.update_lora(lora_config)
        workflow_handler.update_length(video_duration)
        workflow_handler.update_output_name(generation_id)

        # Get final workflow
        workflow = workflow_handler.get_workflow()

        # Store generation data in Supabase
        generation_data = {
            "generation_id": generation_id,
            "user_id": "anonymous",
            "status": "queued",
            "progress": 0,
            "worker_id": None,
            "created_at": datetime.datetime.utcnow().isoformat(),
            "message": {
                "generationId": generation_id,
                "workflow": {
                    "prompt": workflow
                }
            },
            "metadata": {
                "prompt": {
                    "original": subject,
                    "enhanced": subject
                },
                "lora": {
                    "id": selected_index,
                    "strength": 1.0,
                    "name": lora_config["title"]
                },
                "workflow": "img2vid",
                "dimensions": None,
                "input_image_url": image_url,
                "video_length": {"duration": video_duration},
            },
            "error": None,
            "output_url": None,
            "batch_id": None,
            "platform": "huggingface"
        }
        
        # Remove await - the execute() method returns the response directly
        response = supabase.table('generations').insert(generation_data).execute()
        print(f"Stored generation data with ID: {generation_id}")
        
        # Return generation ID for tracking
        return generation_id
        
    except Exception as e:
        print(f"Error in generate_video: {e}")
        raise e

def update_selection(evt: gr.SelectData):
  selected_lora = loras[evt.index]  
  sentence = f"Selected LoRA: {selected_lora['title']}"
  return selected_lora['id'], sentence

async def handle_generation(image_input, subject, duration, selected_index, progress=gr.Progress(track_tqdm=True)):
    try:
        if selected_index is None:
            raise gr.Error("You must select a LoRA before proceeding.")
        
        # First, moderate the prompt
        prompt_moderation = await moderate_prompt(subject)
        print(f"Prompt moderation result: {prompt_moderation}")
        if prompt_moderation["isNSFW"]:
            raise gr.Error(f"Content moderation failed: {prompt_moderation['reason']}")
        
        # Then, moderate the image
        image_moderation = await moderate_image(image_input)
        print(f"Image moderation result: {image_moderation}")
        if image_moderation["isNSFW"]:
            raise gr.Error(f"Content moderation failed: {image_moderation['reason']}")
        
        # Finally, check the combination
        combined_moderation = await moderate_combined(subject, image_input)
        print(f"Combined moderation result: {combined_moderation}")
        if combined_moderation["isNSFW"]:
            raise gr.Error(f"Content moderation failed: {combined_moderation['reason']}")
            
        # Generate the video and get generation ID
        generation_id = await generate_video(image_input, subject, duration, selected_index)
        
        # Poll for status updates
        while True:
            generation = poll_generation_status(generation_id)
            
            if not generation:
                raise ValueError(f"Generation {generation_id} not found")
                
            # Update progress
            if 'progress' in generation:
                progress_value = generation['progress']
                progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {progress_value}; --total: 100;"><span class="progress-text">Processing: {progress_value}%</span></div></div><div class="refresh-warning">Please do not refresh this page while processing</div>'
                
                # Check status
                if generation['status'] == 'completed':
                    # Final yield with completed video
                    yield generation['output_url'], generation_id, gr.update(visible=False)
                    break  # Exit the loop
                elif generation['status'] == 'error':
                    raise ValueError(f"Generation failed: {generation.get('error')}")
                else:
                    # Yield progress update
                    yield None, generation_id, gr.update(value=progress_bar, visible=True)
            
            # Wait before next poll
            await asyncio.sleep(2)
            
    except Exception as e:
        print(f"Error in handle_generation: {e}")
        raise e

css = '''
#gen_btn{height: 100%}
#gen_column{align-self: stretch}
#title{text-align: center}
#title h1{font-size: 3em; display:inline-flex; align-items:center}
#title img{width: 100px; margin-right: 0.5em}
#gallery .grid-wrap{height: auto; min-height: 350px}
#gallery .gallery-item {height: 100%; width: 100%; object-fit: cover}
#lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
.card_internal{display: flex;height: 100px;margin-top: .5em}
.card_internal img{margin-right: 1em}
.styler{--form-gap-width: 0px !important}
#progress{height:30px}
#progress .generating{display:none}
.progress-container {width: 100%;height: 30px;background-color: #2a2a2a;border-radius: 15px;overflow: hidden;margin-bottom: 20px;position: relative;}
.progress-bar {height: 100%;background-color: #7289DA;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
.progress-text {position: absolute;width: 100%;text-align: center;top: 50%;left: 0;transform: translateY(-50%);color: #ffffff;font-weight: bold;}
.refresh-warning {color: #ff7675;font-weight: bold;text-align: center;margin-top: 5px;}

/* Dark mode Discord styling */
.discord-banner {
    background: linear-gradient(135deg, #7289DA 0%, #5865F2 100%);
    color: #ffffff;
    padding: 20px;
    border-radius: 12px;
    margin: 15px 0;
    text-align: center;
    box-shadow: 0 4px 8px rgba(0,0,0,0.3);
}
.discord-banner h3 {
    margin-top: 0;
    font-size: 1.5em;
    text-shadow: 0 2px 4px rgba(0,0,0,0.3);
    color: #ffffff;
}
.discord-banner p {
    color: #ffffff;
    margin-bottom: 15px;
}
.discord-banner a {
    display: inline-block;
    background-color: #ffffff;
    color: #5865F2;
    text-decoration: none;
    font-weight: bold;
    padding: 10px 20px;
    border-radius: 24px;
    margin-top: 10px;
    transition: all 0.3s ease;
    box-shadow: 0 2px 8px rgba(0,0,0,0.3);
    border: none;
}
.discord-banner a:hover {
    transform: translateY(-3px);
    box-shadow: 0 6px 12px rgba(0,0,0,0.4);
    background-color: #f2f2f2;
}
.discord-feature {
    background-color: #2a2a2a;
    border-left: 4px solid #7289DA;
    padding: 12px 15px;
    margin: 10px 0;
    border-radius: 0 8px 8px 0;
    box-shadow: 0 2px 4px rgba(0,0,0,0.2);
    color: #e0e0e0;
}
.discord-feature-title {
    font-weight: bold;
    color: #7289DA;
}
.discord-locked {
    opacity: 0.7;
    position: relative;
    pointer-events: none;
}
.discord-locked::after {
    content: "🔒 Discord members only";
    position: absolute;
    top: 50%;
    left: 50%;
    transform: translate(-50%, -50%);
    background: rgba(114,137,218,0.9);
    color: white;
    padding: 5px 10px;
    border-radius: 20px;
    white-space: nowrap;
    font-size: 0.9em;
    font-weight: bold;
    box-shadow: 0 2px 4px rgba(0,0,0,0.3);
}
.discord-benefits-list {
    text-align: left;
    display: inline-block;
    margin: 10px 0;
    color: #ffffff;
}
.discord-benefits-list li {
    margin: 10px 0;
    position: relative;
    padding-left: 28px;
    color: #ffffff;
    font-weight: 500;
    text-shadow: 0 1px 2px rgba(0,0,0,0.2);
}
.discord-benefits-list li::before {
    content: "✨";
    position: absolute;
    left: 0;
    color: #FFD700;
}
.locked-option {
    opacity: 0.6;
    cursor: not-allowed;
}

/* Warning message styling */
.warning-message {
    background-color: #2a2a2a;
    border-left: 4px solid #ff7675;
    padding: 12px 15px;
    margin: 10px 0;
    border-radius: 0 8px 8px 0;
    box-shadow: 0 2px 4px rgba(0,0,0,0.2);
    color: #e0e0e0;
    font-weight: bold;
}

/* Example images and upload section styling */
.upload-section {
  display: flex;
  gap: 20px;
  margin: 20px 0;
}

.example-images-container {
  flex: 1;
}

.upload-container {
  flex: 1;
  display: flex;
  flex-direction: column;
  justify-content: center;
}

.section-title {
  font-weight: bold;
  margin-bottom: 10px;
  color: #7289DA;
}

.example-images-grid {
  display: grid;
  grid-template-columns: repeat(3, 1fr);
  gap: 10px;
}

.example-image-item {
  border-radius: 8px;
  overflow: hidden;
  cursor: pointer;
  transition: all 0.2s ease;
  border: 2px solid transparent;
}

.example-image-item:hover {
  transform: scale(1.05);
  box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
}

.example-image-item.selected {
  border-color: #7289DA;
}

.upload-button {
  margin-top: 15px;
}
'''

with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate", text_size="lg")) as demo:
    selected_index = gr.State(None)
    current_generation_id = gr.State(None)
    
    # Updated title with April Fool's theme
    gr.Markdown("# Remade AI - April Fool's Edition: Wan 2.1 I2V Effects")
    
    # Insert an April Fool's themed banner at the top
    april_banner = gr.HTML(
        """
        <div style="background-color: #8a0ee3; padding: 15px; border-radius: 8px; text-align: center; font-size: 1.5em; margin-bottom: 15px;">
            🎉 Happy April Fool's Day! Enjoy some playful pranks and fun effects! 🎉
        </div>
        """
    )
    
    # Optionally, update the Discord banner for an April Fool's twist
    discord_banner = gr.HTML(
        """
        <div class="discord-banner">
            <h3>✨ Unlock Premium April Fool's Pranks! ✨</h3>
            <p>Join our Discord community for exclusive prank effects, surprise features, and more playful fun!</p>
            <a href="https://remade.ai/join-discord?utm_source=Huggingface&utm_medium=Social&utm_campaign=hugginface_space&utm_content=april_fools" target="_blank">Join Discord Now</a>
        </div>
        """
    )
    
    selected_info = gr.HTML("")

    with gr.Row():
        with gr.Column(scale=1):
            gallery = gr.Gallery(
                [(item["image"], item["title"]) for item in loras],
                label="Select LoRA",
                allow_preview=False,
                columns=4,
                elem_id="gallery",
                show_share_button=False,
                height="650px",
                object_fit="contain"
            )
            
            # Updated Discord/prank callout
            gr.HTML(
                """
                <div class="discord-feature">
                    <span class="discord-feature-title">✨ Discord Members:</span> Get access to even more mischievous effects beyond these samples!
                </div>
                """
            )
            
            gr.HTML('<div class="section-description">Click an example image or upload your own</div>')
            
            with gr.Row():
                with gr.Column(scale=1):
                    example_gallery = gr.Gallery(
                        [
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_1c2c6e4c-8938-4464-9355-84508bcca24e.jpg", "Old man"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_1f2c6ec9-823f-46d2-982f-73c494e51877.jpg", "Young woman"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_24d949f0-8699-4714-9c82-854e1b963063.jpg", "Puppy"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_af26651e-be1a-40c0-be18-c42b3bf6d211.png", "Mini toy dancers"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_d22a894e-a074-4742-9e23-787f001a3184.jpg", "Chair"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_e6472106-4e9d-4620-b41b-a9bbe4893415.png", "Cartoon boy on bike")
                        ],
                        columns=3,
                        height="300px",
                        object_fit="cover"
                    )
                
                with gr.Column(scale=1):
                    image_input = gr.Image(type="filepath", label="")

            subject = gr.Textbox(label="Describe your subject", placeholder="Cat toy")
            
            duration = gr.Radio(
                ["Short (3s)"],
                label="Duration",
                value="Short (3s)"
            )
            
            # Updated Discord feature callout for additional playful messaging
            gr.HTML(
                """
                <div class="discord-feature">
                    <span class="discord-feature-title">⏱️ Discord Members:</span> Enjoy extended pranks and video durations on our Discord!
                </div>
                """
            )
            
            with gr.Row():
                button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
                audio_button = gr.Button("Add Audio 🔒", interactive=False)
            
        with gr.Column(scale=1):
            warning_message = gr.HTML(
                """
                <div class="warning-message">
                    ⚠️ Please DO NOT refresh the page during generation. Our pranksters are hard at work!
                </div>
                """,
                visible=True
            )
            
            gr.HTML(
                """
                <div class="discord-feature">
                    <span class="discord-feature-title">⚡ Discord Members:</span> Get faster (and prankier) generation speeds!
                </div>
                """
            )
            
            progress_bar = gr.Markdown(elem_id="progress", visible=False)
            output = gr.Video(interactive=False, label="Output video")
    
    gallery.select(
        update_selection,
        outputs=[selected_index, selected_info]
    )
    
    # Modified function to handle example image selection
    def select_example_image(evt: gr.SelectData):
        """Handle example image selection and return image URL, description, and update image source"""
        example_images = [
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_1c2c6e4c-8938-4464-9355-84508bcca24e.jpg",
                "description": "Old man"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_1f2c6ec9-823f-46d2-982f-73c494e51877.jpg",
                "description": "Young woman"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_24d949f0-8699-4714-9c82-854e1b963063.jpg",
                "description": "Puppy"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_af26651e-be1a-40c0-be18-c42b3bf6d211.png",
                "description": "Mini toy dancers"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_d22a894e-a074-4742-9e23-787f001a3184.jpg",
                "description": "Chair"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_e6472106-4e9d-4620-b41b-a9bbe4893415.png",
                "description": "Cartoon boy on bike"
            }
        ]
        
        selected = example_images[evt.index]
        
        # Return the URL, description, and update image source to "example"
        return selected["url"], selected["description"], "example"
    
    # Connect example gallery selection to image_input and subject
    example_gallery.select(
        fn=select_example_image,
        outputs=[image_input, subject]
    )
    
    # Add a custom handler to check if inputs are valid
    def check_inputs(subject, image_input, selected_index):
        if not selected_index:
            raise gr.Error("You must select a LoRA before proceeding.")
        if not subject.strip():
            raise gr.Error("Please describe your subject.")
        if image_input is None:
            raise gr.Error("Please upload an image or select an example image.")
    
    # Use gr.on for the button click with validation
    button.click(
        fn=check_inputs,
        inputs=[subject, image_input, selected_index],
        outputs=None,
    ).success(
        fn=handle_generation,
        inputs=[image_input, subject, duration, selected_index],
        outputs=[output, current_generation_id, progress_bar]
    )
    
    # Add a click handler for the disabled audio button
    audio_button.click(
        fn=lambda: gr.Info("Join our Discord to unlock audio generation features!"),
        inputs=None,
        outputs=None
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=20)
    demo.launch(ssr_mode=False, share=True)