# -*- coding: utf-8 -*-

# UTF8 to cp1251 and ShiftJIS recoder
# by Chtobi and Nazon, 2016

import codecs
import argparse
from os import path

JAPANESE_CODEPAGE = 'shift_jis'

UTF_CODEPAGE = 'utf-8'
RUS_CODEPAGE = 'cp1251'


def nonrus_handler(e):

    if e.object[e.start:e.end] == '~':  # UTF-8: 0xEFBD9E?->?SHIFT-JIS: 0x8160
        japstr_byte = b'\x81\x60'

    elif e.object[e.start:e.end] == '-':  # UTF-8: 0xEFBC8D?->?SHIFT-JIS: 0x817C
        japstr_byte = b'\x81\x7c'

    else:
        japstr_byte = (e.object[e.start:e.end]).encode(JAPANESE_CODEPAGE)

    return japstr_byte, e.end


if __name__ == '__main__':

    arg_parser = argparse.ArgumentParser(prog="Recode to cp1251 and ShiftJIS",
                                         description="Program to encode UTF8 text file to "
                                         "cp1251 for all cyrillic symbols and ShiftJIS for others. "
                                         "Output file will be inputfilename.s",
                                         usage="recode_to_cp1251_shiftjis.py file_name")
    arg_parser.add_argument('file_name', nargs=1, type=argparse.FileType(mode='r', bufsize=-1),
                            help="Input text file name. Only files coded in UTF8 are allowed.\n")

    codecs.register_error('nonrus_handler', nonrus_handler)

    input_name = arg_parser.parse_args().file_name[0].name

    output_name = path.splitext(input_name)[0] + ".s"

    with open(input_name, 'rt', encoding=UTF_CODEPAGE) as input_file:
        with open(output_name, 'wb') as output_file:

            for line in input_file:
                for char1 in line:
                    bytes_out = bytes(line, UTF_CODEPAGE)

                    output_file.write(char1.encode(RUS_CODEPAGE, "nonrus_handler"))

        print("Done.")

# -*- coding: utf-8 -*-

# Qlie engine dpng files merger
# by Chtobi and Nazon, 2016
# Requires ImageMagick magick.exe on the path.

import os
import glob
import re
import argparse
import subprocess

IMGMAGIC = os.path.dirname(os.path.abspath(__file__)) + '\\' + 'magick.exe'

IMGMAGIC_PARAMS1 = ['-background', 'rgba(0,0,0,0)']
IMGMAGIC_PARAMS2 = ['-mosaic']

INPUT_FILES_MASK = '*+DPNG[0-9][0-9][0-9]+*.png'

SPLIT_MASK = '+DPNG'

x_y_ajusts_re = re.compile('(.+)\+DPNG[0-9][0-9][0-9]\+x(\d+)y(\d+)\.')


if __name__ == '__main__':

    arg_parser = argparse.ArgumentParser(prog="DPNG Merger\n"
                                         "Program to merge sliced png files from QLIE engine. "
                                         "All files with mask *+DPNG[0-9][0-9][0-9]+*.png"
                                         "into the input directory will be merged and copied to the"
                                         "output directory.\n",
                                         usage="connect_png.py input_dir [output_dir]\n")
    arg_parser.add_argument("input_dir_param", nargs=1, help="Full path to the input directory.\n")
    arg_parser.add_argument("output_dir_param", nargs='?', default=os.path.dirname(os.path.abspath(__file__)),
                            help="Full path to the output directory. "
                                 "It would be a script parent directory if not specified.\n")

    input_dir = arg_parser.parse_args().input_dir_param[0]
    output_dir = arg_parser.parse_args().output_dir_param[0]

    os.chdir(input_dir)

    all_append_files = glob.glob(INPUT_FILES_MASK)  # Select only files with DPNG

    prep_bunches = []

    for file_in_dir in all_append_files:
        # Check all files and put all splices that should be connected in separate list
        for num, bunch in enumerate(prep_bunches):
            name_first_part = bunch[0].partition(SPLIT_MASK)[0]  # Part of the filename before +DPNG should be unique

            if name_first_part == file_in_dir.partition(SPLIT_MASK)[0]:
                prep_bunches[num].append(file_in_dir)
                break
        else:
            prep_bunches.append([file_in_dir])

    os.chdir(os.path.dirname(os.path.abspath(__file__)))  # Go to the script parent dir

    for prepared_bunch in prep_bunches:

        sorted_bunch = sorted(prepared_bunch)

        # Prepare -page params for imgmagic
        png_pages_params = [["(", "-page", "+{0}+{1}".format(*[(x_y_ajusts_re.match(part_file).group(2)),
                            x_y_ajusts_re.match(part_file).group(3)]), input_dir+part_file, ")"]
                            for part_file in sorted_bunch]

        connect_png_list =             [imgmagick_page for imgmagick_pages in png_pages_params for imgmagick_page in imgmagick_pages]

        output_file = output_dir + sorted_bunch[0].partition(SPLIT_MASK)[0] + ".png"

        subprocess.check_output([IMGMAGIC] + IMGMAGIC_PARAMS1 + connect_png_list + IMGMAGIC_PARAMS2 + [output_file])

