# byte is assumed to be octet
# char is assumed to be 1 byte
#
# struct rawimg {
#   char[8]   magic_num = "rifV0001" // Big Endian. UTF-8 Encoded. 72 69 66 56 30 30 30 31
#   uint64_t  width;                 // Big Endian
#   uint8_t   format;
#   uint8_t[] data;
# }
#
# formats:
#   0x00 = (24 bits / 3 bytes) R8 B8 G8
#   0x01 = (32 bits / 4 bytes) R8 B8 G8 A8
#   ... undefined in V0001
#
# data is a flat array, starting at top left, going right then bottom.

# Python 3.10
# encode.py uses pillow for input image reading
# Usage: python encode.py <input image path> <output image path>

# Pillow
from PIL import Image
# For arguments and error throwing
import sys
# For path checking (argument validation)
import os

# Check if the argument count matches 3 (Zeroth argument is the script name as is convention)
if len(sys.argv) != 3:
    sys.exit("Wrong argument count. Usage: python encode.py <input image path> <output image path>")

# First argument is the image path to open
in_path = sys.argv[1]
# Second argument is the path to write to
out_path = sys.argv[2]

# Check if the input file exists
if not os.path.exists(in_path):
    sys.exit(f"{in_path} does not exist. Usage: python encode.py <input image path> <output image path>")

# Check if the output path exists
if os.access(os.path.dirname(out_path), os.W_OK):
    sys.exit(f"{out_path} does not exist or is not writable. Usage: python encode.py <input image path> <output image path>")

# Open the input file
img_in = Image.open(in_path)

# Open the output file for writing (in binary mode - python specific)
with open(out_path, "wb") as img_out:
    # Write the magic number
    img_out.write("rifV0001".encode("utf-8"))
    # Write the width from the input image
    img_out.write((img_in.width).to_bytes(8, "big"))
    # Write the mode, depending on the "mode" (pixel format) of the input image
    match img_in.mode:
        # If RGB, write mode 0x00
        case "RGB":
            img_out.write(bytes([0x00]))
        # If RGBA, write mode 0x01
        case "RGBA":
            img_out.write(bytes([0x01]))
        # If something else, error out
        case _:
            sys.exit(f"Image {in_path} in invalid pixel format {img_in.mode}.")
    # Iterate over every pixel in the input image
    for pixel in img_in.getdata():
        # Iterate over the channels in the pixel (R, then G, then B...) then write
        for chan in pixel:
            img_out.write(bytes([chan]))