#!/usr/bin/python3
# -*- coding: utf-8 -*-
# 2018/02/19
# Reference: http://gemrb.org/iesdp/file_formats/ie_formats/tlk_v1.htm
import sys, os
import openpyxl
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment

SOURCE_FILE = "~/GOG Games/Planescape Torment Enhanced Edition/game/lang/en_US/dialog.tlk"
TARGET_FILE = "~/GOG Games/Planescape Torment Enhanced Edition/game/lang/zh_TW/dialog.tlk"
INPUT_SHEET = "input.xlsx"
OUTPUT_SHEET = "output.xlsx"
DEFAULT_ENCODING = "utf-8"

def input_xlsx():
    print("Loading the spreadsheet file....")
    strings = []
    wb = openpyxl.load_workbook(INPUT_SHEET)
    ws = wb.active

    for x in range(2, ws.max_row + 1):
        strings.append(ws["C{}".format(x)].value)

    return strings

def output_xlsx(source_strings, target_strings):
    print("Generating the spreadsheet file....")
    wb = Workbook()
    ws = wb.active
    ws.column_dimensions["B"].width = 60
    ws.column_dimensions["C"].width = 60
    ws.column_dimensions["D"].width = 60

    ws["A1"] = "Ref"
    ws["B1"] = "Source"
    ws["C1"] = "Target"
    ws["D1"] = "Comment"
    ws["A1"].font = Font(bold=True)
    ws["B1"].font = Font(bold=True)
    ws["C1"].font = Font(bold=True)
    ws["D1"].font = Font(bold=True)

    row = 1

    for x in range(0, len(source_strings)):
        row = row + 1
        ws["A{}".format(row)] = x + 1
        #ws["B{}".format(row)].alignment = Alignment(wrapText=True)
        ws["B{}".format(row)] = source_strings[x]
        #ws["C{}".format(row)].alignment = Alignment(wrapText=True)
        if x < len(target_strings):
            ws["C{}".format(row)] = target_strings[x]

    ws.freeze_panes = ws["A2"]
    wb.save(OUTPUT_SHEET)
    print("Completed.")

def strings_from_tlk(path, encoding):
    refs = []
    strings = []
    full_path = os.path.expanduser(path)
    print("Loading \"{}\"".format(full_path))

    input_file = open(full_path, "rb")
    input_file.seek(0)
    signature = input_file.read(4).decode(DEFAULT_ENCODING).strip()
    #print(signature)
    version = input_file.read(4).decode(DEFAULT_ENCODING).strip()
    #print(version)
    
    lang_id = input_file.read(2)
    #print(int.from_bytes(lang_id, "little"))
    refs_count = int.from_bytes(input_file.read(4), "little")
    string_offset = int.from_bytes(input_file.read(4), "little")
    #print(string_offset)
    input_file.seek(18)
    print("References: {}".format(refs_count))
    
    for x in range(0, refs_count):
        bit_field = input_file.read(2)
        resource = input_file.read(8)
        volume = input_file.read(4)
        pitch = input_file.read(4)
        offset = int.from_bytes(input_file.read(4), "little")
        length = int.from_bytes(input_file.read(4), "little")
        refs.append([offset, length])
        
    for x in range(0, refs_count):
        input_file.seek(string_offset + refs[x][0])
        strings.append(input_file.read(refs[x][1]).decode(encoding))

    input_file.close()
    return strings

def strings_to_tlk(target_strings):
    refs = []
    source_strings = []
    input_path = os.path.expanduser(SOURCE_FILE)
    output_path = os.path.expanduser(TARGET_FILE)

    print("Loading \"{}\"".format(input_path))
    print("Creating \"{}\"".format(output_path))

    input_file = open(input_path, "rb")
    output_file = open(output_path, "w+b")
    input_file.seek(0)
    output_file.seek(0)
    output_file.write(input_file.read(18))
    output_file.seek(10)
    refs_count = int.from_bytes(output_file.read(4), "little")
    print("References: {}".format(refs_count))
    
    input_file.seek(18)
    output_file.seek(18)

    for x in range(0, refs_count):
        output_file.write(input_file.read(26))

    string_offset = output_file.tell()
    output_file.seek(14)
    output_file.write(string_offset.to_bytes(4, byteorder="little", signed=True))
    output_file.seek(18)

    for x in range(0, refs_count):
        output_file.read(18)
        offset = int.from_bytes(output_file.read(4), "little")
        length = int.from_bytes(output_file.read(4), "little")
        refs.append([offset, length])

    for x in range(0, refs_count):
        input_file.seek(string_offset + refs[x][0])
        source_strings.append(input_file.read(refs[x][1]).decode(DEFAULT_ENCODING))

    ref_pos = 36
    output_file.seek(string_offset)

    for x in range(0, refs_count):
        if x < len(target_strings) and target_strings[x] != None:
            string = bytearray(target_strings[x], DEFAULT_ENCODING) 
        else:
            string = bytearray(source_strings[x], DEFAULT_ENCODING)

        string_pos = output_file.tell()
        offset = string_pos - string_offset
        length = len(string)
        output_file.seek(ref_pos)
        output_file.write(offset.to_bytes(4, byteorder="little", signed=True))
        output_file.write(length.to_bytes(4, byteorder="little", signed=True))
        output_file.seek(string_pos)
        output_file.write(string)
        ref_pos = ref_pos + 26

    input_file.close()
    output_file.close()
    
def tlk_bilingual_export():
    source_strings = strings_from_tlk(SOURCE_FILE, DEFAULT_ENCODING)
    target_strings = strings_from_tlk(TARGET_FILE, DEFAULT_ENCODING)
    #target_strings = strings_from_tlk("Interwise/dialog.tlk", "cp950")

    if len(source_strings) == 0:
        return False

    output_xlsx(source_strings, target_strings)
    
def tlk_import():
    target_strings = input_xlsx()
    strings_to_tlk(target_strings)

if __name__ == "__main__":
    if len(sys.argv) >= 2:
        if sys.argv[1] == "--export":
            tlk_bilingual_export()
        elif sys.argv[1] == "--import":
            tlk_import()