# -*- coding: utf-8 -*-

# Extract b
# AnimatedBMP extractor for Bishoujo Mangekyou game files
# by Chtobi and Nazon, 2016

import glob
import os
import struct
import argparse
from collections import namedtuple

b_hdr = b'abmp12'+bytes(10)

signa_len = 16

b_abdata = (b'abdata10'+bytes(8), b'abdata11'+bytes(8), b'abdata12'+bytes(8), b'abdata13'+bytes(8))

b_imgdat = (b'abimgdat10'+bytes(6), b'abimgdat11'+bytes(6), b'abimgdat14'+bytes(6))

b_img = (b'abimage10'+bytes(7), b'abimage11'+bytes(7), b'abimage12'+bytes(7), b'abimage13'+bytes(7),
         b'abimage14'+bytes(7))

b_sound = (b'absound10'+bytes(7), b'absound11'+bytes(7), b'absound12'+bytes(7))
# not sure about structure of sound11 and sound12

b_snd = (b'absnddat11'+bytes(7), b'absnddat10'+bytes(7), b'absnddat12'+bytes(7))


Abimgdat13_pattern = namedtuple('Abimgdat13', ['signa', 'name_size_len', 'hash_size_len', 'unknown1_len',
                                               'unknown2_len', 'data_size_len'])
Abimgdat13 = Abimgdat13_pattern(signa=b'abimgdat13'+bytes(6), name_size_len=2, hash_size_len=2, unknown1_len=1,
                                unknown2_len=12, data_size_len=4)

Abimgdat14_pattern = namedtuple('Abimgdat14', ['signa', 'name_size_len', 'hash_size_len', 'unknown1_len',
                                               'data_size_len'])
Abimgdat14 = Abimgdat14_pattern(signa=b'abimgdat14'+bytes(6), name_size_len=2, hash_size_len=2, unknown1_len=77,
                                data_size_len=4)

Abimgdat_pattern = namedtuple('Abimgdat', ['name_size_len', 'hash_size_len', 'unknown1_len', 'data_size_len'])
# probably, abimgdat10,abimgdat11 and others
Other_imgdat = Abimgdat_pattern(name_size_len=2, hash_size_len=2, unknown1_len=1, data_size_len=4)

Absnddat11_pattern = namedtuple('Absnddat11', ['signa', 'name_size_len', 'hash_size_len', 'unknown1_len',
                                               'data_size_len'])
Absnddat11 = Absnddat11_pattern(signa=b'absnddat11'+bytes(7), name_size_len=2, hash_size_len=2, unknown1_len=1,
                                data_size_len=4)


def create_parser():
    arg_parser = argparse.ArgumentParser(prog='AnimatedBMP extractor\n',
                                         usage='extract_b input_file_name output_dir\n',
                                         description='AnimatedBMP extractor for QLIE engine *.b files.\n')
    arg_parser.add_argument('input_file_name', nargs='+', help="Input file with full path(wildcards are supported).\n")
    arg_parser.add_argument('output_dir', nargs=1,
                            help="Output directory.\n")

    return arg_parser


