"""
Contains some quick format checks. Should not try to import large dependencies.
If the dependency is necessary, the check function should be moved into the respective MountSource implementation.
It basically is something like libmagic (Linux 'file' command), but only for the types we actually need.
There seem to be no nice, simple, widespread, and maintained Python package for this.

https://github.com/h2non/filetype.py                    MIT  ~700 stars    Last Release: 2022-11-02
    Small, dependency-free, fast Python package to infer binary file types checking the magic numbers signature
https://github.com/riad-azz/py-file-type                MIT  1 star        Last Release: 2023-04-22
     Wrapper for python-magic that includes necessary bin files.
https://github.com/enchant97/python-file-type-guesser   MIT  1 star        Last Release: 2021-12-20
    Guess a file's type from its content and extension.
    This is intended to be a pure Python implementation of something like python-magic.
https://github.com/dveselov/python-libmagic             MIT  ~30 stars     Last Release: 2016-07-07
    Python bindings to libmagic

The situation for file type detection Python packages is abysmal.
And I already skipped as many other low-effort packages on PyPI.

See:
 - https://en.wikipedia.org/wiki/List_of_file_signatures
 - https://en.wikipedia.org/wiki/List_of_archive_formats
"""

import contextlib
import dataclasses
import enum
import re
import struct
import sys
import tarfile
import zipfile
from typing import IO, Callable, Optional, Union

with contextlib.suppress(ImportError):
    import rapidgzip

try:
    import sqlcipher3  # noqa: F401
except ImportError:
    sqlcipher3 = None  # type: ignore


class FileFormatID(enum.Enum):
    # fmt: off
    # Archive Formats (bundles more than one file)
    # "Normal" Archive formats with compression
    SEVEN_ZIP        = 0x101
    RAR              = 0x102
    ZIP              = 0x103
    SQLAR            = 0x104
    CAB              = 0x105
    AIX_SMALL        = 0x106
    AR               = 0x107
    XAR              = 0x108
    CPIO             = 0x109
    AR_THIN          = 0x110
    # "TAR"-like compression formats without compression
    TAR              = 0x201
    ASAR             = 0x202
    # Read-Only File systems
    SQUASHFS         = 0x301
    ISO9660          = 0x302
    # Fully-fledged file systems
    FAT              = 0x401
    EXT4             = 0x402
    # Other archive  formats
    RATARMOUNT_INDEX = 0x901
    WARC             = 0x902
    HTML             = 0x903
    PDF              = 0x904

    # Compression formats (compresses a single file / stream)
    BZIP2            = 0x1001
    GZIP             = 0x1002
    XZ               = 0x1003
    ZSTANDARD        = 0x1004
    ZLIB             = 0x1005
    LZ4              = 0x1006
    LZMA             = 0x1007
    DEFLATE          = 0x1008

    GRZIP            = 0x1021
    LRZIP            = 0x1022
    LZIP             = 0x1023
    LZOP             = 0x1024
    RPM              = 0x1025
    UU               = 0x1026
    Z                = 0x1027

    # Multimedia container formats
    AVI              = 0x1101
    MKV              = 0x1102
    MP4              = 0x1103
    RMVB             = 0x1104
    OGG              = 0x1105
    # fmt: on


FID = FileFormatID


HTML_REGEX = re.compile(br"<(head|title|html|script|style|table|a\s+href)[\s>]")


def is_html_file(fileobj: IO[bytes]) -> bool:
    offset = fileobj.tell()
    try:
        # Do not read more than 512 B because else we might match the first TAR entry contents!
        chunk = fileobj.read(512).lower()
        if b'\x00' in chunk:
            return False
        # Trying to decode the chunk can be inconclusive because we may not know the encoding and
        # because the chunk end might be inside a multi-byte UTF-8 character.

        # See https://github.com/file/file/blob/FILE5_46/magic/Magdir/sgml
        # This is a very trimmed-down heuristic. In the future, we might have to add a proper
        # file detection dependency such as libmagic (for Python).
        return (b'<!doctype html' in chunk) or bool(HTML_REGEX.search(chunk))
    finally:
        fileobj.seek(offset)
    return False


