import cv2
import os
import sys
import numpy as np

def fit_image_to_frame(input_image_path, frame_image_path):
    try:
        # Read the input image and the frame image
        input_image = cv2.imread(input_image_path, cv2.IMREAD_UNCHANGED)  # Read with alpha channel if present
        frame_image = cv2.imread(frame_image_path, cv2.IMREAD_UNCHANGED)  # Read with alpha channel if present

        # Check if images were loaded successfully
        if input_image is None or frame_image is None:
            raise ValueError("Error loading images. Please check the file paths.")

        # Print the shapes and types for debugging
        print(f"Input image shape: {input_image.shape}, Frame image shape: {frame_image.shape}")

        # Resize the input image to 904x904
        input_image = cv2.resize(input_image, (904, 904), interpolation=cv2.INTER_LANCZOS4)

        # Create a new image with the same size as the frame
        combined_image = np.zeros_like(frame_image)  # Initialize with zeros (black background)

        # Calculate position to center the input image
        input_x = (frame_image.shape[1] - input_image.shape[1]) // 2
        input_y = (frame_image.shape[0] - input_image.shape[0]) // 2

        # First, place the input image in the center of the combined image
        if input_image.shape[2] == 4:  # RGBA
            # For RGBA images, we need to handle the alpha channel
            for c in range(0, 3):  # For each color channel (RGB)
                combined_image[input_y:input_y + input_image.shape[0], input_x:input_x + input_image.shape[1], c] = input_image[:, :, c]
            # Copy the alpha channel
            combined_image[input_y:input_y + input_image.shape[0], input_x:input_x + input_image.shape[1], 3] = input_image[:, :, 3]
        else:  # RGB
            # For RGB images, we need to create an alpha channel
            for c in range(0, 3):  # For each color channel (RGB)
                combined_image[input_y:input_y + input_image.shape[0], input_x:input_x + input_image.shape[1], c] = input_image[:, :, c]
            # Set the alpha channel to fully opaque
            combined_image[input_y:input_y + input_image.shape[0], input_x:input_x + input_image.shape[1], 3] = 255

        # Now, overlay the frame on top of the combined image
        # The frame's alpha channel will determine how much of the input image shows through
        frame_alpha = frame_image[:, :, 3] / 255.0  # Normalize alpha channel to [0, 1]
        for c in range(0, 3):  # For each color channel (RGB)
            combined_image[:, :, c] = (1 - frame_alpha) * combined_image[:, :, c] + frame_alpha * frame_image[:, :, c]

        # Ensure the alpha channel is preserved
        combined_image[:, :, 3] = np.maximum(combined_image[:, :, 3], frame_image[:, :, 3])

        return combined_image  # Return the combined image

    except Exception as e:
        print(f"Error: {e}")  # Print the error message for debugging
        return None  # Return None if there was an error

def save_as_jpg(output_image, output_path, quality=95):
    # Save the combined image as a JPG file with specified quality
    # JPG doesn't support alpha channel, so we need to convert to BGR
    if output_image.shape[2] == 4:  # If RGBA
        # Create white background
        background = np.ones_like(output_image, dtype=np.uint8) * 255
        # Only use RGB channels for JPG
        mask = output_image[:,:,3:] / 255.0
        output_image_rgb = output_image[:,:,:3]
        background[:,:,:3] = background[:,:,:3] * (1 - mask) + output_image_rgb * mask
        # Save as JPG
        cv2.imwrite(output_path, background[:,:,:3], [cv2.IMWRITE_JPEG_QUALITY, quality])
    else:
        # Save as JPG
        cv2.imwrite(output_path, output_image, [cv2.IMWRITE_JPEG_QUALITY, quality])

if __name__ == "__main__":
    # Check if at least one input image is provided
    if len(sys.argv) < 3:  # Ensure at least one input image and one frame image
        print("Usage: python FrameTool.py <input_image_path1> <input_image_path2> ... <frame_image_path>")
        sys.exit(1)

    # Get the frame image path from the last argument
    frame_image_path = sys.argv[-1]  # The last argument is the frame image

    # Check if the frame image exists
    if not os.path.exists(frame_image_path):
        print(f"Frame image not found: {frame_image_path}")
        sys.exit(1)

    # Get the directory of the script for output
    script_dir = os.path.dirname(os.path.abspath(__file__))

    # Process each input image
    for input_image_path in sys.argv[1:-1]:  # Exclude the last argument (frame image path)
        # Generate a unique output image path in the script directory
        output_image_path = os.path.join(script_dir, "output.jpg")  # Start with output.jpg

        # Ensure the output file does not overwrite existing files
        counter = 1
        while os.path.exists(output_image_path):
            output_image_path = os.path.join(script_dir, f"output_{counter}.jpg")
            counter += 1

        # Call the function to fit the image to the frame
        combined_image = fit_image_to_frame(input_image_path, frame_image_path)

        # Save the combined image if it was created successfully
        if combined_image is not None:
            save_as_jpg(combined_image, output_image_path)  # Save the combined image as JPG
            print(f"JPG saved successfully at: {output_image_path}")  # Log success
        else:
            print(f"Failed to process {input_image_path}.")