import requests
import base64
import json
import argparse
import os
import glob



# Set up argument parser
parser = argparse.ArgumentParser(description='Test the /gen_latest_pic endpoint by sending the latest image from a directory.')
parser.add_argument('directory', type=str, nargs='?', default='.', 
                    help='Directory to search for the latest image. Defaults to the current directory.')
args = parser.parse_args()


# The URL of the Flask application endpoint
url = 'http://127.0.0.1:5000/gen_latest_pic'

# Find the latest image in the specified directory
image_dir = args.directory
# Add more extensions if needed
image_extensions = ('*.png', '*.jpg', '*.jpeg', '*.gif', '*.bmp')
all_images = []
for ext in image_extensions:
    # Use glob to find files with the current extension
    all_images.extend(glob.glob(os.path.join(image_dir, ext)))

if not all_images:
    print(f"Error: No image files found in the directory '{image_dir}'.")
    exit()

# Get the latest file by modification time
try:
    image_path = max(all_images, key=os.path.getmtime)
    print(f"Using latest image: {image_path}")
except FileNotFoundError:
    print(f"Error: The directory '{image_dir}' does not exist.")
    exit()


# Read the image file in binary mode
with open(image_path, 'rb') as image_file:
    # Encode the image data in base64
    encoded_string = base64.b64encode(image_file.read()).decode('utf-8')

# Prepare the data to be sent in the request
# The 'pid' is optional, so you can include it or not
data = {
    'base_image': encoded_string,
    'pid': 'user123'  # Optional: you can change or remove this
}

# Set the headers for the request to indicate JSON content
headers = {
    'Content-Type': 'application/json'
}

try:
    # Send the POST request
    response = requests.post(url, data=json.dumps(data), headers=headers, timeout=60)

    # Check the response from the server
    if response.status_code == 201:
        print("Request was successful.")
        print("Response JSON:", response.json())
    else:
        print(f"Request failed with status code {response.status_code}")
        try:
            print("Error response:", response.json())
        except json.JSONDecodeError:
            print("Error response could not be decoded as JSON:", response.text)

except requests.exceptions.RequestException as e:
    print(f"An error occurred: {e}")