def is_tar(fileobj: IO[bytes], encoding: str = tarfile.ENCODING) -> bool:
    try:
        # Will only open uncompressed TAR files.
        with tarfile.open(fileobj=fileobj, mode='r:', encoding=encoding):
            return True
    except (tarfile.ReadError, tarfile.CompressionError):
        pass
    return False


def find_asar_header(fileobj: IO[bytes]) -> tuple[int, int, int]:
    """Return triple (start of header JSON, size of header JSON, start of file objects)"""
    # https://github.com/electron/asar/issues/128
    # https://github.com/electron/asar
    # > UInt32: header_size | String: header | Bytes: file1 | ...
    # > The header_size and header are serialized with Pickle class
    # > The header is a JSON string, and the header_size is the size of header's Pickle object.
    # Numbers are in little endian. Pickling writes the data, prepends a uint32 before and pads to 4 B offsets.
    # Because header_size is always 4 bytes, the implicit "magic" / redundant bytes of this file format are:
    # \x04\x00\x00\x00
    # Furthermore because pickling "header" already prepends the length of "header", the actual "header_size"
    # is completely redundant and can also be used as some kind of integrity check...
    # I don't understand where the fourth 32-bit number comes from. For some reason, the header seems to be
    # pickled twice, contrary to the ReadMe!?
    # https://github.com/electron/asar/issues/16
    # https://github.com/electron/asar/issues/226
    # https://knifecoat.com/Posts/ASAR+Format+Spec
    # -> Nice rant about the format as I have to concur about the usage of pickle-js.
    ASAR_MAGIC_SIZE = 4 * 4
    sizeOfPickledSize, sizeOfPickledPickledPickledHeader, sizeOfPickledPickledHeader, sizeOfPickledHeader = (
        struct.unpack('<LLLL', fileobj.read(ASAR_MAGIC_SIZE))
    )
    if sizeOfPickledSize != 4:
        raise ValueError("First magic bytes quadruplet does not match ASAR!")
    if sizeOfPickledPickledPickledHeader != sizeOfPickledPickledHeader + 4:
        raise ValueError("Second magic bytes quadruplet does not match ASAR!")
    padding = (4 - sizeOfPickledHeader % 4) % 4
    if sizeOfPickledPickledHeader != sizeOfPickledHeader + padding + 4:
        raise ValueError("Third magic bytes quadruplet does not match ASAR!")

    dataOffset = ASAR_MAGIC_SIZE + sizeOfPickledHeader + padding

    # It would be nice if we could check that dataOffset < file size, but we cannot get that (cheaply)
    # in case of compressed files or for large headers and because "seek(offset)" will always return that
    # offset even if it is > file size, so we would actually have to read the data.

    return ASAR_MAGIC_SIZE, sizeOfPickledHeader, dataOffset


def is_asar(fileobj: IO[bytes]) -> bool:
    offset = fileobj.tell()
    try:
        # Reading the header and checking it to be correct JSON is to expensive for large archives.
        # The unnecessary pickling already should introduce enough redundancy for checks and detection.
        find_asar_header(fileobj)
        return True
    except Exception:
        pass
    finally:
        fileobj.seek(offset)
    return False


def is_avi(fileobj: IO[bytes]) -> bool:
    offset = fileobj.tell()
    try:
        # AVI files start with 'RIFF' followed by file size and 'AVI '
        header = fileobj.read(12)
        if len(header) < 12:
            return False
        if header[:4] != b'RIFF':
            return False
        return header[8:12] == b'AVI '
    except Exception:
        pass
    finally:
        fileobj.seek(offset)
    return False


def is_squashfs(fileobj: IO[bytes]) -> bool:
    offset = fileobj.tell()
    try:
        # https://dr-emann.github.io/squashfs/squashfs.html#_the_superblock
        magicBytes = fileobj.read(4)
        if magicBytes != b"hsqs":
            return False

        _inodeCount, _modificationTime, blockSize, _fragmentCount = struct.unpack('<IIII', fileobj.read(4 * 4))
        compressor, blockSizeLog2, _flags, _idCount, major, minor = struct.unpack('<HHHHHH', fileobj.read(6 * 2))
        # root_inode, bytes_used, id_table, xattr_table, inode_table, dir_table, frag_table, export_table =
        # struct.unpack('<QQQQQQQQ', fileobj.read(8 * 8))

        # The size of a data block in bytes. Must be a power of two between 4096 (4k) and 1048576 (1 MiB).
        # log2 4096 = 12, log2 1024*1024 = 20
        if blockSizeLog2 < 12 or blockSizeLog2 > 20 or 2**blockSizeLog2 != blockSize:
            return False

        if major != 4 or minor != 0:
            return False

        # Compressions: 0:None, 1:GZIP, 2:LZMA, 3:LZO, 4:XZ, 5:LZ4, 6:ZSTD
        if compressor > 6:
            return False

    except struct.error:
        return False

    finally:
        fileobj.seek(offset)

    return True


