# Routines for working with PSB files

import collections
import pprint
import struct
import sys
import zlib

PSBdata = collections.namedtuple('PSBdata', 'nameblock objectblock textptrblock textblock chunkptrblock chunklenblock chunkblock')

class UnknownBytes(list):
	def __init__(self, type, vals):
		super().__init__(vals)
		self.type = type
		
class Chunk(bytes):
	def __repr__(self):
		return '<Chunk size={}>'.format(len(self))
	def __str__(self):
		return '<Chunk size={}>'.format(len(self))

def get_psb(data):
	offsets = struct.unpack('<IIIIIII', data[0x0C:0x28])
	nameblock = data[offsets[0]:offsets[6]]
	objectblock = data[offsets[6]:offsets[1]]
	textptrblock = data[offsets[1]:offsets[2]]
	textblock = data[offsets[2]:offsets[3]]
	chunkptrblock = data[offsets[3]:offsets[4]]
	chunklenblock = data[offsets[4]:offsets[5]]
	chunkblock = data[offsets[5]:]
	obj = PSBdata(nameblock, objectblock, textptrblock, textblock, chunkptrblock, chunklenblock, chunkblock)
	return obj
	
def modify_psb(obi):
	objlist = list(obj)
	
def make_header(obj):
	header = b''
	header += b'PSB\x00\x02\x00\x00\x00\x00\x00\x00\x00'
	name_offset = 0x28
	object_offset = name_offset + len(obj.nameblock)
	textptr_offset = object_offset + len(obj.objectblock)
	text_offset = textptr_offset + len(obj.textptrblock)
	chunkptr_offset = text_offset + len(obj.textblock)
	chunklen_offset = chunkptr_offset + len(obj.chunkptrblock)
	chunk_offset = chunklen_offset + len(obj.chunklenblock)
	header += struct.pack('<IIIIIII', name_offset, textptr_offset, text_offset, chunkptr_offset, chunklen_offset, chunk_offset, object_offset)
	return header
	
def make_psb(obj):
	blocks = []
	header = make_header(obj)
	blocks.append(header)
	blocks.append(obj.nameblock)
	blocks.append(obj.objectblock)
	blocks.append(obj.textptrblock)
	blocks.append(obj.textblock)
	blocks.append(obj.chunkptrblock)
	blocks.append(obj.chunklenblock)
	blocks.append(obj.chunkblock)
	return b''.join(blocks)

def get_byte(data, pos):
	val = data[pos]
	return val, pos+1 
	
def get_word(data, pos):
	val = struct.unpack('<H', data[pos:pos+2])[0]
	return val, pos+2
	
def get_tword(data, pos):
	val_lower, val_upper = struct.unpack('<HB', data[pos:pos+3])
	val = val_lower | (val_upper << 16)
	return val, pos+3
	
def get_dword(data, pos):
	val = struct.unpack('<I', data[pos:pos+4])[0]
	return val, pos+4
	
def get_float(data, pos):
	val = struct.unpack('<f', data[pos:pos+4])[0]
	return val, pos+4
	
def get_double(data, pos):
	val = struct.unpack('<d', data[pos:pos+8])[0]
	return val, pos+8
	
def get_bytes(data, pos, n):
	return [get_byte(data, pos+k)[0] for k in range(n)], pos+n
	
def get_words(data, pos, n):
	return [get_word(data, pos+2*k)[0] for k in range(n)], pos+2*n
	
def get_twords(data, pos, n):
	return [get_tword(data, pos+3*k)[0] for k in range(n)], pos+3*n
	
def get_dwords(data, pos, n):
	return [get_dword(data, pos+4*k)[0] for k in range(n)], pos+4*n
	
def get_string(data, pos):
	endpos = data.find(b'\x00', pos)
	block = data[pos:endpos]
	pos = len(block) + 1
	text = block.decode('utf-8')
	return text
	
def get_chunk(data, pos, length):
	chunk = Chunk(data[pos:pos+length])
	return chunk
	
def get_name(A, B, C, index):
	chars = []
	a = C[index]
	b = B[a]
	if not b:
		return ''
	while 1:
		c = B[b]
		d = A[c]
		e = b - d
		chars.append(bytes([e]))
		b = c
		if not b:
			break
	chars.reverse()
	return b''.join(chars).decode('utf-8')
	
