from collections.abc import Iterator
from typing import BinaryIO


import os
import time

from .reader import FileReader, Chunk
from ..constants import CH_DATA, CH_ALLOC, zeros


class ChunkerFixed:
    """
    This is a simple chunker for input data with data usually staying at the same
    offset and/or with known block/record sizes:

    - raw disk images
    - block devices
    - database files with simple header + fixed-size records layout

    It optionally supports:

    - a header block of different size
    - using a sparsemap to read only data ranges and seek over hole ranges
      for sparse files.
    - using an externally given filemap to read only specific ranges from
      a file.

    Note: the last block of a data or hole range may be less than the block size,
          this is supported and not considered to be an error.

    Chunks are yielded as memoryviews over per-chunk buffers that the file data is
    read into directly - each byte is copied exactly once, from the reader's block
    buffer into the chunk (all-zero chunks are yielded as allocation CH_ALLOC with
    data None). Call release_chunk_data() after consuming a chunk.
    """

    def __init__(self, block_size: int, header_size: int = 0, sparse: bool = False) -> None:
        self.block_size = block_size
        self.header_size = header_size
        self.chunking_time = 0.0  # likely will stay close to zero - not much to do here.
        self.reader_block_size = 1024 * 1024
        self.reader: FileReader | None = None
        self.sparse = sparse

    def chunkify(
        self, fd: BinaryIO | None = None, fh: int = -1, fmap: list | None = None, st: os.stat_result | None = None
    ) -> Iterator:
        """
        Cut a file into chunks.

        :param fd: Python file object
        :param fh: OS-level file handle (if available),
                   defaults to -1 which means not to use OS-level fd.
        :param fmap: a file map, same format as generated by sparsemap
        :param st: the file's os.stat_result, if the caller has it anyway -
                   saves FileReader a stat call, see there.
        """
        # Initialize the reader with the file descriptors
        self.reader = FileReader(fd=fd, fh=fh, read_size=self.reader_block_size, sparse=self.sparse, fmap=fmap, st=st)

        # Read the header block first, if there is one; after it, read block_size sized blocks.
        wanted = self.header_size if self.header_size > 0 else self.block_size
        while True:
            started_chunking = time.monotonic()
            # The file data is read directly into a fresh per-chunk buffer (zero-filled
            # for ranges stemming from holes), so each byte is copied only once, instead
            # of being assembled in an intermediate buffer and copied out of it again.
            buf = bytearray(wanted)
            got = self.reader.readinto(buf, wanted)
            self.chunking_time += time.monotonic() - started_chunking
            if got == 0:
                break  # EOF
            assert got <= wanted
            data = memoryview(buf)[:got]
            if zeros.startswith(data):
                # All-zero chunk: either read as zeros from the file or stemming from
                # a hole in a sparse file - we can not distinguish that here.
                data.release()
                yield Chunk(None, size=got, allocation=CH_ALLOC)
            else:
                yield Chunk(data, size=got, allocation=CH_DATA)
            wanted = self.block_size