def check_type(file_buf):
    if file_buf.startswith(b'\x89' + b'PNG'):
        return '.png'
    elif file_buf.startswith(b'BM'):
        return '.bmp'
    elif file_buf.startswith(b'JFIF', 6):
        return '.jpg'
    elif file_buf.startswith(b'IMOAVI'):
        return '.imoavi'
    elif file_buf.startswith(b'OggS'):
        return '.ogg'
    elif file_buf.startswith(b'RIFF'):
        return '.wav'
    else:
        return ''


def bytes_shiftjis_to_utf8(shiftjis_bytes):
    shiftjis_str = shiftjis_bytes.decode('shift_jis', 'strict')
    utf_str = shiftjis_str.encode('utf-8', 'strict').decode('utf-8', 'strict')

    return utf_str


def check_signa(f_buffer):

    if f_buffer.endswith(b_abdata):
        return 'abdata'
    elif f_buffer.endswith(b_img):
        return 'abimgdat'
    elif f_buffer.endswith(b_sound):
        return 'absound'


def prepare_filename(out_file_name, out_dir, postfix=''):
    ready_name = out_dir + os.path.basename(out_file_name) + postfix

    return ready_name


def create_file(file_name_hndl, out_buffer):
    if len(out_buffer) != 0:
        with open(file_name_hndl, 'wb') as ext_file:
            ext_file.write(out_buffer)
    else:
        print("Zero file. Skipped.")


def check_file_header(file_handle, bytes_num):
    file_handle.seek(0)

    readed_bytes = file_handle.read(bytes_num)

    if readed_bytes == b_hdr:
        print("File is valid abmp")
        return True
    else:
        print("Can't read header. Probably, wrong file...")
        return False


