# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.


import zipfile

from apfio.base.types import check_type
from apfio.base.zip import ZippedArray
from apfio.imaging.avif import AVIFArray
from apfio.imaging.pil import PILArray
from apfio.vision.format import ImageFormat


def vopen(file, mode='r', shape=None, fmt=None, max_readers=1, zipargs=None, array_type=None, **eparams):
    r"""
    Opens a \*.zip file for reading and/or writing compressed image files from it.

    :param file: A path or file-like object representing the \*.zip file to open, as accepted by :py:class:`~python:zipfile.ZipFile`.

    :param mode: The mode in which the file is to be opened, see :py:class:`~python:zipfile.ZipFile`.

    :param shape: The shape of the array to open.

    :param fmt: An :py:class:`~apfio.vision.format.ImageFormat` supported by :py:mod:`apfio.vision`.

    :param max_readers: The maximum number of zip file readers to use for concurrent reading from the zip file.

    :param zipargs: Either `None`, or a :py:class:`~python:dict` of additional keyword arguments to be passed on to the :py:class:`~python:zipfile.ZipFile`
                    constructor.

    :param array_type: Either `None`, or a type that is a subclass of :py:class:`~apfio.imaging.array.CompressedImageArray`. If given, this procedure
                       will use this type to convert images to bitstrings and vice versa. If omitted, this procedure
                       will automatically select a :py:class:`~apfio.imaging.array.CompressedImageArray` subclass.

    :param eparams: Keyword arguments to be passed on to the image encoder (if images are to be written). For some
                    formats this will be :py:meth:`PIL.Image.save`, for others it might be FFMPEG command line arguments.

    :return: An :py:class:`~apfio.base.array.IndependentPointArray` object the data points of which are three-dimensional numpy arrays. The object
             can safely be accessed concurrently and it is a context manager.
    """

    if mode == 'a':
        raise ValueError("File mode 'a' is not supported, as this mode would require modifying existing entries in"
                         "the zip file, which cannot be done efficiently.")

    if mode not in 'rwx':
        raise ValueError("'mode' can only be 'r', 'w', or 'x'!")

    fmt = check_type(fmt, ImageFormat, allow_none=True)

    if zipargs is None:
        zipargs = dict()

    zf = None
    zarray = None
    iarray = None

    try:
        zf = zipfile.ZipFile(file, mode=mode, **zipargs)
        zarray = ZippedArray(zf, read=mode == 'r', write=mode in 'wx', own=True, max_open=max_readers)
        if fmt is None:
            fmt = ImageFormat.from_ext(zarray.extension)
        else:
            if zarray.write:
                zarray.extension = fmt.extension
            else:
                if zarray.extension != fmt.extension:
                    raise ValueError(f"The 'fmt' parameter for 'vopen' was set to {fmt.extension}, but the zip file "
                                     f"actually uses '{zarray.extension}' as the filename extension for its entries!")

        if shape is None:
            shape = zarray.shape
        else:
            if zarray.write:
                zarray.shape = shape
            else:
                if zarray.shape != shape:
                    raise ValueError(f"The 'shape' parameter for 'vopen' was set to {shape}, but the zip file "
                                     f"actually indicates '{zarray.shape}'!")

        if array_type is None:
            if fmt == ImageFormat.AVIF:
                iarray = AVIFArray(zarray, own=True)
            else:
                iarray = PILArray(zarray, own=True)
        else:
            iarray = array_type(zarray, own=True)
        assert iarray.shape == shape
        if iarray.writable:
            iarray.format = fmt.identifier
        else:
            if iarray.format != fmt.identifier:
                raise IOError(f"The *.zip file indicates '{zarray.extension}' as the filename extension for its entries, "
                              f"but at the same time it indicates '{iarray.format}' as the image format, which does "
                              f"not seem to fit together!")
        iarray.encoder_config = eparams

        return iarray

    except:
        if iarray is not None:
            iarray.close()
        if zarray is not None:
            zarray.close()
        if zf is not None:
            zf.close()
        raise
