# Convert and insert .png images into a .psb file.

import os
import shutil
import sys

import bmp2v4
import rle
import psb

def insert_img_psb(out_psb_file, in_psb_file, dir):
	data, is_comp = psb.get_data(in_psb_file)
	psbobj = psb.get_psb(data)
	objectblock = psbobj.objectblock
	texts = psb.get_texts(psbobj)
	names = psb.get_names(psbobj)
	chunks = psb.get_chunks(psbobj)
	expr = psb.get_expr(objectblock, names, texts, chunks)
	indexes = psb.make_index(chunks)
	source = expr['source']
	for sub_dir in source:
		in_dir = os.path.join(dir, sub_dir)
		icon = source[sub_dir]['icon']
		for file in icon:
			in_file = os.path.join(in_dir, file + '.png')
			if not os.path.exists(in_file):
				continue
			image = icon[file]
			chunk = image['pixel']
			index = indexes[chunk]
			shutil.copy(in_file, 'temp.png')
			os.system('png2bmp.bat temp.png')
			bmp_data = open('temp.bmp', 'rb').read()
			os.remove('temp.png')
			os.remove('temp.bmp')
			bin_data = rle.bmp2bin(bmp_data)
			rle_data = rle.bin2rle(bin_data)
			chunks[index] = rle_data
	chunkptrblock, chunkblock, chunklenblock = psb.make_chunks(chunks)
	opsbobj = psb.PSBdata(psbobj.nameblock, psbobj.objectblock, psbobj.textptrblock, psbobj.textblock, chunkptrblock, chunkblock, chunklenblock)
	odata = psb.make_psb(opsbobj)
	psb.write_data(out_psb_file, odata, is_comp)

if __name__ == '__main__':
	if len(sys.argv) < 4:
		print('insert_img_psb.py <in_psb> <in_dir> <out_psb>')
		sys.exit(1)
	in_psb_file = sys.argv[1]
	in_dir = sys.argv[2]
	out_psb_file = sys.argv[3]
	insert_img_psb(out_psb_file, in_psb_file, in_dir)