if __name__ == '__main__':

    parser = create_parser()
    arguments = parser.parse_args()

    all_b_files = glob.glob(arguments.input_file_name[0])

    output_dir = arguments.output_dir[0]

    for b_file in all_b_files:

        file_buffer = bytearray(b'')

        with open(b_file, 'rb') as bfile_h:

            check_file_header(bfile_h, len(b_hdr))

            read_byte = bfile_h.read(1)

            file_buffer.extend(read_byte)

            while read_byte:

                read_byte = bfile_h.read(1)

                file_buffer.extend(read_byte)

                # Finding content sections signature
                check_result = check_signa(file_buffer)

                if check_result:
                    if check_result == 'abdata':
                        file_buffer = bytearray(b'')

                        read_length = bfile_h.read(4)

                        size = struct.unpack('<L', read_length)[0]

                        file_buffer.extend(bfile_h.read(size))

                        # Adding _abdata to separate from other parts
                        outfile_name = prepare_filename(b_file, output_dir, '_abdata')

                        create_file(outfile_name, file_buffer)

                    elif check_result == 'abimgdat':

                        images_number = struct.unpack('B', bfile_h.read(1))[0]  # Number of pictures in section

                        for i1 in range(images_number):
                            file_buffer = bytearray(b'')
                            file_name = ''

                            imgsec_hdr = bfile_h.read(signa_len)

                            if imgsec_hdr == Abimgdat13.signa:

                                file_name_size = struct.unpack('<H', bfile_h.read(Abimgdat13.name_size_len))[0]

                                # Decode filename to utf8
                                file_name = bytes_shiftjis_to_utf8(bfile_h.read(file_name_size))

                                # CRC size
                                hash_size = struct.unpack('<H', bfile_h.read(Abimgdat13.hash_size_len))[0]
                                # Picture CRC (don't need it)
                                pic_hash = bfile_h.read(hash_size)
                                unknown1 = bfile_h.read(Abimgdat13.unknown1_len)
                                unknown2 = bfile_h.read(Abimgdat13.unknown2_len)

                                pic_size = struct.unpack('<L', bfile_h.read(Abimgdat13.data_size_len))[0]
                                print("pic_size:", pic_size)

                                file_buffer.extend(bfile_h.read(pic_size))
                            elif imgsec_hdr == Abimgdat14.signa:

                                file_name_size = struct.unpack('<H', bfile_h.read(Abimgdat14.name_size_len))[0]
                                file_name = bytes_shiftjis_to_utf8(bfile_h.read(file_name_size))

                                hash_size = struct.unpack('<H', bfile_h.read(Abimgdat14.hash_size_len))[0]
                                pic_hash = bfile_h.read(hash_size)

                                bfile_h.seek(Abimgdat14.unknown1_len, os.SEEK_CUR)

                                pic_size = struct.unpack('<L', bfile_h.read(Abimgdat14.data_size_len))[0]

                                file_buffer.extend(bfile_h.read(pic_size))

                            else:  # probably abimgdat10, abimgdat11...

                                file_name_size = struct.unpack('<H', bfile_h.read(Other_imgdat.name_size_len))[0]

                                file_name = bytes_shiftjis_to_utf8(bfile_h.read(file_name_size))

                                hash_size = struct.unpack('<H', bfile_h.read(Other_imgdat.hash_size_len))[0]
                                pic_hash = bfile_h.read(hash_size)

                                bfile_h.seek(Other_imgdat.unknown1_len, os.SEEK_CUR)

                                pic_size = struct.unpack('<L', bfile_h.read(Other_imgdat.data_size_len))[0]
                                file_buffer.extend(bfile_h.read(pic_size))

                            for i, letter in enumerate(file_name):  # Replace any unusable symbols from filename with _
                                if letter == '<' or letter == '>' or letter == '*' or letter == '/':
                                    file_name = file_name.replace(letter, "_")

                            # Checking file signature and adding proper extension
                            outfile_name = prepare_filename(b_file, output_dir, '_' + file_name +
                                                            check_type(file_buffer))

                            create_file(outfile_name, file_buffer)
                            file_buffer = bytearray(b'')

                    elif check_result == 'absound':

                        sound_files_number = struct.unpack('B', bfile_h.read(1))[0]

                        for i2 in range(sound_files_number):
                            file_buffer = bytearray(b'')
                            file_name = ''

                            sndsec_hdr = bfile_h.read(signa_len)

                            if sndsec_hdr == Absnddat11.signa:
                                file_name_size = struct.unpack('<H', bfile_h.read(Absnddat11.name_size_len))[0]
                                file_name = bytes_shiftjis_to_utf8(bfile_h.read(file_name_size))

                                hash_size = struct.unpack('<H', bfile_h.read(Absnddat11.hash_size_len))[0]
                                snd_hash = bfile_h.read(hash_size)

                                unknown1 = bfile_h.read(Absnddat11.unknown1_len)
                                snd_size = struct.unpack('<L', bfile_h.read(Absnddat11.data_size_len))[0]

                                file_buffer.extend(bfile_h.read(snd_size))
                            else:
                                file_name_size = struct.unpack('<H', bfile_h.read(Absnddat11.name_size_len))[0]
                                file_name = bytes_shiftjis_to_utf8(bfile_h.read(file_name_size))

                                hash_size = struct.unpack('<H', bfile_h.read(Absnddat11.hash_size_len))[0]
                                snd_hash = bfile_h.read(hash_size)

                                unknown1 = bfile_h.read(Absnddat11.unknown1_len)
                                snd_size = struct.unpack('<L', bfile_h.read(Absnddat11.data_size_len))[0]

                                file_buffer.extend(bfile_h.read(snd_size))

                            for i, letter in enumerate(file_name):
                                if letter == '<' or letter == '>' or letter == '*' or letter == '/':
                                    file_name[i] = '_'

                            outfile_name = prepare_filename(b_file, output_dir, '_' + file_name +
                                                            check_type(file_buffer))

                            print("create absound")
                            create_file(outfile_name, file_buffer)
                            file_buffer = bytearray(b'')

# -*- coding: utf-8 -*-

# Extract imoavi
# Imoavi extractor for Bishoujo Mangekyou game files
# by Chtobi and Nazon, 2016


import glob
import os
import struct
import argparse

imoavi_hdr = b'IMOAVI'

hdr_len = len(imoavi_hdr)


def create_file(file_name, out_buffer, wr_mode='wb'):
    if len(out_buffer) != 0:

        with open(file_name, wr_mode) as ext_file:
            ext_file.write(out_buffer)
    else:
        print("Zero file. Skipped.")


def prepare_filename(file_name, out_dir, postfix=''):
    ready_name = out_dir + os.path.basename(file_name) + postfix

    return ready_name


