from __future__ import annotations

import io
import os
import sys
import zlib
from bisect import bisect_left, bisect_right
from threading import Lock
from typing import BinaryIO

STREAM_BUFFER_SIZE = int(os.getenv("DISSECT_STREAM_BUFFER_SIZE", io.DEFAULT_BUFFER_SIZE))


class AlignedStream(io.RawIOBase):
    """Basic buffered stream that provides aligned reads.

    Must be subclassed for various stream implementations. Subclasses can implement:
        - :meth:`~AlignedStream._read`
        - :meth:`~AlignedStream._seek`

    The offset and length for ``_read`` are guaranteed to be aligned for streams of a known size.
    If your stream has an unknown size (i.e. ``size == None``), reads of length ``-1`` (i.e. read until EOF) will be
    passed through to your implementation of ``_read``.
    The only time that overriding ``_seek`` would make sense is if there's no known size of your stream,
    but still want to provide ``SEEK_END`` functionality.

    Most subclasses of ``AlignedStream`` take one or more file-like objects as source.
    Operations on these subclasses, like reading, will modify the source file-like object as a side effect.

    Args:
        size: The size of the stream. This is used in read and seek operations. ``None`` if unknown.
        align: The alignment size. Read operations are aligned on this boundary. Also determines buffer size.

    .. automethod:: _read
    .. automethod:: _seek
    """

    def __init__(self, size: int | None = None, align: int = STREAM_BUFFER_SIZE):
        super().__init__()
        self.size = size
        self.align = align

        self._pos = 0
        self._pos_align = 0

        self._buf = None
        self._lock = Lock()

    def readable(self) -> bool:
        """Indicate that the stream is readable."""
        return True

    def seekable(self) -> bool:
        """Indicate that the stream is seekable."""
        return True

    def seek(self, pos: int, whence: int = io.SEEK_SET) -> int:
        """Seek the stream to the specified position.

        Returns:
            The new stream position after seeking.
        """
        with self._lock:
            pos = self._seek(pos, whence)
            self._set_pos(pos)

            return pos

    def _seek(self, pos: int, whence: int = io.SEEK_SET) -> int:
        """Calculate and return the new stream position after a seek."""
        if whence == io.SEEK_SET:
            if pos < 0:
                raise ValueError(f"negative seek position {pos}")
        elif whence == io.SEEK_CUR:
            pos = max(0, self._pos + pos)
        elif whence == io.SEEK_END:
            if self.size is None:
                raise IOError("invalid whence value for stream with no size")
            pos = max(0, self.size + pos)
        else:
            raise IOError("invalid whence value")

        return pos

    def _set_pos(self, pos: int) -> None:
        """Update the position and aligned position within the stream."""
        new_pos_align = pos - (pos % self.align)

        if self._pos_align != new_pos_align:
            self._pos_align = new_pos_align
            self._buf = None

        self._pos = pos

    def tell(self) -> int:
        """Return current stream position."""
        return self._pos

    def _fill_buf(self) -> None:
        """Fill the alignment buffer if we can."""
        if self._buf or (self.size is not None and (self.size <= self._pos or self.size <= self._pos_align)):
            # Don't fill the buffer if:
            # - We already have a buffer
            # - The stream position is at the end (or beyond) the stream size
            return

        self._buf = self._read(self._pos_align, self.align)

    def read(self, n: int = -1) -> bytes:
        """Read and return up to ``n`` bytes, or read to the end of the stream if ``n`` is ``-1``.

        Returns an empty bytes object on EOF.
        """
        if n is not None and n < -1:
            raise ValueError("invalid number of bytes to read")

        r = []
        size = self.size
        align = self.align

        with self._lock:
            if size is None and n == -1:
                r = []
                if self._buf:
                    buffer_pos = self._pos - self._pos_align
                    r.append(self._buf[buffer_pos:])
                    self._set_pos(self._pos_align + align)

                r.append(self._read(self._pos_align, -1))

                buf = b"".join(r)
                self._set_pos(self._pos + len(buf))
                return buf

            # If we know the stream size, adjust n
            if size is not None:
                remaining = size - self._pos
                n = remaining if n == -1 else min(n, remaining)

            # Short path for when it turns out we don't need to read anything
            if n == 0 or (size is not None and size <= self._pos):
                return b""

            # Read misaligned start from buffer
            if self._pos != self._pos_align:
                self._fill_buf()

                buffer_pos = self._pos - self._pos_align
                remaining = align - buffer_pos
                buffer_len = min(n, remaining)

                r.append(self._buf[buffer_pos : buffer_pos + buffer_len])

                n -= buffer_len
                self._set_pos(self._pos + buffer_len)

            # Aligned blocks
            if n >= align:
                count, n = divmod(n, align)

                read_len = count * align
                r.append(self._read(self._pos, read_len))

                self._set_pos(self._pos + read_len)

            # Misaligned remaining bytes
            if n > 0:
                self._fill_buf()
                r.append(self._buf[:n])
                self._set_pos(self._pos + n)

            return b"".join(r)

    def readinto(self, b: bytearray) -> int:
        """Read bytes into a pre-allocated bytes-like object b.

        Returns an int representing the number of bytes read (0 for EOF).
        """
        buf = self.read(len(b))
        length = len(buf)
        b[:length] = buf
        return length

    def _read(self, offset: int, length: int) -> bytes:
        """Read method that backs this aligned stream."""
        raise NotImplementedError("_read needs to be implemented by subclass")

    def readall(self) -> bytes:
        """Read until end of stream."""
        return self.read()

    def readoffset(self, offset: int, length: int) -> bytes:
        """Convenience method to read from a given offset.

        Args:
            offset: The offset in the stream to read from.
            length: The number of bytes to read.
        """
        self.seek(offset)
        return self.read(length)

    def peek(self, n: int) -> bytes:
        """Convenience method to peek from the current offset without advancing the stream position.

        Args:
            n: The number of bytes to peek.
        """
        pos = self._pos
        data = self.read(n)
        self._set_pos(pos)
        return data

    def close(self) -> None:
        """Close the stream. Does nothing by default."""


