from PIL import Image
import os

WIDTH, HEIGHT = 800, 480

COLORS = {
    0x0: (255, 255, 255),  # TFT_WHITE
    0xF: (0, 0, 0),        # TFT_BLACK
    0xD: (0, 0, 255),      # TFT_BLUE
    0xB: (255, 255, 0),    # TFT_YELLOW
    0x2: (0, 255, 0),      # TFT_GREEN
    0x6: (255, 0, 0),      # TFT_RED
}

COLOR_LIST = list(COLORS.items())

def nearest_color(r, g, b):
    best_idx = 0
    best_dist = float("inf")
    for idx, (cr, cg, cb) in COLOR_LIST:
        d = (r - cr) ** 2 + (g - cg) ** 2 + (b - cb) ** 2
        if d < best_dist:
            best_dist = d
            best_idx = idx
    return best_idx

def convert_image(filepath):
    img = Image.open(filepath).convert("RGB")
    # source is 480x800 portrait → rotate 90° CW → 800x480 landscape
    if img.width < img.height:
        img = img.rotate(90, expand=True)
    img = img.resize((WIDTH, HEIGHT), Image.LANCZOS)
    pixels = list(img.getdata())  # keep until Pillow 14 removes it
    data = bytearray()
    for i in range(0, len(pixels), 2):
        p1 = nearest_color(*pixels[i])
        if i + 1 < len(pixels):
            p2 = nearest_color(*pixels[i + 1])
        else:
            p2 = 0x0
        data.append((p1 << 4) | p2)
    return bytes(data)

def main():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    posters = []
    for fname in sorted(os.listdir(script_dir)):
        if fname.lower().startswith("poster") and fname.lower().endswith(".png"):
            posters.append(fname)
    posters.sort()

    lines = []
    lines.append("#ifndef IMAGES_H")
    lines.append("#define IMAGES_H")
    lines.append("")
    lines.append("#ifdef USE_COLORFULL_EPAPER")
    lines.append("")

    for fname in posters:
        var_name = "gImage_" + os.path.splitext(fname)[0]
        filepath = os.path.join(script_dir, fname)
        print(f"Converting {fname}...")
        data = convert_image(filepath)
        hex_str = "  " + ", ".join(f"0x{b:02X}" for b in data)
        lines.append(f"// {fname} ({WIDTH}x{HEIGHT})")
        lines.append(f"const unsigned char {var_name}[{len(data)}] = {{")
        lines.append(hex_str)
        lines.append("};")
        lines.append("")

    lines.append("#endif // USE_COLORFULL_EPAPER")
    lines.append("#endif // IMAGES_H")

    out_path = os.path.join(script_dir, "images.h")
    with open(out_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines) + "\n")
    print(f"\nDone! Written to {out_path}")

if __name__ == "__main__":
    main()