import sys
import struct
from argparse import ArgumentParser
from os.path import splitext
from json import dump


TJS_NS0_MAGIC = b'TJS/ns0\x00TJS\x00\x00\x00\x00\x00'
sys.setrecursionlimit(200000)


class TJS_NS0:
    def __init__(self, data):
        self.data = data
        self.obj = None
        self.pos = 0

    def check_magic(self):
        if self.data[:16] != TJS_NS0_MAGIC:
            raise ValueError("Invalid TJS/ns0 magic header")
        self.pos += 16

    def parse(self):
        self.check_magic()
        self.obj = self.parse_obj()

    def read_bytes(self, length):
        if self.pos + length > len(self.data):
            raise ValueError("Unexpected end of data")
        result = self.data[self.pos:self.pos + length]
        self.pos += length
        return result
    
    def read_u16(self):
        return struct.unpack('<H', self.read_bytes(2))[0]

    def read_u32(self):
        return struct.unpack('<I', self.read_bytes(4))[0]
    
    def read_i64(self):
        return struct.unpack('<q', self.read_bytes(8))[0]

    def parse_string(self):
        str_len = self.read_u32()
        #print(f"String length: {str_len}, pos={self.pos - 4:08X}")
        return self.read_bytes(str_len * 2).decode('utf-16-le')
    
    def parse_obj(self):
        value_type = self.read_u16()
        first_byte = value_type & 0xFF
        if first_byte == 0x02:
            return self.parse_string()
        elif first_byte == 0x04:
            return self.read_i64()
        elif first_byte == 0x00:
            return value_type >> 8
        elif first_byte == 0xC1:
            return self.parse_dict()
        elif first_byte == 0x81:
            return self.parse_array()
        else:
            raise ValueError(f"Unknown object type: {value_type:04X}, pos={self.pos - 2:08X}")

    def parse_dict(self):
        kv_count = self.read_u32()
        obj = {}
        for _ in range(kv_count):
            key = self.parse_string()
            value = self.parse_obj()
            # print(f"Key: {key}, Value: {value}")
            obj[key] = value
        return obj

    def parse_array(self):
        arr_len = self.read_u32()
        #print(f"Array length: {arr_len}, pos={self.pos - 4:08X}")
        arr = []
        for _ in range(arr_len):
            obj = self.parse_obj()
            arr.append(obj)
        return arr


if __name__ == "__main__":
    parser = ArgumentParser(description="Parse TJS/ns0 files")
    parser.add_argument("input", help="Input TJS/ns0 file")
    args = parser.parse_args()

    with open(args.input, "rb") as f:
        data = f.read()

    tjs_ns0 = TJS_NS0(data)
    tjs_ns0.parse()
    print("TJS/ns0 file parsed successfully.")
    json_path = splitext(args.input)[0] + ".json"
    with open(json_path, "w", encoding="utf-8") as f:
        dump(tjs_ns0.obj, f, ensure_ascii=False, indent=4)
    print(f"Output written to {json_path}")