def find_squashfs_offset(fileobj: IO[bytes], maxSkip=1024 * 1024) -> int:
    """
    Looks for the SquashFS superblock, which can be at something other than offset 0 for AppImage files.
    """
    # https://dr-emann.github.io/squashfs/squashfs.html#_the_superblock
    if is_squashfs(fileobj):
        return 0

    oldOffset = fileobj.tell()
    try:
        magic = b"hsqs"
        data = fileobj.read(maxSkip + len(magic))
        magicOffset = 0
        while True:
            magicOffset = data.find(magic, magicOffset + 1)
            if magicOffset < 0 or magicOffset >= len(data):
                break
            fileobj.seek(magicOffset)
            if is_squashfs(fileobj):
                return magicOffset
    finally:
        fileobj.seek(oldOffset)

    return -1


def _is_cpio(fileobj: IO[bytes]) -> bool:
    # http://justsolve.archiveteam.org/wiki/Cpio
    # https://github.com/libarchive/libarchive/blob/6110e9c82d8ba830c3440f36b990483ceaaea52c/libarchive/
    #   archive_read_support_format_cpio.c#L272
    firstBytes = fileobj.read(5)
    return firstBytes == b'07070' or firstBytes[:2] in [b'\x71\xc7', b'\xc7\x71']


def _is_iso9660(fileobj: IO[bytes]) -> bool:
    # https://www.iso.org/obp/ui/#iso:std:iso:9660:ed-1:v1:en
    # http://www.brankin.com/main/technotes/Notes_ISO9660.htm
    # https://en.wikipedia.org/wiki/ISO_9660
    # https://en.wikipedia.org/wiki/Optical_disc_image
    offset = 32 * 1024 + 1
    udfOffset = 38 * 1024 + 1
    buffer = fileobj.read(max(offset, udfOffset) + 4)
    return buffer[offset : offset + 5] == b'CD001' or buffer[udfOffset : udfOffset + 4] == b'NSR0'


def _check_zlib_header(fileobj: IO[bytes]) -> bool:
    header = fileobj.read(2)
    if len(header) < 2:
        return False
    cmf = header[0]
    if cmf & 0xF != 8:
        return False
    if cmf >> 4 > 7:
        return False
    flags = header[1]
    if ((cmf << 8) + flags) % 31 != 0:
        return False
    usesDictionary = ((flags >> 5) & 1) != 0
    return not usesDictionary


def _check_deflate(fileobj: IO[bytes]) -> bool:
    if 'rapidgzip' not in sys.modules:
        return False

    determineFileType = getattr(rapidgzip, 'determineFileType', None)
    if not callable(determineFileType):
        return False

    if determineFileType(fileobj).lower() != 'deflate':
        return False

    try:
        # determineFileType only checks the header, which is almost nothing for 'deflate'.
        # Therefore try to read a bit.
        with rapidgzip.RapidgzipFile(fileobj, parallelization=1, chunk_size=128 * 1024) as file:
            file.read(108 * 1024)  # Read a bit less than the chunk size.
        return True
    except Exception:
        pass

    return False


def _is_bzip2(fileobj: IO[bytes]) -> bool:
    return fileobj.read(4)[:3] == b'BZh' and fileobj.read(6) == (0x314159265359).to_bytes(6, 'big')


