"""The game's UE3 packages (.upk): summary, LZO-compressed chunks, block-level reads and patches.

Offsets used by the rest of ackit are positions in the *uncompressed* package (what Gildor's
decompress.exe writes), so a sound's location does not depend on how the file is compressed.
Format and rules: docs/formats/upk.md.

Patching rewrites only the blocks that contain changed bytes (stored as one LZO literal run, see
lzo.compress_stored) and moves the chunks after them; every other byte of the file stays identical.
"""
import struct
import threading
from collections import OrderedDict, namedtuple

from . import lzo

TAG = 0x9E2A83C1
COMPRESS_LZO = 2


class PackageError(ValueError):
    pass


Export = namedtuple("Export", "index cls outer name size offset entry")   # entry: table position (uncompressed)


class Block:
    __slots__ = ("chunk", "uoff", "usize", "coff", "csize")

    def __init__(self, chunk, uoff, usize, coff, csize):
        self.chunk, self.uoff, self.usize, self.coff, self.csize = chunk, uoff, usize, coff, csize


class Package:
    """Read access to one package file. `blocks` covers the uncompressed range [header_size, size)."""

    def __init__(self, path, cache_blocks=8):
        self.path = path
        with open(path, "rb") as f:
            head = f.read(0x10000)
            self._parse_summary(head, f)
            self.file_size = f.seek(0, 2)
        self._cache = OrderedDict()
        self._cache_n = cache_blocks
        self._lock = threading.Lock()

    # summary -----------------------------------------------------------------------------------
    def _parse_summary(self, b, f):
        if len(b) < 64 or struct.unpack_from("<I", b, 0)[0] != TAG:
            raise PackageError("%s is not an Unreal package" % self.path)
        self.version, self.licensee = struct.unpack_from("<HH", b, 4)
        o = 12
        n = struct.unpack_from("<i", b, o)[0]
        o += 4 + max(n, 0)
        # PackageFlags, NameCount/Offset, ExportCount/Offset, ImportCount/Offset, DependsOffset,
        # ImportExportGuidsOffset, ImportGuidsCount, ExportGuidsCount, ThumbnailTableOffset
        f12 = struct.unpack_from("<12i", b, o)
        self.name_count, self.name_offset = f12[1], f12[2]
        self.export_count, self.export_offset = f12[3], f12[4]
        self.import_count, self.import_offset = f12[5], f12[6]
        o += 48
        self.guid = b[o:o + 16]
        o += 16
        ng = struct.unpack_from("<i", b, o)[0]
        o += 4 + 12 * ng
        self.engine, self.cooker, self.compression = struct.unpack_from("<3I", b, o)
        o += 12
        nc = struct.unpack_from("<i", b, o)[0]
        o += 4
        if nc and len(b) < o + 16 * nc:
            f.seek(0)
            b = f.read(o + 16 * nc)
        self.chunk_table_offset = o
        self.chunks = [list(struct.unpack_from("<4i", b, o + 16 * i)) for i in range(nc)]  # uoff usize coff csize
        if nc and self.compression != COMPRESS_LZO:
            raise PackageError("%s: unsupported compression %d" % (self.path, self.compression))
        self.blocks = []
        if not nc:
            return
        for ci, (uoff, usize, coff, csize) in enumerate(self.chunks):
            f.seek(coff)
            h = f.read(16)
            tag, bsize, sum_c, sum_u = struct.unpack("<4I", h)
            if tag != TAG or sum_u != usize:
                raise PackageError("%s: bad chunk header at 0x%X" % (self.path, coff))
            nb = (sum_u + bsize - 1) // bsize
            tbl = f.read(8 * nb)
            p = coff + 16 + 8 * nb
            u = uoff
            for i in range(nb):
                c, uu = struct.unpack_from("<2I", tbl, 8 * i)
                self.blocks.append(Block(ci, u, uu, p, c))
                p += c
                u += uu
            if p != coff + csize or u != uoff + usize:
                raise PackageError("%s: inconsistent chunk %d" % (self.path, ci))
        self.block_size = bsize

    @property
    def compressed(self):
        return bool(self.chunks)

    @property
    def header_size(self):
        """Uncompressed size of the summary area (= start of the first chunk's data)."""
        return self.chunks[0][0] if self.chunks else 0

    @property
    def size(self):
        """Uncompressed package size."""
        if not self.chunks:
            return self.file_size
        return self.chunks[-1][0] + self.chunks[-1][1]

    # tables (need the uncompressed data: `read_all`) ---------------------------------------------
    def tables(self, data=None, start=None):
        """(names, imports, exports) of the package. imports: [(class name, object name, outer)],
        exports: [Export]. Layouts: docs/formats/upk.md (licensee 101 adds one u32 to each export)."""
        if data is None:
            start, data = self.read_all()
        buf = memoryview(data)
        names = []
        o = self.name_offset - start
        for _ in range(self.name_count):
            n = struct.unpack_from("<i", buf, o)[0]
            o += 4
            if n >= 0:
                names.append(bytes(buf[o:o + n - 1]).decode("latin-1"))
                o += n
            else:
                names.append(bytes(buf[o:o - 2 * n - 2]).decode("utf-16-le", "replace"))
                o += -2 * n
            o += 8                                   # name flags
        imports = []
        o = self.import_offset - start
        for _ in range(self.import_count):
            _pkg, _pn, cls, _cn, outer, nm, nn = struct.unpack_from("<7i", buf, o)
            imports.append((names[cls], names[nm], outer))
            o += 28
        exports = []
        o = self.export_offset - start
        for i in range(self.export_count):
            cls, sup, outer, nm, nn, arch, _x = struct.unpack_from("<7i", buf, o)
            size, off = struct.unpack_from("<ii", buf, o + 36)
            ngen = struct.unpack_from("<i", buf, o + 48)[0]
            exports.append(Export(i + 1, cls, outer, names[nm] + ("_%d" % (nn - 1) if nn else ""), size, off,
                                  o + start))
            o += 72 + 4 * ngen
        return names, imports, exports

    def fingerprint(self):
        """Cheap identity of the package contents: guid + chunk table (changes with any re-cook/patch)."""
        import hashlib
        h = hashlib.sha1(self.guid)
        h.update(struct.pack("<q", self.file_size))
        for c in self.chunks:
            h.update(struct.pack("<4i", *c))
        return h.hexdigest()[:16]

    # reading -----------------------------------------------------------------------------------
    def _block(self, i, f=None):
        with self._lock:                        # the editor reads one package from several background jobs
            data = self._cache.get(i)
            if data is not None:
                self._cache.move_to_end(i)
                return data
        bl = self.blocks[i]
        own = f is None
        f = f or open(self.path, "rb")
        try:
            f.seek(bl.coff)
            data = lzo.decompress(f.read(bl.csize), bl.usize)
        finally:
            if own:
                f.close()
        with self._lock:
            self._cache[i] = data
            if len(self._cache) > self._cache_n:
                self._cache.popitem(last=False)
        return data

    def _block_index(self, off):
        lo, hi = 0, len(self.blocks) - 1
        while lo <= hi:
            mid = (lo + hi) // 2
            bl = self.blocks[mid]
            if off < bl.uoff:
                hi = mid - 1
            elif off >= bl.uoff + bl.usize:
                lo = mid + 1
            else:
                return mid
        raise PackageError("offset 0x%X is outside the package data" % off)

    def read(self, off, n):
        """Uncompressed bytes [off, off+n)."""
        if not self.chunks:
            with open(self.path, "rb") as f:
                f.seek(off)
                return f.read(n)
        if off < self.header_size:
            raise PackageError("reading the package summary area is not supported")
        out = bytearray()
        i = self._block_index(off)
        with open(self.path, "rb") as f:
            while len(out) < n:
                bl = self.blocks[i]
                d = self._block(i, f)
                s = off + len(out) - bl.uoff
                out += d[s:s + n - len(out)]
                i += 1
        return bytes(out)

    def iter_data(self):
        """(uncompressed offset, bytes) per block, in order (whole-package scans)."""
        if not self.chunks:
            with open(self.path, "rb") as f:
                yield 0, f.read()
            return
        with open(self.path, "rb") as f:
            for bl in self.blocks:
                f.seek(bl.coff)
                yield bl.uoff, lzo.decompress(f.read(bl.csize), bl.usize)

    def read_all(self):
        """The whole uncompressed data area as one bytes object (offset `data_start`)."""
        parts = list(self.iter_data())
        return (parts[0][0] if parts else 0), b"".join(p for _, p in parts)

    # patching ----------------------------------------------------------------------------------
    def patched(self, edits):
        """edits: [(uncompressed offset, bytes)] (same length as what they overwrite) -> new file bytes.

        Uncompressed packages are patched in place. For compressed ones the touched blocks are
        re-encoded and the chunk headers and the summary's chunk table are updated."""
        with open(self.path, "rb") as f:
            raw = f.read()
        if not self.chunks:
            out = bytearray(raw)
            for off, data in edits:
                if off + len(data) > len(out):
                    raise PackageError("patch beyond the end of %s" % self.path)
                out[off:off + len(data)] = data
            return bytes(out)
        touched = {}
        for off, data in edits:
            if off < self.header_size or off + len(data) > self.size:
                raise PackageError("patch outside the package data (0x%X)" % off)
            pos = 0
            while pos < len(data):
                i = self._block_index(off + pos)
                bl = self.blocks[i]
                if i not in touched:
                    touched[i] = bytearray(lzo.decompress(raw[bl.coff:bl.coff + bl.csize], bl.usize))
                s = off + pos - bl.uoff
                k = min(len(data) - pos, bl.usize - s)
                touched[i][s:s + k] = data[pos:pos + k]
                pos += k
        # rebuild chunk by chunk
        head = bytearray(raw[:self.chunks[0][2]])
        body = bytearray()
        by_chunk = {}
        for i, bl in enumerate(self.blocks):
            by_chunk.setdefault(bl.chunk, []).append(i)
        for ci, (uoff, usize, coff, csize) in enumerate(self.chunks):
            new_coff = len(head) + len(body)
            idx = by_chunk[ci]
            if not any(i in touched for i in idx):
                body += raw[coff:coff + csize]
                self._set_chunk(head, ci, uoff, usize, new_coff, csize)
                continue
            datas = []
            for i in idx:
                bl = self.blocks[i]
                if i in touched:
                    datas.append((lzo.compress_stored(touched[i]), bl.usize))
                else:
                    datas.append((raw[bl.coff:bl.coff + bl.csize], bl.usize))
            hdr = struct.pack("<4I", TAG, self.block_size, sum(len(d) for d, _ in datas), usize)
            hdr += b"".join(struct.pack("<2I", len(d), u) for d, u in datas)
            chunk = hdr + b"".join(d for d, _ in datas)
            body += chunk
            self._set_chunk(head, ci, uoff, usize, new_coff, len(chunk))
        return bytes(head + body)

    def _set_chunk(self, head, ci, uoff, usize, coff, csize):
        struct.pack_into("<4i", head, self.chunk_table_offset + 16 * ci, uoff, usize, coff, csize)


def is_package(path):
    try:
        with open(path, "rb") as f:
            return struct.unpack("<I", f.read(4))[0] == TAG
    except (OSError, struct.error):
        return False