# codec.py
import struct
import zlib
import time
import json

MAGIC = b"SIMIMG"       # 6-byte file signature.
VERSION = b'\x01'       # 1-byte version.
MODE_RGB  = b'\x01'     # For 24-bit RGB (3 bytes per pixel).
MODE_RGBA = b'\x02'     # For 32-bit RGBA (4 bytes per pixel).

# Constants for the compression flag.
NO_COMPRESSION  = b'\x00'
COMPRESSION_ZLIB = b'\x01'  # Indicates that zlib compression is used.

def encode(image_pixels, use_zlib=True):
    """
    Encodes an image (2D list of (R, G, B) tuples OR (R, G, B, A) tuples) into our custom SIMIMG format.
    It first applies an RLE-based compression to the pixel data and then optionally compresses the RLE data using zlib.
    
    The file format now includes a variable-length metadata block stored in JSON.
    
    Fixed header layout (17 bytes total):
        MAGIC       (6 bytes)
        VERSION     (1 byte)
        Width       (4 bytes, little-endian)
        Height      (4 bytes, little-endian)
        Mode        (1 byte) -> MODE_RGB (0x01) for 24-bit RGB OR MODE_RGBA (0x02) for 32-bit RGBA.
        Compression (1 byte) -> 0x01 indicates zlib compression, 0x00 indicates none

    Followed by:
        Metadata length (4 bytes, little-endian)
        Metadata JSON (variable length)

    Then:
        Body: The RLE-compressed (and optionally zlib-compressed) pixel data.
    
    The metadata stored includes:
        - encoder: Identifier string.
        - timestamp: Unix epoch time when the image was encoded.
        - original_rle_size: Size of the RLE data (before zlib compression).
        - compressed_body_size: Size of the final body after compression.
        - compression_ratio: Ratio = compressed_body_size / original_rle_size.
        - width, height, mode, and compression type.

    Returns:
        A bytes object containing the full encoded image with metadata.
    """
    height = len(image_pixels)
    width  = len(image_pixels[0]) if height > 0 else 0

    # Determine the pixel format by checking the first pixel.
    if height > 0 and width > 0:
        first_pixel = image_pixels[0][0]
        if len(first_pixel) == 4:
            mode = MODE_RGBA
            mode_str = "32-bit RGBA"
        else:
            mode = MODE_RGB
            mode_str = "24-bit RGB"
    else:
        mode = MODE_RGB
        mode_str = "24-bit RGB"

    # Determine the compression flag.
    compression_flag = COMPRESSION_ZLIB if use_zlib else NO_COMPRESSION

    # Build the fixed header (17 bytes total).
    fixed_header = (
        MAGIC +
        VERSION +
        struct.pack("<I", width) +
        struct.pack("<I", height) +
        mode +
        compression_flag
    )

    # Flatten the 2D list of pixels (row-major order).
    flat_pixels = [pixel for row in image_pixels for pixel in row]
    n = len(flat_pixels)

    # Build RLE compressed body.
    rle_data = bytearray()
    literal_buffer = []
    i = 0
    while i < n:
        run_length = 1
        # Check for a run (up to maximum of 128 consecutive identical pixels).
        while (i + run_length < n and
               flat_pixels[i + run_length] == flat_pixels[i] and
               run_length < 128):
            run_length += 1

        if run_length >= 2:
            if literal_buffer:
                control_byte = len(literal_buffer)  # Literal count (range 1-127).
                rle_data.append(control_byte)
                for pix in literal_buffer:
                    rle_data.extend(bytes(pix))
                literal_buffer = []
            # For a run block, control = 128 + (run_length - 1).
            control = 128 + (run_length - 1)
            rle_data.append(control)
            rle_data.extend(bytes(flat_pixels[i]))
            i += run_length
        else:
            literal_buffer.append(flat_pixels[i])
            i += 1
            if len(literal_buffer) == 128:
                control_byte = len(literal_buffer)
                rle_data.append(control_byte)
                for pix in literal_buffer:
                    rle_data.extend(bytes(pix))
                literal_buffer = []

    # Flush any remaining literal pixels.
    if literal_buffer:
        control_byte = len(literal_buffer)
        rle_data.append(control_byte)
        for pix in literal_buffer:
            rle_data.extend(bytes(pix))

    original_rle_size = len(rle_data)

    # Optionally compress the RLE data with zlib.
    if use_zlib:
        compressed_body = zlib.compress(bytes(rle_data))
    else:
        compressed_body = bytes(rle_data)
    compressed_body_size = len(compressed_body)
    
    compression_ratio = (compressed_body_size / original_rle_size if original_rle_size > 0 else 1.0)
    
    # Build metadata dictionary.
    metadata = {
        "encoder": "SimImg Codec v1.0",
        "timestamp": time.time(),
        "original_rle_size": original_rle_size,
        "compressed_body_size": compressed_body_size,
        "compression_ratio": compression_ratio,
        "width": width,
        "height": height,
        "mode": mode_str,
        "compression": "ZLIB" if use_zlib else "None"
    }
    metadata_json = json.dumps(metadata)
    metadata_bytes = metadata_json.encode("utf-8")
    metadata_length = len(metadata_bytes)
    
    # Pack metadata length as 4 bytes little-endian.
    metadata_header = struct.pack("<I", metadata_length)
    
    # Final file layout: fixed header + metadata length + metadata + body.
    return fixed_header + metadata_header + metadata_bytes + compressed_body