def _check_lz4_header(fileobj: IO[bytes]) -> bool:
    SKIPPABLE_FRAME_MAGIC = 0x184D2A50
    try:
        (magic,) = struct.unpack('<L', fileobj.read(4))

        # https://github.com/lz4/lz4/blob/dev/doc/lz4_Frame_format.md#skippable-frames
        while magic & 0xFFFF_FFF0 == SKIPPABLE_FRAME_MAGIC:
            (frame_size,) = struct.unpack('<L', fileobj.read(4))
            fileobj.seek(fileobj.tell() + frame_size)
            (magic,) = struct.unpack('<L', fileobj.read(4))

        # https://github.com/lz4/lz4/blob/dev/doc/lz4_Frame_format.md#general-structure-of-lz4-frame-format
        return magic == 0x184D2204
    except struct.error:
        return False


def _check_zstandard_header(fileobj: IO[bytes]) -> bool:
    SKIPPABLE_FRAME_MAGIC = 0x184D2A50
    try:
        (magic,) = struct.unpack('<L', fileobj.read(4))

        # https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#skippable-frames
        while magic & 0xFFFF_FFF0 == SKIPPABLE_FRAME_MAGIC:
            (frame_size,) = struct.unpack('<L', fileobj.read(4))
            fileobj.seek(fileobj.tell() + frame_size)
            (magic,) = struct.unpack('<L', fileobj.read(4))

        return magic == 0xFD2FB528
    except struct.error:
        return False


def _check_encrypted_sqlar(fileobj: IO[bytes]) -> bool:
    # File must at least contain the AES encryption salt, probably even many more bytes.
    salt = fileobj.read(16)
    return len(salt) >= 16


@dataclasses.dataclass
class FileFormatInfo:
    # Extensions without the initial '.'
    extensions: list[str]
    # If the first bytes of a format are constant, then they should be stated here.
    # This is used to build a lookup table for faster format detection as a first stage to thin down
    # the formats to actually call checkHeader for.
    # Some formats like ZIP (only has magic footer) and ISO9660 don't have a magic footer!
    magicBytes: Optional[bytes]
    # File format checkers should err on the side of false positives because else the file
    # will be rejected without actually trying to open it!
    # For simplicity, assume that all of these checkHeader will read from the current offset and be
    # at an arbitrary offset after the check. Normally, the offset should be larger than before in order
    # to mostly work with non-seekable files (although how is the caller supposed to parse the format after
    # an opaque check that skips an arbitrary amount of header?).
    checkHeader: Optional[Callable[[IO[bytes]], bool]] = None


COMPRESSION_FORMATS: dict[FileFormatID, FileFormatInfo] = {
    FID.BZIP2: FileFormatInfo(['bz2', 'bzip2'], b'BZh', _is_bzip2),
    FID.GZIP: FileFormatInfo(['gz', 'gzip'], b'\x1f\x8b'),
    FID.XZ: FileFormatInfo(['xz'], b"\xfd7zXZ\x00"),
    FID.ZSTANDARD: FileFormatInfo(['zst', 'zstd', 'pzstd'], None, _check_zstandard_header),
    FID.ZLIB: FileFormatInfo(['zz', 'zlib'], None, _check_zlib_header),
    FID.DEFLATE: FileFormatInfo(['deflate'], None, _check_deflate),
    # https://github.com/libarchive/libarchive/blob/6110e9c82d8ba830c3440f36b990483ceaaea52c/libarchive/
    # archive_read_support_filter_grzip.c#L46
    # It is almost impossible to find the original sources or a specification:
    # https://t2sde.org/packages/grzip
    # ftp://ftp.ac-grenoble.fr/ge/compression/grzip-0.3.0.tar.bz2
    FID.GRZIP: FileFormatInfo(['grz', 'grzip'], b'GRZipII\x00\x02\x04:)'),
    # https://github.com/ckolivas/lrzip/blob/master/doc/magic.header.txt
    FID.LRZIP: FileFormatInfo(['lrz', 'lrzip'], b'LRZI\x00'),
    # https://github.com/lz4/lz4/blob/dev/doc/lz4_Frame_format.md#general-structure-of-lz4-frame-format
    FID.LZ4: FileFormatInfo(['lz4'], None, _check_lz4_header),
    # https://www.ietf.org/archive/id/draft-diaz-lzip-09.txt
    FID.LZIP: FileFormatInfo(['lz', 'lzip'], b'LZIP\x01'),
    # https://github.com/jljusten/LZMA-SDK/blob/master/DOC/lzma-specification.txt
    # https://github.com/frizb/FirmwareReverseEngineering/blob/master/IdentifyingCompressionAlgorithms.md#lzma
    FID.LZMA: FileFormatInfo(['lzma'], b'\x5d\x00\x00'),
    FID.LZOP: FileFormatInfo(['lzo', 'lzop'], b'\x89\x4c\x5a\x4f\x00\x0d\x0a\x1a\x0a'),
    # https://refspecs.linuxbase.org/LSB_4.1.0/LSB-Core-generic/LSB-Core-generic/pkgformat.html
    FID.RPM: FileFormatInfo(['rpm'], b'\xed\xab\xee\xdb'),
    # https://en.wikipedia.org/wiki/Uuencoding
    FID.UU: FileFormatInfo(['uu'], b'begin '),
    # https://github.com/file/file/blob/master/magic/Magdir/compress
    FID.Z: FileFormatInfo(['Z'], b'\x1f\x9d'),
}