type_table = {
0x0D: ('{:02X}', get_byte, get_bytes),
0x0E: ('{:04X}', get_word, get_words),
0x0F: ('{:06X}', get_tword, get_twords),
0x10: ('{:08X}', get_dword, get_dwords),
}
	
def get_array(data, pos):
	type1, pos = get_byte(data, pos)
	fmt1, get_fcn1, gets_fcn1 = type_table[type1]
	nvals, pos = get_fcn1(data, pos)
	type2, pos = get_byte(data, pos)
	fmt2, get_fcn2, gets_fcn2 = type_table[type2]
	vals, pos = gets_fcn2(data, pos, nvals)
	return type1, nvals, type2, vals, pos
	
def vals_string(vals, fmt):
	return ' '.join(fmt.format(x) for x in vals)
	
def array_string(type1, nvals, type2, vals):
	fmt1 = type_table[type1][0]
	fmt2 = type_table[type2][0]
	s1 = fmt1.format(nvals)
	s2 = vals_string(vals, fmt2)
	return '{:02X} {} {:02X} {}'.format(type1, s1, type2, s2)
	
def get_data(psb_file):
	is_comp = False
	data = open(psb_file, 'rb').read()
	if data[:4] == b'mdf\x00':
		odata = zlib.decompress(data[8:])
		is_comp = True
	elif data[:4] == b'PSB\x00':
		odata = data
		is_comp = False
	return odata, is_comp
	
def write_data(psb_file, odata, is_comp):
	fo = open(psb_file, 'wb')
	if is_comp:
		fo.write(b'mdf\x00')
		fo.write(struct.pack('<I', len(odata)))
		fo.write(zlib.compress(odata))
	else:
		fo.write(odata)
	fo.close()
	
def get_texts(obj):
	texts = []
	type1, nvals, type2, textptrs, pos = get_array(obj.textptrblock, 0)
	textblock = obj.textblock
	for textptr in textptrs:
		texts.append(get_string(textblock, textptr))
	return texts
	
def get_names(obj):
	names = []
	type1, nvals, type2, A, pos = get_array(obj.nameblock, 0)
	type1, nvals, type2, B, pos = get_array(obj.nameblock, pos)
	type1, nvals, type2, C, pos = get_array(obj.nameblock, pos)
	for k in range(nvals):
		names.append(get_name(A, B, C, k))
	return names
	
def get_chunks(obj):
	chunks = []
	type1, nvals, type2, ptrs, pos = get_array(obj.chunkptrblock, 0)
	type1, nvals, type2, lengths, pos = get_array(obj.chunklenblock, 0)
	for ptr, length in zip(ptrs, lengths):
		chunk = get_chunk(obj.chunkblock, ptr, length)
		chunks.append(chunk)
	return chunks
		
def get_expr(data, names, texts, chunks, pos=0):
	base = pos
	type, pos = get_byte(data, pos)
	if type == 0x01:
		return None
	elif type == 0x04:
		return 0
	elif type == 0x05:
		val, pos = get_byte(data, pos)
		return val
	elif type == 0x06:
		val, pos = get_word(data, pos)
		return val
	elif type == 0x07:
		val, pos = get_tword(data, pos)
		return val
	elif type == 0x08:
		val, pos = get_dword(data, pos)
		return val
	elif type == 0x09:
		vals, pos = get_bytes(data, pos, 5)
		return UnknownBytes(type, vals)
	elif type == 0x15:
		val, pos = get_byte(data, pos)
		text = texts[val]
		return text
	elif type == 0x16:
		val, pos = get_word(data, pos)
		text = texts[val]
		return text
	elif type == 0x17:
		val, pos = get_tword(data, pos)
		text = texts[val]
		return text
	elif type == 0x18:
		val, pos = get_cword(data, pos)
		text = texts[val]
		return text
	elif type == 0x19:
		val, pos = get_byte(data, pos)
		return chunks[val]
	elif type == 0x1A:
		val, pos = get_word(data, pos)
		return chunks[val]
	elif type == 0x1B:
		val, pos = get_tword(data, pos)
		return chunks[val]
	elif type == 0x1C:
		val, pos = get_dword(data, pos)
		return chunks[val]
	elif type == 0x1D:
		return 0.0
	elif type == 0x1E:
		val, pos = get_float(data, pos)
		return val
	elif type == 0x1F:
		val, pos = get_double(data, pos)
		return val
	elif type == 0x20:
		type1, nvals, type2, vals, pos = get_array(data, pos)
		a = []
		for k in range(nvals):
			expr = get_expr(data, names, texts, chunks, pos+vals[k])
			a.append(expr)
		return a
	elif type == 0x21:
		type11, nvals1, type21, vals1, pos = get_array(data, pos)
		type12, nvals2, type22, vals2, pos = get_array(data, pos)
		d = {}
		for k in range(nvals2):
			name = names[vals1[k]]
			expr = get_expr(data, names, texts, chunks, pos+vals2[k])
			d[name] = expr
		return d
	else:
		raise Exception('Unknown type {:02X} @ {:06X}'.format(type, base))
		
