#!/usr/bin/env python

"""
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.


extractScript.py
================
A (basic, initial) script extractor for the PC-Engine game 'Cyber Knight'.

In order to run, it requires a headerless copy of the Cyber Knight ROM file
and a complete translation table (distributed with this programe).

John Snowdon <john@target-earth.net>
"""

import struct
import binascii

ROM_NAME = "Cyber Knight (J).pce"
TABLE_NAME = "CyberKnightTranslation.csv"
OUT_NAME = "out.sjs"
METHOD_1 = 1
METHOD_2 = 2
METHOD_3 = 3
SWITCH_MODE = '\x5c'

def load_table():
	"""
	load_table - load the translation table.
	The translation table is a tab delimited data file
	with the following columns:
	hex code, actual char pre-0x5c byte, char set type (A/K/H/S), post-0x5c byte, char set type, notes
	
	where A/S/H/K = ASCII, Symbol, Hiragana, Katakana
	pre-0x5c = the character shown if the byte come before a 0x5c control byte
	post-0x5c = the character shown if the byte comes after a 0x5c control byte
	"""
	trans_table = {}
	
	f = open(TABLE_NAME, "r")
	for line in f:
		columns = line.split('\t')
		byte_code = columns[0].replace('"', '')
		trans_table[byte_code] = {}
		trans_table[byte_code]["byte_code"] = byte_code
		trans_table[byte_code]["pre_shift"] = columns[1].replace('"', '')
		trans_table[byte_code]["pre_shift_type"] = columns[2].replace('"', '')
		trans_table[byte_code]["post_shift"] = columns[3].replace('"', '')
		trans_table[byte_code]["post_shift_type"] = columns[4].replace('"', '')
		trans_table[byte_code]["notes"] = columns[5].replace('"', '')
	f.close()
	return trans_table
	

def translate_string(byte_sequence, trans_table):
	"""
	translate_string - construct the actual text, using multi-byte characters
	where appropriate, that represent the hex codes found in the rom.
	e.g. 0x1A 0x5F 0x76 0x61 0x62 0x63 0x64 0x65 0x00 = <control><control>vabcde<end>
	"""

	# method1 has two leading control bytes and a null byte as terminator
	byte_sequence["text"] = []
	if (byte_sequence["method"] == METHOD_1):
		previous_b = ""
		switch_mode = False
		for i in range(2, len(byte_sequence["bytes"]) - 1):
			b = str(binascii.hexlify(byte_sequence["bytes"][i])).upper()
			if switch_mode:
				if b == SWITCH_MODE:
					switch_mode = False
				else:
					if b in trans_table.keys():
						byte_sequence["text"].append(trans_table[b]["post_shift"])
					else:
						# warning - byte sequence not in table
						print "WARNING: Untranslated byte <%s>" % b
						byte_sequence["text"].append("<%s>" % b)
			else:
				if b == SWITCH_MODE:
					switch_mode = True
				else:
					if b in trans_table.keys():
						byte_sequence["text"].append(trans_table[b]["pre_shift"])
					else:
						# warning - byte sequence not in table
						print "WARNING: Untranslated byte <%s>" % b
						byte_sequence["text"].append("<%s>" % b)
	return byte_sequence


def method1(ROM_NAME, rom_start_address, rom_end_address):
	""" 
	method1 - extract text from a given byte range using
	the notation of 2 control bytes, a variable number of
	text bytes and then a single null closing byte.
	e.g. 0x1A 0x2B 0x60 0x61 0x62 0x63 0x64 0x65 0x00
	"""
		
	ttable = load_table()
		
	f = open(ROM_NAME, "rb")
	f.seek(rom_start_address, 0)
	rom_addr = rom_start_address
	
	byte_strings = []
	byte_sequence = {}
	byte_sequence["start_pos"] = rom_addr
	byte_sequence["bytes"] = []
	byte_sequence["size"] = 0
	byte_sequence["method"] = METHOD_1
	
	while (rom_addr <= rom_end_address):
		# Read a byte from the file at the current position
		byte = struct.unpack('c', f.read(1))[0]
		if byte != "\x00":
			# Add the byte
			byte_sequence["bytes"].append(byte)
		else:
			# Add the end byte and record the string
			byte_sequence["bytes"].append(byte)
			byte_sequence["size"] = len(byte_sequence["bytes"])
			
			# Generate the actual text string (which we will print for translation)
			byte_sequence = translate_string(byte_sequence, ttable)
			
			# Record the data
			byte_strings.append(byte_sequence)
			
			# Start a new byte sequence
			byte_sequence = {}
			byte_sequence["start_pos"] = rom_addr
			byte_sequence["bytes"] = []
			byte_sequence["size"] = 0
			byte_sequence["method"] = METHOD_1

		# Increment position ID
		rom_addr += 1
	f.close()
	return byte_strings
	
def method2(ROM_NAME, rom_start_address, rom_end_address):
	"""
	method2 - extract text from a given byte range using
	the notation of each string being wrapped in a single 
	control byte to start, and a single control byte to end.
	e.g. 0x3C 0x60 0x61 0x62 0x63 0x64 0x65 0x04
	"""
	pass

	
byte_strings = method1(ROM_NAME, 82400, 84000)
f = open(OUT_NAME, "w")
for b in byte_strings:
	f.write("Position: %s\n" % b["start_pos"])
	f.write("Method: %s\n" % b["method"])
	f.write("Length: %s\n" % b["size"])	
	for c in b["text"]:
		f.write(c)
	f.write('\n\n')	
f.close()