class RangeStream(AlignedStream):
    """Create a stream with a specific range from another file-like object.

    ASCII representation::

        Source file-like object
        |................................................|
                RangeStream with offset and size
                |............................|

    Args:
        fh: The source file-like object.
        offset: The offset the stream should start from on the source file-like object.
        size: The size the stream should be.
        align: The alignment size.
    """

    def __init__(self, fh: BinaryIO, offset: int, size: int | None, align: int = STREAM_BUFFER_SIZE):
        super().__init__(size, align)
        self._fh = fh
        self.offset = offset

    def _seek(self, pos: int, whence: int = io.SEEK_SET) -> int:
        if self.size is None and whence == io.SEEK_END:
            if (pos := self._fh.seek(pos, whence)) is None:
                pos = self._fh.tell()
            return max(0, pos - self.offset)
        return super()._seek(pos, whence)

    def _read(self, offset: int, length: int) -> bytes:
        read_length = min(length, self.size - offset) if self.size else length
        self._fh.seek(self.offset + offset)
        return self._fh.read(read_length)


class RelativeStream(RangeStream):
    """Create a relative stream from another file-like object.

    ASCII representation::

        Source file-like object
        |................................................|
                RelativeStream with offset
                |........................................|

    Args:
        fh: The source file-like object.
        offset: The offset the stream should start from on the source file-like object.
        size: Optional size the stream should be.
        align: The alignment size.
    """

    def __init__(self, fh: BinaryIO, offset: int, size: int | None = None, align: int = STREAM_BUFFER_SIZE):
        super().__init__(fh, offset, size, align)


class BufferedStream(RelativeStream):
    """Create a buffered stream from another file-like object.

    Optionally start from a specific offset.

    Args:
        fh: The source file-like object.
        offset: The offset the stream should start from.
        size: The size the stream should be.
        align: The alignment size.
    """

    def __init__(self, fh: BinaryIO, offset: int = 0, size: int | None = None, align: int = STREAM_BUFFER_SIZE):
        super().__init__(fh, offset, size, align)


