untrans

#!/usr/bin/env python3

#
# UnTrans - Extract all translatable text and graphics of Wild Arms to text and image files
#
# Copyright (C) Christian Bauer <www.cebix.net>
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#

__version__ = "1.0"

import sys
import os
import struct
import io
import shutil

sys.stdout.reconfigure(encoding = "locale", errors = "backslashreplace")
sys.stderr.reconfigure(encoding = "locale", errors = "backslashreplace")

from PIL import Image
from PIL import ImagePalette

import wa
from wa.map import Op


# Save a list of unicode strings to a UTF-8 file in the text output directory.
def saveTrans(transPath, subDir, fileName, lines):

    # Create the output directory if necessary
    outputDir = os.path.join(transPath, subDir)
    if not os.path.isdir(outputDir):
        os.mkdir(outputDir)

    # Create the file
    filePath = os.path.join(transPath, subDir, fileName)
    f = open(filePath, "w", encoding = "utf-8")

    # Write the lines, converted to UTF-8
    f.writelines([l + '\n' for l in lines])
    f.close()


# Convert the text from a script instruction to the format used in map
# translation files and append it to a list.
def appendScriptString(l, header, instr, version, kanjiBitmap = None):
    string = wa.text.decode(instr.getText(), version, kanjiBitmap)
    string = string.replace("{CR}", "\n")
    string = string.replace("{CLEAR}", "{CLEAR}\n")
    l.append(header)
    l.append(string)


