import gradio as gr
import os
from PIL import Image
import numpy as np

# Custom theme for professional UI design
custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="lg",
    radius_size="md"
).set(
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_700",
    block_title_text_weight="600",
)

def process_single_image(image, prompt):
    """
    Process a single image with the given prompt
    """
    # Convert image to numpy array
    img_array = np.array(image)

    # Here you would add your actual processing logic
    # For demonstration, we'll just return the image with a watermark

    # Create a simple watermark with the prompt text
    result_img = Image.fromarray(img_array)
    result_img = add_watermark(result_img, prompt)

    return result_img

def add_watermark(image, text):
    """
    Add watermark text to an image
    """
    # Create a copy of the image
    img_copy = image.copy()

    # Create a drawing context
    draw = ImageDraw.Draw(img_copy)

    # Get image dimensions
    width, height = img_copy.size

    # Set font and text position
    try:
        font = ImageFont.truetype("arial.ttf", 30)
    except:
        font = ImageFont.load_default()

    # Calculate text size and position
    text_width, text_height = draw.textsize(text, font=font)
    position = (width - text_width - 10, height - text_height - 10)

    # Add semi-transparent background for text
    draw.rectangle(
        [position[0] - 5, position[1] - 5,
         position[0] + text_width + 5, position[1] + text_height + 5],
        fill=(255, 255, 255, 150)
    )

    # Draw the text
    draw.text(position, text, fill=(0, 0, 0), font=font)

    return img_copy

def process_batch(images, prompt):
    """
    Process multiple images with the same prompt
    """
    results = []
    for image in images:
        if image is not None:
            result = process_single_image(image, prompt)
            results.append(result)
        else:
            results.append(None)
    return results

def process_batch_files(image_files, prompt):
    """
    Process multiple image files with the same prompt
    """
    results = []
    for file_path in image_files:
        try:
            # Open the image file
            with Image.open(file_path) as img:
                result = process_single_image(img, prompt)
                results.append(result)
        except Exception as e:
            print(f"Error processing {file_path}: {e}")
            results.append(None)
    return results

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# 🖼️ Batch Image Processor")
    gr.Markdown("Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)")

    with gr.Tabs():
        # Single Image Processing Tab
        with gr.Tab("Single Image"):
            gr.Markdown("### Process a single image with your prompt")

            with gr.Row():
                with gr.Column():
                    single_image_input = gr.Image(label="Upload Image", type="pil")
                    single_prompt = gr.Textbox(label="Prompt", placeholder="Enter your processing prompt...")
                    single_process_btn = gr.Button("Process Image", variant="primary")

                with gr.Column():
                    single_output = gr.Image(label="Processed Image")

            single_process_btn.click(
                fn=process_single_image,
                inputs=[single_image_input, single_prompt],
                outputs=single_output,
                api_visibility="public"
            )

        # Batch Processing Tab
        with gr.Tab("Batch Processing"):
            gr.Markdown("### Process multiple images with the same prompt")

            with gr.Row():
                with gr.Column():
                    batch_images_input = gr.Gallery(
                        label="Upload Multiple Images",
                        type="pil",
                        height="auto"
                    )
                    batch_prompt = gr.Textbox(label="Prompt", placeholder="Enter your processing prompt...")
                    batch_process_btn = gr.Button("Process All Images", variant="primary")

                with gr.Column():
                    batch_output = gr.Gallery(label="Processed Images")

            batch_process_btn.click(
                fn=process_batch,
                inputs=[batch_images_input, batch_prompt],
                outputs=batch_output,
                api_visibility="public"
            )

        # File Batch Processing Tab
        with gr.Tab("File Batch Processing"):
            gr.Markdown("### Process multiple image files with the same prompt")

            with gr.Row():
                with gr.Column():
                    file_batch_input = gr.File(
                        label="Upload Image Files",
                        file_count="multiple",
                        file_types=["image"]
                    )
                    file_batch_prompt = gr.Textbox(label="Prompt", placeholder="Enter your processing prompt...")
                    file_batch_process_btn = gr.Button("Process All Files", variant="primary")

                with gr.Column():
                    file_batch_output = gr.Gallery(label="Processed Images")

            file_batch_process_btn.click(
                fn=process_batch_files,
                inputs=[file_batch_input, file_batch_prompt],
                outputs=file_batch_output,
                api_visibility="public"
            )

    # Examples section
    gr.Markdown("## 📚 Examples")
    examples = gr.Examples(
        examples=[
            ["https://gradio-builds.s3.amazonaws.com/assets/cheetah-003.jpg", "Wildlife processing"],
            ["https://gradio-builds.s3.amazonaws.com/assets/TheCheethcat.jpg", "Animal detection"],
            ["https://gradio-static-files.s3.amazonaws.com/world.mp4", "Video frame processing"]
        ],
        inputs=[single_image_input, single_prompt],
        outputs=single_output,
        cache_examples=True
    )

# Launch the application with custom theme and settings
demo.launch(
    theme=custom_theme,
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
        {"label": "Gradio Documentation", "url": "https://gradio.app/docs"}
    ],
    show_error=True,
    debug=False
)