class MappingStream(AlignedStream):
    """Create a stream from multiple mapped file-like objects.

    Args:
        size: The size the stream should be.
        align: The alignment size.
    """

    def __init__(self, size: int | None = None, align: int = STREAM_BUFFER_SIZE):
        super().__init__(size, align)
        self._runs: list[tuple[int, int, BinaryIO, int]] = []

    def add(self, offset: int, size: int, fh: BinaryIO, file_offset: int = 0) -> None:
        """Add a file-like object to the stream.

        Args:
            offset: The offset in the stream this fh maps to.
            size: The size that this mapped fh spans in the stream.
            fh: The file-like object to map.
            file_offset: The offset in the fh to start from.

        Note that there is no check on overlapping offsets and/or sizes.
        """
        self._runs.append((offset, size, fh, file_offset))
        self._runs = sorted(self._runs, key=lambda run: run[0])
        self._buf = None
        self.size = self._runs[-1][0] + self._runs[-1][1]

    def _get_run_idx(self, offset: int) -> tuple[int, int, BinaryIO, int]:
        """Find a mapping run for a given offset.

        Args:
            offset: The offset to find a mapping for.

        Returns:
            The run tuple if found.

        Raises:
            IOError: If no mapping is found for the given offset.
        """
        for idx, run in enumerate(self._runs):
            if run[0] <= offset < run[0] + run[1]:
                return idx

        raise EOFError(f"No mapping for offset {offset}")

    def _read(self, offset: int, length: int) -> bytes:
        result = []

        run_idx = self._get_run_idx(offset)
        runlist_len = len(self._runs)
        size = self.size

        while length > 0:
            if run_idx >= runlist_len:
                # We somehow requested more data than we have runs for
                break

            run_offset, run_size, run_fh, run_file_offset = self._runs[run_idx]

            if run_offset > offset:
                # We landed in a gap, stop reading
                break

            run_pos = offset - run_offset
            run_remaining = run_size - run_pos

            if run_remaining < 0:
                break

            read_count = min(size - offset, min(run_remaining, length))

            run_fh.seek(run_file_offset + run_pos)
            result.append(run_fh.read(read_count))

            offset += read_count
            length -= read_count
            run_idx += 1

        return b"".join(result)


class RunlistStream(AlignedStream):
    """Create a stream from multiple runs on another file-like object.

    This is common in filesystems, where file data information is stored in "runs".
    A run is a ``(block_offset, block_count)`` tuple, meaning the amount of consecutive blocks from a
    specific starting block. A block_offset of ``None`` represents a sparse run, meaning it must simply
    return all ``\\x00`` bytes.

    Args:
        fh: The source file-like object.
        runlist: The runlist for this stream in block units.
        size: The size of the stream. This can be smaller than the total sum of blocks (to account for slack space).
        block_size: The block size in bytes.
        align: Optional alignment that differs from the block size, otherwise ``block_size`` is used as alignment.
    """

    def __init__(
        self, fh: BinaryIO, runlist: list[tuple[int, int]], size: int, block_size: int, align: int | None = None
    ):
        super().__init__(size, align or block_size)

        if isinstance(fh, RunlistStream):
            self._fh = fh._fh
        else:
            self._fh = fh

        self._runlist = []
        self._runlist_offsets = []

        self.runlist = runlist
        self.block_size = block_size

    @property
    def runlist(self) -> list[tuple[int, int]]:
        return self._runlist

    @runlist.setter
    def runlist(self, runlist: list[tuple[int, int]]) -> None:
        self._runlist = runlist
        self._runlist_offsets = []

        offset = 0
        # Create a list of starting offsets for each run so we can bisect that quickly when reading
        for _, block_count in self._runlist:
            if offset != 0:
                self._runlist_offsets.append(offset)
            offset += block_count

        self._buf = None

    def _read(self, offset: int, length: int) -> bytes:
        r = []

        block_offset = offset // self.block_size

        run_idx = bisect_right(self._runlist_offsets, block_offset)
        runlist_len = len(self.runlist)
        size = self.size

        while length > 0:
            if run_idx >= runlist_len:
                # We somehow requested more data than we have runs for
                break

            # If run_idx == 0, we only have a single run
            run_block_pos = 0 if run_idx == 0 else self._runlist_offsets[run_idx - 1]
            run_block_offset, run_block_count = self.runlist[run_idx]

            run_size = run_block_count * self.block_size
            run_pos = offset - run_block_pos * self.block_size
            run_remaining = run_size - run_pos

            # Sometimes the self.size is way larger than what we actually have runs for?
            # Stop reading if we reach a negative run_remaining
            if run_remaining < 0:
                break

            read_count = min(size - offset, min(run_remaining, length))

            # Sparse run
            if run_block_offset is None:
                r.append(b"\x00" * read_count)
            else:
                self._fh.seek(run_block_offset * self.block_size + run_pos)
                r.append(self._fh.read(read_count))

            offset += read_count
            length -= read_count
            run_idx += 1

        return b"".join(r)


