import struct
import sys
import json


class BPScript:
    def __init__(self, path: str):
        self.path = path
        with open(path, 'rb') as f:
            self.data = bytearray(f.read())
        self.len = len(self.data)
        self.read_header()
        self.iPos = self.header_size

    def read_header(self):
        self.header_size = struct.unpack('<L', self.data[0:4])[0]
        self.instr_size = struct.unpack('<L', self.data[4:8])[0]
        if self.len != self.header_size + self.instr_size:
            raise ValueError("Invalid BPScript file size")
        self.iPos = self.header_size
        self.last_pos = 0
        while self.iPos < self.len - 4:
            d = struct.unpack('<L', self.data[self.iPos:self.iPos + 4])[0]
            if d == 0x17:
                self.last_pos = self.iPos
                self.iPos += 4
            else:
                self.iPos += 1

    def extract_string(self):
        offset = struct.unpack('<h', self.data[self.iPos:self.iPos + 2])[0]
        start_pos = self.iPos + offset - 1
        self.iPos += 2
        if start_pos < self.last_pos or start_pos >= self.len or ((start_pos > self.last_pos + 1) and self.data[start_pos-1] != 0):
            self.iPos -= 2
            return None
        pos = start_pos
        while True:
            if self.data[pos] == 0x00:
                break
            pos += 1
        data = self.data[start_pos:pos]
        data_len = len(data)
        try:
            data = data.decode('cp932')
        except UnicodeDecodeError:
            data = data.decode('utf-8')
            data = "utf8:" + data
        if not data:
            return None
        return start_pos, self.iPos - 2, data_len

    def extract_values(self):
        str_pos = {}
        self.iPos = self.header_size
        while self.iPos < self.last_pos:
            ins = self.data[self.iPos]
            self.iPos += 1
            if ins == 0x5:
                try:
                    t = self.extract_string()
                    if t is not None:
                        if t[0] not in str_pos:
                            str_pos[t[0]] = (t[2], [t[1]])
                        else:
                            str_pos[t[0]][1].append(t[1])
                except UnicodeDecodeError:
                    self.iPos -= 2
        return str_pos

    def setup_values(self, values: dict, append: bool = True):
        data = bytearray(self.data)
        goffsets = self.extract_values()
        offsets = {}
        for k, v in goffsets.items():
            offsets[k] = v[0]
        print(offsets)
        self.iPos = 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]
        instr_size = self.instr_size
        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
            if append:
                new_offsets[key] = instr_size + self.header_size
                instr_size += new_len
            else:
                for j in range(i + 1, len(all_keys)):
                    off = all_keys[j]
                    new_offsets[off] += new_len - ori_len
        print(new_offsets)
        old_len = last_key + last_len
        if append:
            new_len = self.header_size + instr_size
        else:
            new_len = new_offsets[last_key] + last_len
        # 对齐到16字节
        # new_len = new_len + (16 - new_len % 16)
        print(old_len, new_len, self.len)
        if old_len == self.len:
            data = bytearray(data + b'\0' * (new_len - self.len))
        else:
            if append:
                data = bytearray(data + b'\0' * (new_len - self.len))
            else:
                data = bytearray(data[0:all_keys[0]] + b'\0' * (new_len - self.len) + data[all_keys[0]:])
        self.len = len(data)
        data[4:8] = struct.pack('<L', self.len - self.header_size)
        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])
        for k, v in goffsets.items():
            for i in v[1]:
                data[i:i+2] = struct.pack('<H', new_offsets[k] - i + 1)
        self.data = data
        print(len(data))
    
    def save_as(self, path: str):
        with open(path, 'wb') as f:
            f.write(self.data)


append = True
if '--no-append' in sys.argv:
    append = False
    sys.argv.remove('--no-append')
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 = BPScript(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, append)
scr.save_as(output)