def decode(data):
    """
    Decodes data from our custom SIMIMG image format, which includes a metadata block.
    If the compression flag is set, the body is decompressed using zlib.
    
    Returns:
        A 2D list (rows) of pixels, where each pixel is a (R, G, B) tuple for RGB images
        or a (R, G, B, A) tuple for RGBA images.
    
    Raises:
        ValueError if the file format is incorrect.
    """
    # Verify the fixed header.
    if not data.startswith(MAGIC):
        raise ValueError("Invalid file format: missing SIMIMG signature.")
    
    # Read fixed header (17 bytes total).
    offset = len(MAGIC) + 1  # Skip MAGIC (6 bytes) and VERSION (1 byte).
    width = struct.unpack("<I", data[offset:offset+4])[0]
    offset += 4
    height = struct.unpack("<I", data[offset:offset+4])[0]
    offset += 4
    mode = data[offset]
    offset += 1
    # Determine pixel size from mode.
    if mode == MODE_RGB[0]:
        pixel_size = 3
        mode_str = "24-bit RGB"
    elif mode == MODE_RGBA[0]:
        pixel_size = 4
        mode_str = "32-bit RGBA"
    else:
        raise ValueError(f"Unsupported mode in image file: 0x{mode:02X}")
    compression_flag = data[offset]
    offset += 1

    # Read metadata length (4 bytes, little-endian).
    metadata_length = struct.unpack("<I", data[offset:offset+4])[0]
    offset += 4
    # Read metadata.
    metadata_bytes = data[offset:offset+metadata_length]
    offset += metadata_length
    try:
        metadata = json.loads(metadata_bytes.decode("utf-8"))
    except json.JSONDecodeError:
        metadata = {}
    
    # The remainder is the body.
    body = data[offset:]
    
    # If zlib compression is used, decompress the body.
    if compression_flag == COMPRESSION_ZLIB[0]:
        try:
            rle_data = zlib.decompress(body)
        except zlib.error as e:
            raise ValueError("Decompression failed: " + str(e))
    else:
        rle_data = body

    # Decode the RLE data.
    flat_pixels = []
    rle_offset = 0
    rle_length = len(rle_data)
    while rle_offset < rle_length:
        control = rle_data[rle_offset]
        rle_offset += 1
        if control >= 128:
            run_length = control - 128 + 1
            pixel = tuple(rle_data[rle_offset:rle_offset+pixel_size])
            rle_offset += pixel_size
            flat_pixels.extend([pixel] * run_length)
        else:
            literal_count = control
            for _ in range(literal_count):
                pixel = tuple(rle_data[rle_offset:rle_offset+pixel_size])
                rle_offset += pixel_size
                flat_pixels.append(pixel)
    
    if width * height != len(flat_pixels):
        raise ValueError("Pixel data does not match width/height in header.")
    
    # Reconstruct a 2D list (row-major).
    image = []
    for y in range(height):
        row = flat_pixels[y * width : (y + 1) * width]
        image.append(row)
    
    return image

# Aliases for clarity.
compress = encode
decompress = decode

if __name__ == '__main__':
    # Test: Create and round-trip a small image.
    # Here we test with a simple RGB image.
    test_image_rgb = [
        [(255, 0, 0)] * 10,
        [(0, 255, 0)] * 10,
        [(0, 0, 255)] * 10,
    ]
    encoded_rgb = encode(test_image_rgb, use_zlib=True)
    decoded_rgb = decode(encoded_rgb)
    assert decoded_rgb == test_image_rgb, "Round-trip encoding/decoding failed for RGB."
    
    # Test with a simple RGBA image (with transparency).
    test_image_rgba = [
        [(255, 0, 0, 128)] * 10,
        [(0, 255, 0, 128)] * 10,
        [(0, 0, 255, 128)] * 10,
    ]
    encoded_rgba = encode(test_image_rgba, use_zlib=True)
    decoded_rgba = decode(encoded_rgba)
    assert decoded_rgba == test_image_rgba, "Round-trip encoding/decoding failed for RGBA."
    
    print("Codec test passed.")