def make_data(expr):
	names = set()
	texts = set()
	chunks = set()	
	def make_data_r(expr):
		if isinstance(expr, str):
			texts.add(expr)
		elif isinstance(expr, bytes):
			chunks.add(expr)
		elif isinstance(expr, list):
			for entry in expr:
				make_data_r(entry)
		elif isinstance(expr, dict):
			for key in expr:
				names.add(key)
				make_data_r(expr[key])
	make_data_r(expr)
	return sorted(names), sorted(texts), sorted(chunks)
	
def make_index(data):
	index = {}
	for id, entry in enumerate(data):
		index[entry] = id
	return index
	
def pack_none():
	return b'\x01', 1
	
def pack_int(i):
	if i == 0:
		return b'\x04', 1
	elif i < 0x100:
		return b'\x05' + struct.pack('B', i), 2
	elif i < 0x10000:
		return b'\x06' + struct.pack('<H', i), 3
	elif i < 0x1000000:
		return b'\x07' + struct.pack('<HB', i&0xFFFF, i>>16), 4
	else:
		return b'\x08' + struct.pack('<I', i), 5
		
def pack_float(f):
	return b'\x1F' + struct.pack('<d', f), 9
	
def pack_unknown(c):
	return struct.pack('B', c.type) + struct.pack('B'*len(c), *tuple(c)), len(c)+1
	
def pack_str(s, texts_index):
	i = texts_index[s]
	if i < 0x100:
		return b'\x15' + struct.pack('B', i), 2
	elif i < 0x10000:
		return b'\x16' + struct.pack('<H', i), 3
	elif i < 0x1000000:
		return b'\x17' + struct.pack('<HB', i&0xFFFF, i>>16), 4
	else:
		return b'\x18' + struct.pack('<I', i), 5
		
def pack_chunk(s, chunks_index):
	i = chunks_index[s]
	if i < 0x100:
		return b'\x19' + struct.pack('B', i), 2
	elif i < 0x10000:
		return b'\x1A' + struct.pack('<H', i), 3
	elif i < 0x1000000:
		return b'\x1B' + struct.pack('<HB', i&0xFFFF, i>>16), 4
	else:
		return b'\x1C' + struct.pack('<I', i), 5

def pack_array(vals):
	i = len(vals)
	block = b''
	if i < 0x100:
		block += b'\x0D' + struct.pack('B', i)
	elif i < 0x10000:
		block += b'\x0E' + struct.pack('<H', i)
	elif i < 0x1000000:
		block += b'\x0F' + struct.pack('<HB', i&0xFFFF, i>>16)
	else:
		block += b'\x10' + struct.pack('<I', i)
	if i > 0:
		m = max(vals)
		if m < 0x100:
			block += b'\x0D' + struct.pack('B'*i, *tuple(vals))
		elif m < 0x10000:
			block += b'\x0E' + struct.pack('<'+'H'*i, *tuple(vals))
		elif m < 0x1000000:
			lowers = [x&0xFFFF for x in vals]
			uppers = [x>>16 for x in vals]
			all = [x for y in zip(lowers,uppers) for x in y]
			block += b'\x0F' + struct.pack('<'+'HB'*i, *tuple(all))
		else:
			block += b'\x10' + struct.pack('<'+'I'*i, *tuple(vals))
	else:
		block += b'\x0D'
	return block

		
