import struct
import json
import sys


class BGIScript:
    def __init__(self, path: str):
        self.path = path
        with open(path, 'rb') as f:
            self.data = bytearray(f.read())
        self.read_header()
        self.get_instr_end()

    def read_header(self):
        self.header_size = struct.unpack('<L', self.data[28:32])[0]
        self.iPos = 28 + self.header_size

    def get_instr_end(self):
        finder = self.iPos
        self.len = len(self.data)
        self.last_pos = self.len
        while finder + 4 < self.len:
            d = struct.unpack('<L', self.data[finder:finder + 4])[0]
            if d == 0x1b:
                self.last_pos = finder
            finder += 4
        print(self.len)

    def extract_string(self):
        offset = struct.unpack('<L', self.data[self.iPos:self.iPos + 4])[0]
        self.iPos += 4
        start_pos = 28 + self.header_size + offset
        if start_pos < self.last_pos:
            return None
        pos = start_pos
        while True:
            if self.data[pos] == 0x00:
                break
            pos += 1
        data = self.data[start_pos:pos]
        try:
            data = data.decode('cp932')
        except UnicodeDecodeError:
            data = data.decode('utf-8')
            data = "utf8:" + data
        return start_pos, pos - start_pos
    
    def extract_value(self):
        offset = struct.unpack('<L', self.data[self.iPos:self.iPos + 4])[0]
        self.iPos += 4
        return offset

    def extract_values(self):
        strings = []
        while self.iPos < self.last_pos:
            ins = struct.unpack('<L', self.data[self.iPos:self.iPos + 4])[0]
            self.iPos += 4
            if ins == 0x03:
                string = self.extract_string()
                if string is not None:
                    strings.append(string)
            elif ins in [0x00, 0x01, 0x02, 0x08, 0x09, 0x0a, 0x19, 0x3f]:
                val = self.extract_value()
        return dict(strings)

    def setup_values(self, values: dict):
        data = bytearray(self.data)
        offsets = self.extract_values()
        self.iPos = 28 + self.header_size
        all_keys = list(int(i) for i in values.keys())
        all_keys.sort()
        new_offsets = {}
        for key in all_keys:
            new_offsets[key] = key
        last_len = 0
        last_key = all_keys[-1]
        for i in range(len(all_keys)):
            key = all_keys[i]
            ori_len = offsets[key] + 1
            v = values[f"{key}"]
            if v.startswith("utf8:"):
                v = v[5:]
                v = v.encode("utf-8")
            else:
                v = v.encode("cp932")
            v += b'\0'
            new_len = len(v)
            last_len = new_len
            if new_len == ori_len:
                continue
            for j in range(i + 1, len(all_keys)):
                off = all_keys[j]
                new_offsets[off] += new_len - ori_len
        old_len = last_key + last_len
        new_len = new_offsets[last_key] + last_len
        print(new_len, old_len)
        data = bytearray(data[:all_keys[0]] + b'\0' * (new_len - old_len) + data[all_keys[0] + last_len:])
        print(len(data))
        for i in all_keys:
            start = new_offsets[i]
            v = values[f"{i}"]
            if v.startswith("utf8:"):
                v = v[5:]
                v = v.encode("utf-8")
            else:
                v = v.encode("cp932")
            v += b'\0'
            end = start + len(v)
            data[start:end] = v
            print(data[start:end])
        while self.iPos < self.last_pos:
            ins = struct.unpack('<L', self.data[self.iPos:self.iPos + 4])[0]
            self.iPos += 4
            if ins == 0x03:
                off = struct.unpack('<L', self.data[self.iPos:self.iPos + 4])[0] + 28 + self.header_size
                new_off = new_offsets[int(off)]
                data[self.iPos:self.iPos + 4] = struct.pack('<L', new_off - 28 - self.header_size)
                self.iPos += 4
            elif ins in [0x00, 0x01, 0x02, 0x08, 0x09, 0x0a, 0x19, 0x3f]:
                val = self.extract_value()
        print(len(data))
        print(new_offsets)
        self.data = data

    def save_as(self, path: str):
        with open(path, 'wb') as f:
            f.write(self.data)


base = sys.argv[1]
json_f = f"{base}.json"
if len(sys.argv) >= 3:
    json_f = sys.argv[2]
output = f"../patched/{base}"
if len(sys.argv) >= 4:
    output = sys.argv[3]
scr = BGIScript(base)
print(scr.header_size)
print(scr.iPos)
print(scr.last_pos)
with open(json_f, "r", encoding="utf-8") as f:
    data = json.load(f)
scr.setup_values(data)
scr.save_as(output)