# Extract images from a .psb file and convert to .png.

import os
import shutil
import sys

import bmp2v4
import rle
import psb


def extract_img_psb(dir, psb_file):
	expr, names, texts, chunks, is_comp = psb.unpack(psb_file)
	source = expr['source']
	for sub_dir in source:
		out_dir = os.path.join(dir, sub_dir)
		if not os.path.exists(out_dir):
			os.mkdir(out_dir)
		icon = source[sub_dir]['icon']
		for file in icon:
			image = icon[file]
			rle_data = image['pixel']
			width = image['width']
			height = image['height']
			try:
				compress = image['compress']
			except KeyError:
				print('Warning: compress tag not found')
			if compress == 'RL':
				bin_data = rle.rle2bin(rle_data)
				if len(bin_data) != 4*width*height:
					print('Warning: image_size={} mismatch width width={}, height={}'.format(len(bin_data), width, height))
					continue
				bmp_data = rle.bin2bmp(bin_data, width, height)
			else:
				raise Exception('Unsupported compression: {}'.format(compress))
			out_file = os.path.join(out_dir, file + '.bmp')
			open(out_file, 'wb').write(bmp_data)
			shutil.move(out_file, 'temp.bmp')
			bmp2v4.convert('temp.bmp') # convert bmp to v4 format so image-magick can handle alpha channel
			os.system('bmp2png.bat temp.bmp')
			os.remove('temp.bmp')
			shutil.move('temp.png', os.path.splitext(out_file)[0] + '.png')

if __name__ == '__main__':
	if len(sys.argv) < 3:
		print('extract_img_psb.py <psb_file> <out_dir>')
		sys.exit(1)
	psb_file = sys.argv[1]
	out_dir = sys.argv[2]
	if not os.path.exists(out_dir):
		os.mkdir(out_dir)
	extract_img_psb(out_dir, psb_file)
