"""Bounded-memory UnityFS I/O for this Unity 2019 game.

Uncompressed media stays in a temporary disk file. Asset metadata is still read
and written by UnityPy; only the outer bundle and resource streams are spooled.
"""
from contextlib import contextmanager
import io
import hashlib
from pathlib import Path
import struct
import tempfile
import zlib


def temporary_file():
    directory = Path(__file__).resolve().parent.parent / "data/tomie/tmp"
    directory.mkdir(parents=True, exist_ok=True)
    return tempfile.TemporaryFile(dir=directory)


class FileSlice(io.RawIOBase):
    def __init__(self, backing, offset, size):
        self.backing, self.offset, self.size, self.position = backing, offset, size, 0

    def tell(self):
        return self.position

    def seek(self, offset, whence=0):
        position = offset if whence == 0 else self.position + offset if whence == 1 else self.size + offset
        if not 0 <= position <= self.size:
            raise ValueError("Resource seek outside its bounds")
        self.position = position
        return position

    def read(self, size=-1):
        size = self.size - self.position if size < 0 else min(size, self.size - self.position)
        self.backing.seek(self.offset + self.position)
        result = self.backing.read(size)
        self.position += len(result)
        return result


def read_fs(bundle, reader):
    from UnityPy.enums.BundleFile import ArchiveFlagsOld
    from UnityPy.files.BundleFile import DirectoryInfoFS
    from UnityPy.streams import EndianBinaryReader
    bundle.size = reader.read_long()
    compressed, uncompressed, flags = reader.read_u_int(), reader.read_u_int(), reader.read_u_int()
    if bundle.parse_version()[:2] != (2019, 4) or flags & 0x200:
        raise ValueError("Streaming bundle reader only supports unencrypted Unity 2019.4 assets")
    bundle.dataflags = ArchiveFlagsOld(flags)
    bundle._uses_block_alignment = True
    reader.align_stream(16)
    start = reader.Position
    if flags & 0x80:
        reader.Position = reader.Length - compressed
    info = EndianBinaryReader(bundle.decompress_data(reader.read_bytes(compressed), uncompressed, bundle.dataflags))
    if flags & 0x80:
        reader.Position = start
    info.read_bytes(16)
    blocks = [(info.read_u_int(), info.read_u_int(), info.read_u_short()) for _ in range(info.read_int())]
    nodes = [DirectoryInfoFS(info.read_long(), info.read_long(), info.read_u_int(), info.read_string_to_null())
             for _ in range(info.read_int())]
    bundle._block_info_flags = blocks[0][2] if blocks else 0
    backing = temporary_file()
    crc = 0
    try:
        for index, (raw_size, packed_size, block_flags) in enumerate(blocks):
            data = bundle.decompress_data(reader.read_bytes(packed_size), raw_size, block_flags, index)
            if len(data) != raw_size:
                raise ValueError("Decompressed block size mismatch")
            crc = zlib.crc32(data, crc)
            backing.write(data)
        backing.seek(0)
        bundle.stream_crc = crc & 0xffffffff
        return nodes, EndianBinaryReader(backing)
    except BaseException:
        backing.close()
        raise


@contextmanager
def streaming_unity():
    from UnityPy.files.BundleFile import BundleFile
    from UnityPy.files.File import File
    from UnityPy.files.SerializedFile import SerializedFile
    from UnityPy.helpers import ImportHelper
    from UnityPy.streams import EndianBinaryReader
    old_fs, old_files = BundleFile.read_fs, File.read_files

    def read_files(parent, reader, nodes):
        for node in nodes:
            name = node.path
            if name.endswith((".resS", ".resource")) and hasattr(reader, "stream"):
                source = EndianBinaryReader(FileSlice(reader.stream, node.offset, node.size))
            else:
                reader.Position = node.offset
                source = EndianBinaryReader(reader.read(node.size), offset=reader.BaseOffset + node.offset)
            obj = ImportHelper.parse_file(source, parent, name, is_dependency=parent.is_dependency)
            if isinstance(obj, (EndianBinaryReader, SerializedFile)) and parent.environment:
                parent.environment.register_cab(name, obj)
            obj.flags = getattr(node, "flags", 0)
            parent.files[name] = obj

    BundleFile.read_fs, File.read_files = read_fs, read_files
    try:
        yield
    finally:
        BundleFile.read_fs, File.read_files = old_fs, old_files


def load(path):
    import UnityPy
    with Path(path).open("rb") as stream, streaming_unity():
        return UnityPy.load(stream)


def resource_hashes(bundle):
    from UnityPy.streams import EndianBinaryReader
    result = {}
    for name, obj in bundle.files.items():
        if not isinstance(obj, EndianBinaryReader):
            continue
        before = obj.Position
        obj.Position = 0
        fingerprint = hashlib.sha256()
        while obj.Position < obj.Length:
            fingerprint.update(obj.read(min(1 << 20, obj.Length - obj.Position)))
        obj.Position = before
        result[name] = fingerprint.hexdigest()
    return result


def save(bundle, path):
    from lz4.block import compress
    from UnityPy.streams import EndianBinaryReader, EndianBinaryWriter
    if bundle.signature != "UnityFS" or bundle.parse_version()[:2] != (2019, 4):
        raise ValueError("Unsupported bundle format")
    header = EndianBinaryWriter()
    header.write_string_to_null(bundle.signature)
    header.write_u_int(bundle.version)
    header.write_string_to_null(bundle.version_player)
    header.write_string_to_null(bundle.version_engine)
    size_position = header.Position
    header.write_long(0)
    header.write_u_int(0)
    header.write_u_int(0)
    header.write_u_int(194)  # LZ4, directory included, block info at end.
    header.align_stream(16)
    nodes, blocks = [], []
    crc, offset = 0, 0
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w+b") as output, temporary_file() as raw:
        output.write(header.bytes)
        for name, obj in bundle.files.items():
            start = raw.tell()
            if isinstance(obj, EndianBinaryReader):
                obj.Position = 0
                while obj.Position < obj.Length:
                    raw.write(obj.read(min(1 << 20, obj.Length - obj.Position)))
            else:
                raw.write(obj.bytes if isinstance(obj, EndianBinaryWriter) else obj.save())
            size = raw.tell() - start
            nodes.append((offset, size, getattr(obj, "flags", 0), name))
            offset += size
        raw.seek(0)
        while data := raw.read(128 << 10):
            packed = compress(data, store_size=False)
            output.write(packed)
            blocks.append((len(data), len(packed), 2))
            crc = zlib.crc32(data, crc)
        info = EndianBinaryWriter(b"\0" * 16)
        info.write_int(len(blocks))
        for raw_size, packed_size, flags in blocks:
            info.write_u_int(raw_size)
            info.write_u_int(packed_size)
            info.write_u_short(flags)
        info.write_int(len(nodes))
        for node_offset, size, flags, name in nodes:
            info.write_long(node_offset)
            info.write_long(size)
            info.write_u_int(flags)
            info.write_string_to_null(name)
        raw_info = info.bytes
        packed_info = compress(raw_info, store_size=False)
        output.write(packed_info)
        size = output.tell()
        output.seek(size_position)
        output.write(struct.pack(">qIII", size, len(packed_info), len(raw_info), 194))
    return crc & 0xffffffff