class OverlayStream(AlignedStream):
    """Create a stream from another file-like object with the ability to overlay other streams or bytes.

    Useful for patching large file-like objects without having to cache the entire contents.
    First wrap the original stream in this class, and then call ``add()`` with the offset and data to overlay.

    Args:
        fh: The source file-like object.
        size: The size the stream should be.
        align: The alignment size.
    """

    def __init__(self, fh: BinaryIO, size: int | None = None, align: int = STREAM_BUFFER_SIZE):
        super().__init__(size, align)
        self._fh = fh
        self.overlays: dict[int, tuple[int, BinaryIO]] = {}
        self._lookup: list[int] = []

    def add(self, offset: int, data: bytes | BinaryIO, size: int | None = None) -> None:
        """Add an overlay at the given offset.

        Args:
            offset: The offset in bytes to add an overlay at.
            data: The bytes or file-like object to overlay.
            size: Optional size specification of the overlay, if it can't be inferred.
        """
        if not hasattr(data, "read"):
            size = size or len(data)
            data = io.BytesIO(data)
        elif size is None:
            size = data.size if hasattr(data, "size") else data.seek(0, io.SEEK_END)

        if not size:
            return None

        if size < 0:
            raise ValueError("Size must be positive")

        # Check if there are overlapping overlays
        for other_offset, (other_size, _) in self.overlays.items():
            if other_offset < offset + size and offset < other_offset + other_size:
                raise ValueError(f"Overlap with existing overlay: ({other_offset, other_size})")

        self.overlays[offset] = (size, data)
        self._lookup.append(offset)
        self._lookup.sort()

        # Clear the buffer if we add an overlay at our current position
        if self._buf and (self._pos_align <= offset + size and offset <= self._pos_align + len(self._buf)):
            self._buf = None

        return self

    def _read(self, offset: int, length: int) -> bytes:
        result = []

        fh = self._fh
        overlays = self.overlays
        lookup = self._lookup

        overlay_len = len(overlays)
        overlay_idx = bisect_left(lookup, offset)

        while length > 0:
            prev_overlay_offset = None if overlay_idx == 0 else lookup[overlay_idx - 1]
            next_overlay_offset = None if overlay_idx >= overlay_len else lookup[overlay_idx]

            if prev_overlay_offset is not None:
                prev_overlay_size, prev_overlay_data = overlays[prev_overlay_offset]
                prev_overlay_end = prev_overlay_offset + prev_overlay_size

                if prev_overlay_end > offset:
                    # Currently in an overlay
                    offset_in_prev_overlay = offset - prev_overlay_offset
                    prev_overlay_remaining = prev_overlay_size - offset_in_prev_overlay
                    prev_overlay_read_size = min(length, prev_overlay_remaining)

                    prev_overlay_data.seek(offset_in_prev_overlay)
                    result.append(prev_overlay_data.read(prev_overlay_read_size))

                    offset += prev_overlay_read_size
                    length -= prev_overlay_read_size

            if length == 0:
                break

            if next_overlay_offset is not None:
                next_overlay_size, next_overlay_data = overlays[next_overlay_offset]
                gap_to_next_overlay = next_overlay_offset - offset

                if 0 <= gap_to_next_overlay < length:
                    if gap_to_next_overlay:
                        fh.seek(offset)
                        result.append(fh.read(gap_to_next_overlay))

                    # read remaining from overlay
                    next_overlay_read_size = min(next_overlay_size, length - gap_to_next_overlay)
                    next_overlay_data.seek(0)
                    result.append(next_overlay_data.read(next_overlay_read_size))

                    offset += next_overlay_read_size + gap_to_next_overlay
                    length -= next_overlay_read_size + gap_to_next_overlay
                else:
                    # Next overlay is too far away, complete read
                    fh.seek(offset)
                    result.append(fh.read(length))
                    break
            else:
                # No next overlay, complete read
                fh.seek(offset)
                result.append(fh.read(length))
                break

            overlay_idx += 1

        return b"".join(result)