def create_parser():
    arg_parser = argparse.ArgumentParser(prog='Imoavi extractor\n',
                                         usage='extract_imoavi input_file_name output_dir\n',
                                         description='Imoavi extractor for QLIE engine *.imoavi files.\n')
    arg_parser.add_argument('input_file_name', nargs='+', help="Input file with full path(wildcards are supported).\n")
    arg_parser.add_argument('output_dir', nargs='+', help="Output directory.\n")

    return arg_parser


if __name__ == '__main__':

    parser = create_parser()
    arguments = parser.parse_args()

    all_imoavi = glob.glob(arguments.input_file_name[0])

    output_dir = arguments.output_dir[0]

    for imoavi_f in all_imoavi:
        file_buffer = bytearray(b'')

        with open(imoavi_f, 'rb') as imoavi_h:

            # Read imoavi file header
            imoavi_h.read(hdr_len)
            imoavi_h.seek(2, os.SEEK_CUR)  # 0x00

            imoavi_h.seek(1, os.SEEK_CUR)  # 0x64
            imoavi_h.seek(3, os.SEEK_CUR)  # 0x00

            imoavi_h.seek(5, os.SEEK_CUR)  # SOUND
            imoavi_h.seek(3, os.SEEK_CUR)  # 0x00
            imoavi_h.seek(1, os.SEEK_CUR)  # 0x64
            imoavi_h.seek(11, os.SEEK_CUR)
            imoavi_h.seek(5, os.SEEK_CUR)  # Movie
            imoavi_h.seek(3, os.SEEK_CUR)  # 00 ??
            imoavi_h.seek(1, os.SEEK_CUR)  # 0x64
            imoavi_h.seek(3, os.SEEK_CUR)  # 0x00 ??
            imoavi_h.seek(4, os.SEEK_CUR)  # ??
            imoavi_h.seek(1, os.SEEK_CUR)  # Number of jpg files in section
            imoavi_h.seek(4, os.SEEK_CUR)  # 0x00
            imoavi_h.seek(1, os.SEEK_CUR)  # 0x05 ???
            imoavi_h.seek(2, os.SEEK_CUR)  # 0x00 ??
            imoavi_h.seek(4, os.SEEK_CUR)  # 720 ??
            imoavi_h.seek(4, os.SEEK_CUR)  # Full size without header?

            to_next_size = struct.unpack('<L', imoavi_h.read(4))[0]  # Bytes till next header

            imoavi_h.seek(16, os.SEEK_CUR)  # 0x00

            jpg_size = struct.unpack('<L', imoavi_h.read(4))[0]

            imoavi_h.seek(4, os.SEEK_CUR)  # 0x00

            file_num = 0

            file_buffer.extend(imoavi_h.read(jpg_size))

            outfile_name = prepare_filename(imoavi_f, output_dir, '_' + (str(file_num)).zfill(3) + '.jpg')

            create_file(outfile_name, file_buffer)

            while to_next_size != 0:
                file_buffer = bytearray(b'')

                to_next_size = struct.unpack('<L', imoavi_h.read(4))[0]

                if to_next_size == 24:  # 0x1C header for index part

                    file_buffer.extend(imoavi_h.read(to_next_size))

                    outfile_name = prepare_filename(imoavi_f, output_dir, '_' + '.index')
                    create_file(outfile_name, file_buffer, 'ab')  # concatenate with index file
                else:

                    imoavi_h.seek(2, os.SEEK_CUR)  # unknown
                    imoavi_h.seek(2, os.SEEK_CUR)  # Unknown, almost always FF FF or FF FE

                    file_num = struct.unpack('B', imoavi_h.read(1))[0]  # File number

                    imoavi_h.seek(11, os.SEEK_CUR)  # 0x00

                    jpg_size = struct.unpack('<L', imoavi_h.read(4))[0]

                    imoavi_h.seek(4, os.SEEK_CUR)  # 0x00

                    file_buffer.extend(imoavi_h.read(jpg_size))

                    outfile_name = prepare_filename(imoavi_f, output_dir, '_' + (str(file_num)).zfill(3) + '.jpg')

                    create_file(outfile_name, file_buffer)