ARCHIVE_FORMATS: dict[FileFormatID, FileFormatInfo] = {
    # "Normal" Archive formats with compression
    FID.SEVEN_ZIP: FileFormatInfo(['7z', '7zip', 'cb7'], b'7z\xbc\xaf\x27\x1c'),
    FID.RAR: FileFormatInfo(['rar', 'cbr'], b'Rar!\x1a\x07'),
    # is_zipfile might yields some false positives, we want it to err on the positive side.
    # See: https://bugs.python.org/issue42096
    # https://en.wikipedia.org/wiki/Comic_book_archive
    FID.ZIP: FileFormatInfo(['zip', 'cbz'], None, zipfile.is_zipfile),
    # SQLAR can be encrypted, for which it will not have any magic bytes! Instead, the first 16 B are the salt.
    FID.SQLAR: (
        FileFormatInfo(['sqlar'], b'SQLite format 3\x00', None)
        if sqlcipher3 is None
        else FileFormatInfo(['sqlar'], None, _check_encrypted_sqlar)
    ),
    # https://download.microsoft.com/download/4/d/a/4da14f27-b4ef-4170-a6e6-5b1ef85b1baa/[ms-cab].pdf
    FID.CAB: FileFormatInfo(['cab'], b'MSCF'),
    # https://www.ibm.com/docs/en/aix/7.2.0?topic=formats-ar-file-format-small
    FID.AIX_SMALL: FileFormatInfo(['ar'], b'<aiaff>\n'),
    # https://www.ibm.com/docs/en/aix/7.2.0?topic=formats-ar-file-format-big#ar_big
    # https://en.wikipedia.org/wiki/Ar_(Unix)
    FID.AR: FileFormatInfo(['a', 'ar', 'lib', 'deb'], b'!<arch>\n'),
    FID.AR_THIN: FileFormatInfo(['a', 'ar', 'lib'], b'!<thin>\n'),
    FID.XAR: FileFormatInfo(['xar'], b'xar!'),
    FID.CPIO: FileFormatInfo(['cpio'], None, _is_cpio),
    # "TAR"-like compression formats without compression
    FID.TAR: FileFormatInfo(
        ['tar', 'tb2', 'tbz', 'tbz2', 'tz2', 'taz', 'tgz', 'txz', 'tzst', 'cbt']
        + ['t' + extension for format_info in COMPRESSION_FORMATS.values() for extension in format_info.extensions]
        + ['tar.' + extension for format_info in COMPRESSION_FORMATS.values() for extension in format_info.extensions],
        None,
        is_tar,
    ),
    FID.ASAR: FileFormatInfo(['asar'], b'\x04\x00\x00\x00', is_asar),
    FID.SQUASHFS: FileFormatInfo(['squashfs', 'AppImage', 'snap'], None, lambda x: find_squashfs_offset(x) >= 0),
}


DISK_FORMATS: dict[FileFormatID, FileFormatInfo] = {
    FID.ISO9660: FileFormatInfo(['iso'], None, _is_iso9660),
    # Fully-fledged file systems
    FID.FAT: FileFormatInfo(['fat', 'img', 'dd', 'fat12', 'fat16', 'fat32', 'raw'], None),
    FID.EXT4: FileFormatInfo(['ext4', 'img', 'dd', 'raw'], None),
}


