import gradio as gr
import pandas as pd
from PIL import Image
import os
import uuid
import boto3

image_directory = "images"
os.makedirs(image_directory, exist_ok=True)

aws_access_key = os.environ['AWS_ACCESS_KEY']
aws_secret_key = os.environ['AWS_SECRET_KEY']
aws_region = os.environ['AWS_REGION']

s3 = boto3.client('s3', aws_access_key_id=aws_access_key, aws_secret_access_key=aws_secret_key, region_name=aws_region)

s3_bucket_name = os.environ['S3_BUCKET_NAME']

data = []


base_csv_file_path = "descriptions"
csv_file_path = f"{base_csv_file_path}.csv"


def file_exists(file_path):
    return os.path.isfile(file_path)


def generate_unique_csv_file():
    count = 1
    while file_exists(f"{base_csv_file_path}{count}.csv"):
        count += 1
    return f"{base_csv_file_path}{count}.csv"


def describe_image(image, filename, description, copyright):
    global data

   
    image_filename = f"{str(uuid.uuid4())}.jpg"

    
    local_image_path = os.path.join(image_directory, image_filename)
    image.save(local_image_path)

   
    s3.upload_file(local_image_path, s3_bucket_name, f"images/{image_filename}")

   
    os.remove(local_image_path)

    
    data.append({"Image": image_filename, "Filename": filename, "Description": description, "Copyright": copyright})        
    data_df = pd.DataFrame(data)

    
    if file_exists(csv_file_path):
        
        new_csv_file_path = generate_unique_csv_file()
        data_df.to_csv(new_csv_file_path, index=False)
        s3.upload_file(new_csv_file_path, s3_bucket_name, new_csv_file_path)
    else:
       
        data_df.to_csv(csv_file_path, index=False)
        s3.upload_file(csv_file_path, s3_bucket_name, csv_file_path)

        gr.Info('Data saved.')
    

    return None 

def auth_fn(username, password):
    return username == os.getenv("SM_USERNAME") and password == os.getenv("SM_PASSWORD")


iface = gr.Interface(
    fn=describe_image,
    inputs=[
        gr.Image(type="pil", label="Upload an Image"),
        "text",
        "text",
        "text"
    ],
    outputs=gr.Info("saved"),  # Display the success message
    live=False,
    title="Image Description App",
    description="Upload an image and provide three descriptions for it.",
    allow_flagging=False,
    theme=gr.themes.Base(primary_hue="green").set(
        button_primary_background_fill="*primary_400",
        button_primary_background_fill_hover="*primary_300",
    )
)

iface.launch(auth=auth_fn)