import fitz
from PIL import Image
from io import BytesIO

def create_test_pdf():
    doc = fitz.open()
    p = doc.new_page()
    img = Image.new('RGBA', (500, 500), color=(255, 0, 0, 128))
    buf = BytesIO()
    img.save(buf, format='PNG')
    p.insert_image(fitz.Rect(0,0,100,100), stream=buf.getvalue())
    doc.save("test_img.pdf")
    doc.close()

def compress():
    doc = fitz.open("test_img.pdf")
    img_info = doc[0].get_images()[0]
    xref = img_info[0]
    smask = img_info[1]
    
    base_image = doc.extract_image(xref)
    image_bytes = base_image["image"]
    pil_img = Image.open(BytesIO(image_bytes))
    
    # Resize
    pil_img = pil_img.resize((250, 250), Image.Resampling.LANCZOS)
    
    if pil_img.mode == 'RGBA' and smask > 0:
        print("Handling SMask!")
        r, g, b, a = pil_img.split()
        rgb_img = Image.merge("RGB", (r, g, b))
        
        rgb_io = BytesIO()
        rgb_img.save(rgb_io, format='JPEG', quality=30)
        
        a_io = BytesIO()
        a.save(a_io, format='JPEG', quality=30) # Grayscale JPEG for mask
        
        # Update main image
        doc.update_stream(xref, rgb_io.getvalue(), compress=0)
        doc.xref_set_key(xref, "Width", "250")
        doc.xref_set_key(xref, "Height", "250")
        doc.xref_set_key(xref, "Filter", "/DCTDecode")
        doc.xref_set_key(xref, "ColorSpace", "/DeviceRGB")
        doc.xref_set_key(xref, "DecodeParms", "null")
        
        # Update smask
        doc.update_stream(smask, a_io.getvalue(), compress=0)
        doc.xref_set_key(smask, "Width", "250")
        doc.xref_set_key(smask, "Height", "250")
        doc.xref_set_key(smask, "Filter", "/DCTDecode")
        doc.xref_set_key(smask, "ColorSpace", "/DeviceGray")
        doc.xref_set_key(smask, "DecodeParms", "null")
    
    doc.save("test_img_compressed.pdf", garbage=4, deflate=True)
    doc.close()

create_test_pdf()
compress()