import os
import base64
import json
import time

from utils import encode_image_to_jpeg_base64
# Import your VertexAIService class
from vertex_ai_service import VertexAIService  # Replace with the module name where VertexAIService is defined


def process_image_file(file_path, client):
    """
    Processes an image file and returns the path to the JSON file with results.

    :param file_path: Path to the image file
    :param client: Instance of VertexAIService
    :return: Path to the saved JSON file
    """
    input_image64 = encode_image_to_jpeg_base64(file_path)

    # Processing settings
    prompt = "Read the text"
    system = "You are receipt recognizer"

    # Call image processing
    result_img = client.process_image(input_image64, "gemini-1.5-flash", prompt, system, 0.0)

    # Create path for JSON file
    json_file_path = os.path.splitext(file_path)[0] + ".json"

    # Write the result to a JSON file
    with open(json_file_path, 'w', encoding='utf-8') as json_file:
        json_file.write(result_img)

    return json_file_path


def process_folder_images_sequentially(folder_path, json_key_path):
    """
    Processes all images in the specified folder sequentially.

    :param folder_path: Path to the folder containing images
    :param json_key_path: Path to the JSON key for Vertex AI authentication
    """
    # Initialize Vertex AI client
    client = VertexAIService(json_key_path=json_key_path)

    # Get a list of all images in the folder
    image_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(('.jpg', '.jpeg', '.png', 'webp'))]

    # Process each image file sequentially
    for file_path in image_files:
        try:
            json_file_path = process_image_file(file_path, client)
            print(f"Processed: {file_path}, result saved to: {json_file_path}")
            time.sleep(60)
        except Exception as exc:
            print(f"Failed to process file {file_path}. Error: {exc}")


# Call the function to process the folder
if __name__ == '__main__':
    folder_path = './examples'  # Set the path to your folder
    json_key_path = 'secrets/GOOGLE_VERTEX_AI_KEY.json'  # Set the path to your JSON key
    process_folder_images_sequentially(folder_path, json_key_path)