DOCUMENT_FORMATS: dict[FileFormatID, FileFormatInfo] = {
    # https://www.iso.org/standard/68004.html
    # https://iipc.github.io/warc-specifications/specifications/warc-format/warc-1.1/#file-and-record-model
    FID.WARC: FileFormatInfo(['warc'], b'WARC/1.'),
    # HTML files with embedded data URLs
    FID.HTML: FileFormatInfo(['html', 'htm'], None, is_html_file),
    # PDF files with embedded files
    FID.PDF: FileFormatInfo(['pdf'], b'%PDF-', None),
    FID.RATARMOUNT_INDEX: FileFormatInfo(['index.sqlite'], b'SQLite format 3\x00'),
}


MULTIMEDIA_FORMATS: dict[FileFormatID, FileFormatInfo] = {
    FID.AVI: FileFormatInfo(['avi'], b'RIFF', is_avi),
    FID.MKV: FileFormatInfo(['mkv', 'mka', 'mks', 'mk3d', 'webm'], b'\x1a\x45\xdf\xa3', None),
    FID.MP4: FileFormatInfo(['mp4'], b'ftypMSNV', None),
    FID.RMVB: FileFormatInfo(['rmvb', 'rmb'], b'.RMF', None),
    FID.OGG: FileFormatInfo(['ogm', 'ogv'], b'OggS\x00', None),
}


def _get_extensions(formats: dict[FileFormatID, FileFormatInfo]) -> set[str]:
    return {extension for info in formats.values() for extension in info.extensions}


# Some formats, such as ZIP and AR, have further subcategories, sometimes only discerned by their extension.
# This makes categorizing file formats a bit more complicated.
# https://en.wikipedia.org/wiki/XPInstall
# https://learn.microsoft.com/de-de/visualstudio/extensibility/anatomy-of-a-vsix-package
# https://en.wikipedia.org/wiki/Apk_(file_format)
_BINARY_EXTENSIONS = {
    'a',
    'lib',
    'deb',
    'AppImage',
    'jar',
    'nupkg',
    # Various extension packages for firefox, VSCode, Android, ...
    'xpi',
    'vsix',
    'apk',
    # Microsoft application formats
    'appx',
    'appxbundle',
    'appv',
}
_DOCUMENT_EXTENSIONS = {
    # https://en.wikipedia.org/wiki/Office_Open_XML_file_formats
    # https://support.microsoft.com/en-gb/office/
    #   file-formats-that-are-supported-in-powerpoint-252c6fa0-a4bc-41be-ac82-b77c9773f9dc
    # https://www.loc.gov/preservation/digital/formats/fdd/fdd000363.shtml
    # https://learn.microsoft.com/de-de/office/client-developer/visio/introduction-to-the-visio-file-formatvsdx
    # The suffix 'm' stands for "macro-enabled".
    # pptx - PowerPoint Presentation
    # potx - PowerPoint Design Templates
    # ppsx - PowerPoint Show
    # dotx - Word Document Templates
    # xltx - Excel Templates
    f'{suffix}{macro}'
    for suffix in ('doc', 'dot', 'ppt', 'pot', 'pps', 'vsd', 'vss', 'vst', 'xls', 'xlt')
    for macro in ('x', 'm')
} | {
    # https://en.wikipedia.org/wiki/OpenDocument
    'odt',
    'ods',
    'odg',
    'epub',
    # https://en.wikipedia.org/wiki/3D_Manufacturing_Format
    '3mf',
    # https://en.wikipedia.org/wiki/Design_Web_Format
    'dwf',
    'dwfx',
    # https://www.circuit-diagram.org/docs/cddx
    'cddx',
    # https://learn.microsoft.com/en-us/azure/cloud-services-extended-support/cloud-services-model-and-package
    'cspkg',
    # https://en.wikipedia.org/wiki/Open_Packaging_Conventions
    'familyx',
    'fdix',
    'appv',
    'semblio',
    'xps',
    'oxps',
    'mmzx',
    'aasx',
    'slx',
    'smpk',
    'scdoc',
    'jtx',
    'amlx',
    'pbix',
    'pbit',
    # https://www.loc.gov/preservation/digital/formats/fdd/fdd000547.shtml
    'kmlz',
    # https://vrarwiki.com/wiki/USDZ
    'usdz',
    # https://www.loc.gov/preservation/digital/formats/fdd/fdd000586.shtml
    'wacz',
}
EXTENSIONS_BY_CATEGORY: dict[str, set[str]] = {
    'archive': _get_extensions(ARCHIVE_FORMATS) - _BINARY_EXTENSIONS,
    'compressed': _get_extensions(COMPRESSION_FORMATS),
    'disk': _get_extensions(DISK_FORMATS),
    # https://7-zip.opensource.jp/chm/general/formats.htm
    'document': _get_extensions(DOCUMENT_FORMATS) | _DOCUMENT_EXTENSIONS,
    'binary': _BINARY_EXTENSIONS,
    'multimedia': _get_extensions(MULTIMEDIA_FORMATS),
}
EXTENSIONS_BY_CATEGORY['all'] = set().union(*EXTENSIONS_BY_CATEGORY.values())