def pack_list(vals):
	block = b'\x20' + pack_array(vals)
	return block, len(block)
	
def pack_object(names, vals, names_index):
	block = b'\x21'
	namevals = [names_index[name] for name in names]
	block += pack_array(namevals)
	block += pack_array(vals)
	return block, len(block)
	
def make_texts(texts):
	ptrs = []
	blocks = []
	pos = 0
	for text in texts:
		block = text.encode('utf-8') + b'\x00'
		ptrs.append(pos)
		blocks.append(block)
		pos += len(block)
	textptrblock = pack_array(ptrs)
	textblock = b''.join(blocks)
	return textptrblock, textblock
	
def make_chunks(chunks):
	ptrs = []
	lengths = []
	blocks = []
	pos = 0
	for chunk in chunks:
		length = len(chunk)
		ptrs.append(pos)
		lengths.append(length)
		blocks.append(chunk)
		pos += length
	chunkptrblock = pack_array(ptrs)
	chunklenblock = pack_array(lengths)
	chunkblock = b''.join(blocks)
	return chunkptrblock, chunklenblock, chunkblock
		
def pack_expr(expr, names=None, texts=None, chunks=None):
	m_names, m_texts, m_chunks = make_data(expr)
	if names is None:
		names = m_names
	if texts is None:
		texts = m_texts
	if chunks is None:
		chunks = m_chunks
	names_index = make_index(names)
	texts_index = make_index(texts)
	chunks_index = make_index(chunks)
	pos = 0
	def pack_expr_r(expr):
		if expr is None:
			return pack_none()
		elif isinstance(expr, int):
			return pack_int(expr)
		elif isinstance(expr, float):
			return pack_float(expr)
		elif isinstance(expr, UnknownBytes):
			return pack_unknown(expr)
		elif isinstance(expr, str):
			return pack_str(expr, texts_index)
		elif isinstance(expr, Chunk):
			return pack_chunk(expr, chunks_index)
		elif isinstance(expr, list):
			blocks = []
			deltas = []
			for entry in expr:
				block, delta = pack_expr_r(entry)
				blocks.append(block)
				deltas.append(delta)
			vals = [sum(deltas[:k]) for k in range(len(deltas))]
			packed_list, list_delta = pack_list(vals)
			return packed_list+b''.join(blocks), list_delta+sum(deltas)
		elif isinstance(expr, dict):
			blocks = []
			deltas = []
			for name in sorted(expr):
				entry = expr[name]
				block, delta = pack_expr_r(entry)
				blocks.append(block)
				deltas.append(delta)
			vals = [sum(deltas[:k]) for k in range(len(deltas))]
			packed_object, object_delta = pack_object(sorted(expr), vals, names_index)
			return packed_object+b''.join(blocks), object_delta+sum(deltas)
		else:
			raise Exception('Unhandled type: ' + str(type(expr)))
	return pack_expr_r(expr)[0], names, texts, chunks
	
def unpack(psb_file):
	data, is_comp = get_data(psb_file)
	obj = get_psb(data)
	objectblock = obj.objectblock
	texts = get_texts(obj)
	names = get_names(obj)
	chunks = get_chunks(obj)
	expr = get_expr(objectblock, names, texts, chunks)
	return expr, names, texts, chunks, is_comp
		
def dump(outfile, expr):
	s = pprint.pformat(expr)
	open(outfile, 'w', encoding='utf-8').write(s)
	
def bytestr(data):
	return ' '.join('{:02X}'.format(x) for x in data)

if __name__ == '__main__':
	import pdb
	sys.setrecursionlimit(5000)
	if len(sys.argv) < 2:
		print('Usage: psb.py <psb_file>')
		sys.exit(1)
	psb_file = sys.argv[1]
	data, is_comp = get_data(psb_file)
	obj = get_psb(data)
	objectblock = obj.objectblock
	texts = get_texts(obj)
	names = get_names(obj)
	chunks = get_chunks(obj)
	expr = get_expr(objectblock, names, texts, chunks)
	pdb.set_trace()