class ZlibStream(AlignedStream):
    """Create a zlib stream from another file-like object.

    Basically the same as ``gzip.GzipFile`` but for raw zlib streams.
    Due to the nature of zlib streams, seeking backwards requires resetting the decompression context.

    Args:
        fh: The source file-like object.
        size: The size the stream should be.
    """

    def __init__(self, fh: BinaryIO, size: int | None = None, align: int = STREAM_BUFFER_SIZE, **kwargs):
        self._fh = fh

        self._zlib = None
        self._zlib_args = kwargs
        self._zlib_offset = 0
        self._zlib_prepend = b""
        self._zlib_prepend_offset = None
        self._rewind()

        super().__init__(size, align)

    def _rewind(self) -> None:
        self._fh.seek(0)
        self._zlib = zlib.decompressobj(**self._zlib_args)
        self._zlib_offset = 0
        self._zlib_prepend = b""
        self._zlib_prepend_offset = None

    def _seek_zlib(self, offset: int) -> None:
        if offset < self._zlib_offset:
            self._rewind()

        while self._zlib_offset < offset:
            read_size = min(offset - self._zlib_offset, self.align)
            if self._read_zlib(read_size) == b"":
                break

    def _read_fh(self, length: int) -> bytes:
        if self._zlib_prepend_offset is None:
            return self._fh.read(length)

        if self._zlib_prepend_offset + length <= len(self._zlib_prepend):
            offset = self._zlib_prepend_offset
            self._zlib_prepend_offset += length
            return self._zlib_prepend[offset : self._zlib_prepend_offset]

        offset = self._zlib_prepend_offset
        self._zlib_prepend_offset = None
        return self._zlib_prepend[offset:] + self._fh.read(length - len(self._zlib_prepend) + offset)

    def _read_zlib(self, length: int) -> bytes:
        if length < 0:
            return self.readall()

        result = []
        while length > 0:
            buf = self._read_fh(io.DEFAULT_BUFFER_SIZE)
            decompressed = self._zlib.decompress(buf, length)

            if self._zlib.unconsumed_tail != b"":
                self._zlib_prepend = self._zlib.unconsumed_tail
                self._zlib_prepend_offset = 0

            if buf == b"":
                break

            result.append(decompressed)
            length -= len(decompressed)

        buf = b"".join(result)
        self._zlib_offset += len(buf)
        return buf

    def _read(self, offset: int, length: int) -> bytes:
        self._seek_zlib(offset)
        return self._read_zlib(length)

    def readall(self) -> bytes:
        self._seek_zlib(self.tell())

        chunks = []
        # sys.maxsize means the max length of output buffer is unlimited,
        # so that the whole input buffer can be decompressed within one
        # .decompress() call.
        while data := self._read_zlib(sys.maxsize):
            chunks.append(data)

        return b"".join(chunks)


class BitStream:
    """Bit-level stream reader.

    Args:
        fh: File-like object to read bits from.
    """

    def __init__(self, fh: BinaryIO):
        self.fh = fh
        self._byte_offset = fh.tell()

        self.buffer = 0
        self.bits = 0

    def readable(self) -> bool:
        """Indicate that the stream is readable."""
        return True

    def seekable(self) -> bool:
        """Indicate that the stream is seekable."""
        return True

    def writable(self) -> bool:
        """Indicate that the stream is not writable."""
        return False

    def seek(self, pos: int, whence: int = io.SEEK_SET) -> int:
        """Seek the stream to the specified position in bits.

        Returns:
            The new stream position after seeking.
        """
        if whence == io.SEEK_SET:
            byte_pos, bit_pos = divmod(pos, 8)
        elif whence == io.SEEK_CUR:
            current_pos = self.tell()
            byte_pos, bit_pos = divmod(current_pos + pos, 8)
        elif whence == io.SEEK_END:
            self.fh.seek(0, io.SEEK_END)
            end_pos = self.fh.tell() * 8
            byte_pos, bit_pos = divmod(end_pos + pos, 8)
        else:
            raise IOError("invalid whence value")

        self._byte_offset = byte_pos
        self.bits = 0
        self.buffer = 0

        if bit_pos > 0:
            self.read(bit_pos)

        return self.tell()

    def tell(self) -> int:
        """Get the current position in the stream in bits."""
        return (self._byte_offset * 8) - self.bits

    def read(self, n: int) -> int:
        """Read n bits from the stream.

        Args:
            n: Number of bits to read.
        """
        value = self.peek(n)
        self.remove(n)
        return value

    def peek(self, n: int) -> int:
        """Peek n bits from the stream without advancing.

        Args:
            n: Number of bits to peek.
        """
        if n == 0:
            return 0

        if n > self.bits:
            while self.bits < n:
                num_bytes = (n - self.bits + 7) // 8
                self.fh.seek(self._byte_offset)
                if not (buf := self.fh.read(num_bytes)):
                    break

                new_bits = int.from_bytes(buf, "big")
                num_new_bits = len(buf) * 8
                self.buffer = (self.buffer << num_new_bits) | new_bits
                self.bits += num_new_bits
                self._byte_offset += len(buf)

        return self.buffer >> (self.bits - min(n, self.bits))

    def remove(self, n: int) -> None:
        """Remove n bits from the stream.

        Args:
            n: Number of bits to remove.
        """
        self.bits -= min(n, self.bits)
        self.buffer &= (1 << (self.bits)) - 1
