"""Freelancer UTF container (.3db, .cmp, .mat, .txm): read, edit and write.

A UTF file is a tree of named nodes; a directory node has children, a leaf node carries bytes.
Header (0x38 bytes): "UTF ", version 0x101, tree offset/size, node size 44, string table
offset/allocated/used, data offset. Node (44 bytes): next sibling offset, name offset, flags
(0x10 directory, 0x80 leaf), 0, child offset (directory) or data offset (leaf), allocated,
size, size again, three FILETIME words.

    python tools/utf.py dump file.3db
"""
import struct
import sys


class Node:
    def __init__(self, name, data=None, children=None):
        self.name = name
        self.data = data              # bytes for a leaf
        self.children = children      # list for a directory

    @property
    def is_dir(self):
        return self.children is not None

    def find(self, path):
        """The node at a '\\'-separated path below this one (case-insensitive), or None."""
        node = self
        for part in [p for p in path.replace('/', '\\').split('\\') if p]:
            if not node.is_dir:
                return None
            nxt = None
            for c in node.children:
                if c.name.lower() == part.lower():
                    nxt = c
                    break
            if nxt is None:
                return None
            node = nxt
        return node

    def walk(self, prefix=''):
        path = prefix + '\\' + self.name if prefix else self.name
        yield path, self
        if self.is_dir:
            for c in self.children:
                yield from c.walk(path)


def read(path):
    raw = open(path, 'rb').read()
    sig, ver, tree_ofs, tree_size, _res, node_size, str_ofs, str_alloc, str_used, data_ofs = \
        struct.unpack_from('<4sIIIIIIIII', raw, 0)
    if sig != b'UTF ':
        raise ValueError(path + ': not a UTF file')
    strings = raw[str_ofs:str_ofs + str_alloc]

    def name_at(o):
        end = strings.index(b'\0', o)
        return strings[o:end].decode('latin-1')

    def node_at(o):
        nxt, name_o, flags, _z, ptr, alloc, size, _size2 = struct.unpack_from('<8I', raw, tree_ofs + o)
        name = name_at(name_o)
        if flags & 0x10:
            children = []
            c = ptr
            while True:
                child, c_next = node_at(c)
                children.append(child)
                if c_next == 0:
                    break
                c = c_next
            return Node(name, children=children), nxt
        return Node(name, data=raw[data_ofs + ptr:data_ofs + ptr + size]), nxt

    root, _ = node_at(0)
    return root


def write(path, root):
    """Writes the tree: nodes in depth-first order, names deduplicated, data 4-byte aligned."""
    nodes = []
    for _, n in root.walk():
        nodes.append(n)
    index = {id(n): i for i, n in enumerate(nodes)}

    names = bytearray(b'\0')
    name_ofs = {}
    for n in nodes:
        if n.name not in name_ofs:
            name_ofs[n.name] = len(names)
            names += n.name.encode('latin-1') + b'\0'
    while len(names) % 4:
        names += b'\0'

    data = bytearray()
    data_ofs = {}
    for n in nodes:
        if not n.is_dir:
            data_ofs[id(n)] = len(data)
            data += n.data
            while len(data) % 4:
                data += b'\0'

    # the next sibling of each node
    sibling = {}
    for n in nodes:
        if n.is_dir:
            for a, b in zip(n.children, n.children[1:]):
                sibling[id(a)] = b

    node_size = 44
    tree = bytearray()
    for n in nodes:
        nxt = sibling.get(id(n))
        nxt_o = index[id(nxt)] * node_size if nxt is not None else 0
        if n.is_dir:
            child_o = index[id(n.children[0])] * node_size if n.children else 0
            tree += struct.pack('<8I3I', nxt_o, name_ofs[n.name], 0x10, 0, child_o, 0, 0, 0, 0, 0, 0)
        else:
            size = len(n.data)
            alloc = (size + 3) & ~3
            tree += struct.pack('<8I3I', nxt_o, name_ofs[n.name], 0x80, 0, data_ofs[id(n)], alloc, size, size,
                                0, 0, 0)

    header_size = 0x38
    tree_ofs = header_size
    str_ofs = tree_ofs + len(tree)
    dat_ofs = str_ofs + len(names)
    header = struct.pack('<4sIIIIIIIIIII', b'UTF ', 0x101, tree_ofs, len(tree), 0, node_size, str_ofs, len(names),
                         len(names), dat_ofs, 0, 0) + b'\0' * 8
    open(path, 'wb').write(header + tree + names + data)


def dump(path):
    root = read(path)
    for p, n in root.walk():
        if n.is_dir:
            print(p + '\\')
        else:
            preview = n.data[:24].hex(' ')
            print('%s  [%d] %s' % (p, len(n.data), preview))


if __name__ == '__main__':
    if len(sys.argv) >= 3 and sys.argv[1] == 'dump':
        for f in sys.argv[2:]:
            print('==', f)
            dump(f)
    else:
        print(__doc__)