#!/usr/bin/env python
# -*- coding: utf-8 -*-

# By HANSEN
import sys
import os
import array
import struct

ROMNAME = '1046 - Megaman & Bass (E).gba'
CFGFILE = 'mega.ini'
INIT = 'CompOffset=0x794870\n'

def encode():
	#
	# Rotina para SIMULAR a compressão!!
	#
	n = raw_input('>> Numero: ')
	if n.isdigit():
		n = int(n)
	else:
		print '<< Erro'
		return
	
	cfg = open(CFGFILE, 'r+b')
	opt = cfg.readlines().pop()	
	add = int(opt.split('=')[1], 16)
	cfg.close()
	
	fd = open(ROMNAME, 'r+b')
	fd.seek(add, 0)
	
	raw = open('raw/%03d.gba' % n, 'rb')
	
	coded = array.array('c')
	for x in range(os.path.getsize(raw.name) / 8):
		coded.append('\x00')
		for y in range(8):
			coded.append(raw.read(1))
	raw.close()
			
	coded.tofile(fd)
	nopt = fd.tell()
	while nopt % 4 != 0:
		nopt += 1
	
	cfg = open(CFGFILE, 'wb')
	cfg.write('CompOffset=0x%06X' % nopt)
	cfg.close()
		
	#GERAR NO ARQUIVO
	header_address = 0x5079D8 + n * 8
	base_address = 0x64BA38
	fd.seek(header_address)
	address = add - base_address	
	fd.write(struct.pack('<L', address))
	fd.write(struct.pack('<L', os.path.getsize(raw.name)))
	
	fd.close()
	return

def decode(image_number = 0):
	fd = open(ROMNAME, 'rb')
	while True:
		try:
			output = array.array('c')

			base_address = 0x64BA38	
			header_address = 0x5079D8 + image_number * 8

			fd.seek(header_address, 0)
			data_address = struct.unpack('<L', fd.read(4))[0] + base_address
			data_size = struct.unpack('<L', fd.read(4))[0]
			
			if data_address == 0x6f95c4:
				print image_number
			
			fd.seek(data_address, 0)

			while len(output) < data_size:
				flag = struct.unpack('B', fd.read(1))[0]
				for x in range(8):
					if ((flag << x) & 0x80) == 0:
						output.append(fd.read(1))
					else:
						par1 = struct.unpack('B', fd.read(1))[0]
						par2 = struct.unpack('B', fd.read(1))[0]
						size = par1 >> 2
						distance = ((par1 & 3) << 8) | par2	
						for x in range(size):
							output.append(output[len(output)  - distance])
					if len(output) == data_size: break
							
			outputf = open('raw/%03d.gba' % image_number, 'wb')
			output.tofile(outputf)
			outputf.close()
		except:
			fd.close()
			return
			
		image_number += 1

if __name__ == '__main__':
	if ROMNAME not in os.listdir(''):
		print '<<< 1046 - Megaman & Bass (E).gba não foi encontrado no diretório raiz.'.decode('utf-8')
		raw_input()
		sys.exit(1)
		
	if CFGFILE not in os.listdir(''):
		cfg = open(CFGFILE, 'wb')
		cfg.write(INIT)
		cfg.close()
		
	if 'raw' not in os.listdir(''):
		os.mkdir('raw')
		
	print 'MegaEncoder - by Hansen\n'.decode('utf-8')
	print '>> Opções:\n\t1.Extrair dados\n\t2.Inserir dados\n\t3.Sair'.decode('utf-8')
	while True:
		task = raw_input('\n>> Escolha: ')
		if task == '1':
			print '>> Extraindo... Tenha paciência!'.decode('utf-8')
			decode()
			print '>> Todos os arquivos extraidos com sucesso!'.decode('utf-8')
			sys.exit(1)
		elif task == '2':
			encode()
			sys.exit(1)
		elif task == '3':
			sys.exit(1)
