import hashlib, io
from PIL import Image
PACKED_IMAGE_LIMIT_BYTES = 21_866_496
SCALES = [1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2]

def sha(b):
    return hashlib.sha256(b).hexdigest()

def png(im):
    b = io.BytesIO()
    im.save(b, format='PNG', compress_level=1, optimize=False)
    return b.getvalue()

def normalized(plan, frames, expected=None):
    factors = [expected['common_geometry']['scale']] if expected else SCALES
    for factor in factors:
        blocks = []
        count = 0
        for block in plan['blocks']:
            assets = []
            for start in range(0, len(block['indices']), 2):
                indices = block['indices'][start:start + 2]
                originals = [frames[i] for i in indices]
                w, h = originals[0].size
                nw = round(w * factor)
                nh = round(h * factor)
                packed = Image.new('RGB', (nw * len(indices), nh))
                rgb = []
                for pos, image in enumerate(originals):
                    if factor != 1:
                        image = image.resize((nw, nh), Image.Resampling.LANCZOS)
                    packed.paste(image, (pos * nw, 0))
                    rgb.append(sha(image.tobytes()))
                blob = png(packed)
                count += len(blob) + (-len(blob) % 3)
                assets.append(dict(indices=indices, size=list(packed.size), sha256=sha(blob), rgb_sha256=rgb, bytes=len(blob), _blob=blob))
            blocks.append(dict(role=block['role'], images=assets))
        if count <= PACKED_IMAGE_LIMIT_BYTES:
            break
    else:
        raise RuntimeError('No common caption input geometry fits')
    if expected:
        for old, new in zip(expected['frame_blocks'], blocks):
            assert old['role'] == new['role'] and len(old['images']) == len(new['images'])
            for a, b in zip(old['images'], new['images']):
                assert a['indices'] == b['indices'] and a['rgb_sha256'] == b['rgb_sha256'] and (a['size'] == b['size'])
    return dict(plan, frame_blocks=blocks, common_geometry=dict(scale=factor, source_pixel_unchanged=factor == 1, aligned_image_bytes=count, image_size_limit_bytes=PACKED_IMAGE_LIMIT_BYTES, sampling_frames_removed=0))

def source_frames(normalized):
    out = {}
    for block in normalized['frame_blocks']:
        for a in block['images']:
            image = Image.open(io.BytesIO(a['_blob'])).convert('RGB')
            n = len(a['indices'])
            w = image.width // n
            for pos, i in enumerate(a['indices']):
                frame = image.crop((pos * w, 0, (pos + 1) * w, image.height))
                assert sha(frame.tobytes()) == a['rgb_sha256'][pos]
                out[i] = frame
    return out