# Extract strings and font from the main game executable.
def extractExec(image, transPath):
    print("Dumping executable...")

    # Retrieve the file
    discFile = image.openFile("EXE", "WILDARMS.EXE")
    data = bytearray(discFile.read())
    discFile.close()

    # Conversion between file offsets and memory addresses
    baseAddr = 0x80011420 - 0x800

    #
    # Extract the string tables with pointer blocks
    #

    offsetList = wa.data.execFileData(image.version)

    for tableOffset, numStrings, dataOffset, dataSize, specialBytes, specialHack, transDir, transFileName in offsetList:

        # Fetch the pointer table
        pointers = struct.unpack_from("<%dL" % numStrings, data, tableOffset)

        # Extract the strings
        lines = []

        specialCount = 2

        for p in pointers:

            # Hack for one table whose first two strings have two additional data bytes
            if specialHack and specialCount > 0:
                skipBytes = specialBytes + 2
                specialCount -= 1
            else:
                skipBytes = specialBytes

            # Extract string data until the first null byte
            o = p - baseAddr + skipBytes
            s = data[o:data.index(b'\0', o)]

            string = wa.text.decode(s, image.version)
            lines.append(string)

        # Save the translation file
        saveTrans(transPath, transDir, transFileName, lines)

    #
    # Extract the simple string tables
    #

    offsetList = wa.data.execFileData2(image.version)

    for offset, numStrings, stringSize, encoding, transDir, transFileName in offsetList:

        # Extract the strings
        lines = []

        for i in range(numStrings):
            base = offset + i*stringSize
            s = data[base:base + stringSize]

            if encoding is not None:
                string = s.rstrip(b'\0').decode(encoding)
            else:
                string = wa.text.decode(s, image.version)

            lines.append(string)

        # Save the translation file
        saveTrans(transPath, transDir, transFileName, lines)

    #
    # Extract the fonts
    #

    fontList = wa.data.fontData(image.version)

    for offset, numChars, charWidth, charHeight, lineSpacing, outCharsPerRow, transDir, transFileName in fontList:
        fontHeight = charHeight + lineSpacing

        # Create image from font bitmap, arranging the characters in table fashion
        imageWidth = outCharsPerRow * charWidth
        imageHeight = (numChars + outCharsPerRow - 1) // outCharsPerRow * fontHeight

        img = Image.new("1", (imageWidth, imageHeight))

        for c in range(numChars):
            xBase = (c % outCharsPerRow) * charWidth
            yBase = (c // outCharsPerRow) * fontHeight

            for y in range(charHeight):
                for x in range(charWidth):
                    byte = x // 8
                    bit = 7 - (x % 8)

                    if bit == 7:
                        # Fetch next data byte
                        d = data[offset]
                        offset += 1

                    img.putpixel((xBase + x, yBase + y), (d >> bit) & 1)

        # Write output image in PNG format
        outputDir = os.path.join(transPath, transDir)
        if not os.path.isdir(outputDir):
            os.mkdir(outputDir)

        img.save(os.path.join(outputDir, transFileName), "PNG")

    #
    # Extract texts from the embedded scripts
    #

    tableOffset, numScripts, dataOffset, dataSize = wa.data.execScriptData(image.version)

    pointers = struct.unpack_from("<%dL" % numScripts, data, tableOffset)

    index = 1
    lines = []

    # Extract the scripts
    for p in pointers:
        offset = p - baseAddr

        while True:
            instr = wa.map.parseInstruction(data, offset, image.version, baseAddr)

            if instr.op == Op.MESSAGE:
                header = "\u25b6 %d (dialog)" % index
                appendScriptString(lines, header, instr, image.version)
                index += 1
            elif instr.op == Op.RETURN:
                break

            offset += instr.length

    # Save to output file
    saveTrans(transPath, "exe", "script.txt", lines)


# Extract strings from the UT0.OVR overlay.
def extractUtil(image, transPath):
    print("Dumping overlay...")

    # Retrieve the file
    discFile = image.openFile("SYS", "UT0.OVR")
    data = discFile.read()
    discFile.close()

    # Conversion between file offsets and memory addresses
    baseAddr = 0x801b0000

    # Extract the string tables
    offsetList = wa.data.utilFileData(image.version)

    for tableOffset, numStrings, dataOffset, dataSize, stringSize, transDir, transFileName in offsetList:

        # Fetch the pointer table
        pointers = struct.unpack_from("<%dL" % numStrings, data, tableOffset)

        # Extract the strings
        lines = []
        for p in pointers:

            # Extract string data until the first null byte
            o = p - baseAddr
            s = data[o:data.index(b'\0', o)]

            string = wa.text.decode(s, image.version)
            lines.append(string)

        # Save the translation file
        saveTrans(transPath, transDir, transFileName, lines)


# Extract strings from maps.
def extractMaps(image, transPath):
    print("Dumping maps...")

    # Retrieve the map file
    mapFile = image.openFile("BIN", "CDSTG.BIN")

    # Process all maps
    for mapNumber in range(128):

        # Read the map data block
        block = mapFile.read(0x91000)

        # Map 25 is dummied out
        if mapNumber == 25:
            continue

        transFileName = "%03d.txt" % mapNumber
        extraFileName = "%03d_extra.txt" % mapNumber

        # Extract the scripts
        mapData = wa.map.MapData(block, mapNumber, image.version)
        script = mapData.getScript1() + mapData.getScript2()

        # Look for message and string instructions and get their text
        index = 1
        lines = []
        for instr in script:
            if instr.op not in [Op.MESSAGE, Op.STRING]:
                continue

            if instr.op == Op.MESSAGE:
                header = "\u25b6 %d (dialog)" % index
            else:
                header = "\u25b6 %d (string)" % index

            appendScriptString(lines, header, instr, image.version, mapData.kanjiBitmap)
            index += 1

        # Save to output file
        saveTrans(transPath, "map", transFileName, lines)

        # Save extra strings in the MIPS code
        codeStrings = mapData.getCodeStrings()

        if codeStrings:
            lines = [wa.text.decode(s, image.version) for s in codeStrings]
            saveTrans(transPath, "map", extraFileName, lines)

    mapFile.close()


# Convert 16-bit little-endian ABGR pixels to RGB (PIL's "BGR;15" format).
def convertABGR(data):
    output = bytearray()

    for i in range(0, len(data), 2):
        pixel = struct.unpack_from("<H", data, i)[0]

        if pixel == 0:
            # Replacement color (magenta) for transparent pixels
            r = 0x1f
            g = 0
            b = 0x1f
        else:
            r = pixel & 0x1f
            g = (pixel >> 5) & 0x1f
            b = (pixel >> 10) & 0x1f

        pixel = (r << 10) | (g << 5) | b

        output.extend(struct.pack("<H", pixel))

    return bytes(output)


# Extract textures from menu graphics archives.
def extractTextures(image, transPath):
    print("Dumping textures...")

    transDir = "gfx"

    for subDir, fileName, archiveSize, lastSectionSize, textureList in wa.data.textureData:

        # Retrieve the archive
        data = image.openFile(subDir, fileName).read()
        file = io.BytesIO(data[:archiveSize])
        archive = wa.archive.Archive(file)

        # Extract all textures
        for pixelSection, clutSection, dimensions, clutOffset, transFileName in textureList:

            # Retrieve CLUT and pixel data
            clutData = archive.getSection(clutSection)[clutOffset:clutOffset + 32]
            pixelData = wa.lzss.decompress(archive.getSection(pixelSection))

            # Expand 4-bit pixel data to 8-bit
            output = bytearray()
            for x in pixelData:
                output.append(x & 0x0f)
                output.append(x >> 4)

            # Convert to image
            img = Image.frombytes("P", dimensions, output, "raw", "P", 0, 1)
            img.palette = ImagePalette.raw("BGR;15", convertABGR(clutData))

            # Write output image in PNG format
            outputDir = os.path.join(transPath, transDir)
            if not os.path.isdir(outputDir):
                os.mkdir(outputDir)

            img.save(os.path.join(outputDir, transFileName), "PNG")


# Print usage information and exit.
def usage(exitcode, error = None):
    print("Usage: %s [OPTION...] <game_dir_or_image> <trans_dir>" % os.path.basename(sys.argv[0]))
    print("  -a, --altchars                  Use alternate character set for text")
    print("  -V, --version                   Display version information and exit")
    print("  -?, --help                      Show this help message")

    if error is not None:
        print("\nError:", error, file=sys.stderr)

    sys.exit(exitcode)


# Parse command line arguments
gamePath = None
transPath = None
altCharset = False

for arg in sys.argv[1:]:
    if arg == "--version" or arg == "-V":
        print("UnTrans", __version__)
        sys.exit(0)
    elif arg == "--help" or arg == "-?":
        usage(0)
    elif arg == "--altchars" or arg == "-a":
        altCharset = True
    elif arg[0] == "-":
        usage(64, "Invalid option '%s'" % arg)
    else:
        if gamePath is None:
            gamePath = arg
        elif transPath is None:
            transPath = arg
        else:
            usage(64, "Unexpected extra argument '%s'" % arg)

if gamePath is None:
    usage(64, "No disc image or game data input directory specified")
if transPath is None:
    usage(64, "No translation output directory specified")

if altCharset:
    wa.text.setAltCharset()

try:

    # Open the input image
    image = wa.openImage(gamePath)

    # Create the output directory
    if os.path.isfile(transPath):
        raise EnvironmentError("Cannot create translation directory '%s': Path refers to a file" % transPath)

    if os.path.isdir(transPath):
        answer = None
        while answer not in ["y", "n"]:
            answer = input("Output directory '%s' exists. Delete and overwrite it (y/n)? " % transPath)

        if answer == 'y':
            shutil.rmtree(transPath)
        else:
            sys.exit(0)

    print("Creating translation directory '%s'..." % transPath)

    try:
        os.makedirs(transPath)
    except OSError as e:
        print("Cannot create translation directory '%s': %s" % (transPath, e.strerror), file=sys.stderr)
        sys.exit(1)

    # Extract everything
    extractExec(image, transPath)
    extractUtil(image, transPath)
    extractMaps(image, transPath)
    extractTextures(image, transPath)

    print("Done.")

except Exception as e:

    # Pokemon exception handler
    print(e, file=sys.stderr)
    sys.exit(1)