# viewer.py
import sys
import struct
import json
import pygame
import numpy as np
import codec  # Import our codec module

def parse_header_and_metadata(data):
    """
    Parses the fixed header and metadata from the SIMIMG file.
    
    Fixed header (17 bytes):
        - Bytes 0-5: MAGIC (e.g. b"SIMIMG")
        - Byte 6   : VERSION
        - Bytes 7-10  : Width (little-endian unsigned int)
        - Bytes 11-14 : Height (little-endian unsigned int)
        - Byte 15  : Mode (e.g., 0x01 for 24-bit RGB)
        - Byte 16  : Compression flag (e.g., 0x02 indicates Custom Compression)
    
    Then a 4-byte little-endian integer indicates the length of the metadata
    block, followed by a JSON string containing metadata.
    
    Returns:
        (header, metadata): Two dictionaries containing the fixed header info and
                            the metadata.
    """
    header = {}
    header['magic'] = data[0:6].decode('ascii', errors='replace')
    header['version'] = data[6]
    header['width'] = struct.unpack("<I", data[7:11])[0]
    header['height'] = struct.unpack("<I", data[11:15])[0]
    mode_val = data[15]
    header['mode'] = "24-bit RGB" if mode_val == 1 else f"Unknown (0x{mode_val:02X})"
    comp_flag = data[16]
    if comp_flag == 1:
        header['compression'] = "ZLIB"
    elif comp_flag == 2:
        header['compression'] = "Custom"
    else:
        header['compression'] = "None"
    
    # Metadata block follows after the fixed 17 bytes.
    # First 4 bytes represent metadata length.
    meta_length = struct.unpack("<I", data[17:21])[0]
    meta_bytes = data[21:21+meta_length]
    try:
        metadata = json.loads(meta_bytes.decode("utf-8"))
    except Exception as e:
        metadata = {"error": str(e)}
    
    return header, metadata

def view_image(file_path):
    """
    Loads a SIMIMG file, decodes it, and displays it using Pygame.
    Overlays all header and metadata information in the top left corner.
    """
    # Read the entire file.
    with open(file_path, 'rb') as f:
        data = f.read()
    
    # Parse header and metadata.
    header_info, metadata = parse_header_and_metadata(data)
    
    # Decode the image using codec.
    image = codec.decode(data)
    height = len(image)
    width = len(image[0]) if height > 0 else 0
    
    # Build a NumPy array for the pixel data.
    array_img = np.zeros((height, width, 3), dtype=np.uint8)
    for y in range(height):
        for x in range(width):
            array_img[y, x] = image[y][x]
    
    # Initialize Pygame.
    pygame.init()
    screen = pygame.display.set_mode((width, height))
    pygame.display.set_caption("SIMIMG Viewer")
    
    # Create a Pygame surface from the image data.
    surface = pygame.image.frombuffer(array_img.tobytes(), (width, height), "RGB")
    
    # Create an overlay containing all header and metadata info.
    font = pygame.font.SysFont(None, 20)
    overlay_lines = []
    # Add header info.
    for key, value in header_info.items():
        overlay_lines.append(f"{key}: {value}")
    # Add metadata info.
    for key, value in metadata.items():
        overlay_lines.append(f"{key}: {value}")
    
    # Render text surfaces.
    info_surfaces = [font.render(line, True, (255, 255, 255)) for line in overlay_lines]
    padding = 5
    panel_width = max(text.get_width() for text in info_surfaces) + 2 * padding
    panel_height = sum(text.get_height() for text in info_surfaces) + (len(info_surfaces) - 1) * 2 + 2 * padding
    
    overlay = pygame.Surface((panel_width, panel_height), pygame.SRCALPHA)
    overlay.fill((0, 0, 0, 150))  # Semi-transparent black background.
    
    current_y = padding
    for text in info_surfaces:
        overlay.blit(text, (padding, current_y))
        current_y += text.get_height() + 2  # 2 pixels spacing
    
    clock = pygame.time.Clock()
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
        
        # Draw the image.
        screen.blit(surface, (0, 0))
        # Blit the overlay with all metadata in the top left corner.
        screen.blit(overlay, (0, 0))
        pygame.display.flip()
        clock.tick(30)
    
    pygame.quit()

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: python viewer.py <image_file.simimg>")
        sys.exit(1)
    
    file_path = sys.argv[1]
    view_image(file_path)