FILE_FORMATS: dict[FileFormatID, FileFormatInfo] = {
    fid: info
    for formats in (ARCHIVE_FORMATS, COMPRESSION_FORMATS, DISK_FORMATS, DOCUMENT_FORMATS, MULTIMEDIA_FORMATS)
    for fid, info in formats.items()
}


# Check that all defined format IDs have FileFormatInfo, so that we can assume FORMATS[FID] to not fail.
for _formatInfo in FileFormatID:
    assert _formatInfo in FILE_FORMATS, f"Missing file format information for: {_formatInfo}"


# Maps the first 2 bytes to possible candidates. 2 B was chosen because most formats have
# at least 2 B of constant magic and if the table was implemented in contiguous memory, it would be nicely
# addressable by a 16-bit index (the 2 B). Statistically, this cuts down the number of expensive checkHeader
# calls by 65536.
_MAGIC_BYTES_TO_FORMATS: dict[bytes, list[FileFormatID]] = {}
_FORMATS_WITHOUT_MAGIC_BYTES: list[FileFormatID] = []


def recompute_cached_magic_bytes():
    """Should be called when injecting new file formats from outside."""
    for fid, info in FILE_FORMATS.items():
        if info.magicBytes is None or len(info.magicBytes) < 2:
            _FORMATS_WITHOUT_MAGIC_BYTES.append(fid)
        else:
            firstTwoBytes = info.magicBytes[:2]
            if firstTwoBytes not in _MAGIC_BYTES_TO_FORMATS:
                _MAGIC_BYTES_TO_FORMATS[firstTwoBytes] = []
            _MAGIC_BYTES_TO_FORMATS[firstTwoBytes].append(fid)


recompute_cached_magic_bytes()


def might_be_format(fileobj: IO[bytes], fid: Union[FileFormatID, FileFormatInfo]) -> bool:
    formatInfo = fid if isinstance(fid, FileFormatInfo) else FILE_FORMATS[fid]
    oldOffset = fileobj.tell()
    try:
        if formatInfo.magicBytes and fileobj.read(len(formatInfo.magicBytes)) != formatInfo.magicBytes:
            return False
        if formatInfo.checkHeader:
            fileobj.seek(oldOffset)
            return formatInfo.checkHeader(fileobj)
    finally:
        fileobj.seek(oldOffset)

    # For no magic bytes and no header check, reject it because the necessary module to read it seems to be missing.
    return bool(formatInfo.magicBytes or formatInfo.checkHeader)


def detect_formats(fileobj: IO[bytes]) -> set[FileFormatID]:
    oldOffset = fileobj.tell()
    try:
        firstTwoBytes = fileobj.read(2)
        # I don't think there is any file format that can be recognized if the file is smaller than 2 B.
        if len(firstTwoBytes) < 2:
            return set()

        formatsToTest = _MAGIC_BYTES_TO_FORMATS.get(firstTwoBytes, [])
    finally:
        fileobj.seek(oldOffset)

    return {fid for fid in formatsToTest + _FORMATS_WITHOUT_MAGIC_BYTES if might_be_format(fileobj, fid)}


def replace_format_check(fid: FileFormatID, checkHeader: Optional[Callable[[IO[bytes]], bool]] = None):
    for formats in [ARCHIVE_FORMATS, COMPRESSION_FORMATS, FILE_FORMATS]:
        if fid in formats:
            info = dataclasses.asdict(formats[fid])
            info['checkHeader'] = checkHeader
            formats[fid] = FileFormatInfo(**info)
