#!/usr/bin/env pytest
###############################################################################
# $Id$
#
# Project:  GDAL/OGR Test Suite
# Purpose:  Test Zarr driver
# Author:   Even Rouault <even.rouault@spatialys.com>
#
###############################################################################
# Copyright (c) 2021, Even Rouault <even.rouault@spatialys.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
###############################################################################

import array
import base64
import json
import math
import struct
import sys

import gdaltest
import pytest

from osgeo import gdal, osr


@pytest.mark.parametrize(
    "dtype,structtype,gdaltype,fill_value,nodata_value",
    [
        ["!b1", "B", gdal.GDT_Byte, None, None],
        ["!i1", "b", gdal.GDT_Int16, None, None],
        ["!i1", "b", gdal.GDT_Int16, -1, -1],
        ["!u1", "B", gdal.GDT_Byte, None, None],
        [
            "!u1",
            "B",
            gdal.GDT_Byte,
            "1",
            1,
        ],  # not really legit to have the fill_value as a str
        ["<i2", "h", gdal.GDT_Int16, None, None],
        [">i2", "h", gdal.GDT_Int16, None, None],
        ["<i4", "i", gdal.GDT_Int32, None, None],
        [">i4", "i", gdal.GDT_Int32, None, None],
        ["<i8", "q", gdal.GDT_Int64, None, None],
        ["<i8", "q", gdal.GDT_Int64, -(1 << 63), -(1 << 63)],
        [
            "<i8",
            "q",
            gdal.GDT_Int64,
            str(-(1 << 63)),
            -(1 << 63),
        ],  # not really legit to have the fill_value as a str
        [">i8", "q", gdal.GDT_Int64, None, None],
        ["<u2", "H", gdal.GDT_UInt16, None, None],
        [">u2", "H", gdal.GDT_UInt16, None, None],
        ["<u4", "I", gdal.GDT_UInt32, None, None],
        [">u4", "I", gdal.GDT_UInt32, None, None],
        ["<u4", "I", gdal.GDT_UInt32, 4000000000, 4000000000],
        [
            "<u8",
            "Q",
            gdal.GDT_UInt64,
            str((1 << 64) - 1),
            (1 << 64) - 1,
        ],  # not really legit to have the fill_value as a str, but libjson-c can't support numeric values in int64::max(), uint64::max() range.
        [">u8", "Q", gdal.GDT_UInt64, None, None],
        ["<f4", "f", gdal.GDT_Float32, None, None],
        [">f4", "f", gdal.GDT_Float32, None, None],
        ["<f4", "f", gdal.GDT_Float32, 1.5, 1.5],
        ["<f4", "f", gdal.GDT_Float32, "NaN", float("nan")],
        ["<f4", "f", gdal.GDT_Float32, "Infinity", float("infinity")],
        ["<f4", "f", gdal.GDT_Float32, "-Infinity", float("-infinity")],
        ["<f8", "d", gdal.GDT_Float64, None, None],
        [">f8", "d", gdal.GDT_Float64, None, None],
        ["<f8", "d", gdal.GDT_Float64, "NaN", float("nan")],
        ["<f8", "d", gdal.GDT_Float64, "Infinity", float("infinity")],
        ["<f8", "d", gdal.GDT_Float64, "-Infinity", float("-infinity")],
        ["<c8", "f", gdal.GDT_CFloat32, None, None],
        [">c8", "f", gdal.GDT_CFloat32, None, None],
        ["<c16", "d", gdal.GDT_CFloat64, None, None],
        [">c16", "d", gdal.GDT_CFloat64, None, None],
    ],
)
@pytest.mark.parametrize("use_optimized_code_paths", [True, False])
def test_zarr_basic(
    dtype, structtype, gdaltype, fill_value, nodata_value, use_optimized_code_paths
):

    j = {
        "chunks": [2, 3],
        "compressor": None,
        "dtype": dtype,
        "fill_value": fill_value,
        "filters": None,
        "order": "C",
        "shape": [5, 4],
        "zarr_format": 2,
    }

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))
        if gdaltype not in (gdal.GDT_CFloat32, gdal.GDT_CFloat64):
            tile_0_0_data = struct.pack(dtype[0] + (structtype * 6), 1, 2, 3, 5, 6, 7)
            tile_0_1_data = struct.pack(dtype[0] + (structtype * 6), 4, 0, 0, 8, 0, 0)
        else:
            tile_0_0_data = struct.pack(
                dtype[0] + (structtype * 12), 1, 11, 2, 0, 3, 0, 5, 0, 6, 0, 7, 0
            )
            tile_0_1_data = struct.pack(
                dtype[0] + (structtype * 12), 4, 0, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0
            )
        gdal.FileFromMemBuffer("/vsimem/test.zarr/0.0", tile_0_0_data)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/0.1", tile_0_1_data)
        with gdaltest.config_option(
            "GDAL_ZARR_USE_OPTIMIZED_CODE_PATHS",
            "YES" if use_optimized_code_paths else "NO",
        ):
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
        assert ar
        assert ar.GetDimensionCount() == 2
        assert [ar.GetDimensions()[i].GetSize() for i in range(2)] == [5, 4]
        assert ar.GetBlockSize() == [2, 3]
        if nodata_value is not None and math.isnan(nodata_value):
            assert math.isnan(ar.GetNoDataValue())
        else:
            assert ar.GetNoDataValue() == nodata_value

        assert ar.GetOffset() is None
        assert ar.GetScale() is None
        assert ar.GetUnit() == ""

        # Check reading one single value
        assert ar[1, 2].Read(
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64)
        ) == struct.pack("d" * 1, 7)

        if structtype == "b":
            structtype_read = "h"
        else:
            structtype_read = structtype

        # Read block 0,0
        if gdaltype not in (gdal.GDT_CFloat32, gdal.GDT_CFloat64):
            assert ar[0:2, 0:3].Read(
                buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64)
            ) == struct.pack("d" * 6, 1, 2, 3, 5, 6, 7)
            assert struct.unpack(structtype_read * 6, ar[0:2, 0:3].Read()) == (
                1,
                2,
                3,
                5,
                6,
                7,
            )
        else:
            assert ar[0:2, 0:3].Read(
                buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_CFloat64)
            ) == struct.pack("d" * 12, 1, 11, 2, 0, 3, 0, 5, 0, 6, 0, 7, 0)
            assert struct.unpack(structtype * 12, ar[0:2, 0:3].Read()) == (
                1,
                11,
                2,
                0,
                3,
                0,
                5,
                0,
                6,
                0,
                7,
                0,
            )

        # Read block 0,1
        assert ar[0:2, 3:4].Read(
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64)
        ) == struct.pack("d" * 2, 4, 8)

        # Read block 1,1 (missing)
        nv = nodata_value if nodata_value else 0
        assert ar[2:4, 3:4].Read(
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64)
        ) == struct.pack("d" * 2, nv, nv)

        # Read whole raster
        assert ar.Read(
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64)
        ) == struct.pack(
            "d" * 20,
            1,
            2,
            3,
            4,
            5,
            6,
            7,
            8,
            nv,
            nv,
            nv,
            nv,
            nv,
            nv,
            nv,
            nv,
            nv,
            nv,
            nv,
            nv,
        )

        if gdaltype not in (gdal.GDT_CFloat32, gdal.GDT_CFloat64):
            assert ar.Read() == array.array(
                structtype_read,
                [
                    1,
                    2,
                    3,
                    4,
                    5,
                    6,
                    7,
                    8,
                    nv,
                    nv,
                    nv,
                    nv,
                    nv,
                    nv,
                    nv,
                    nv,
                    nv,
                    nv,
                    nv,
                    nv,
                ],
            )
        else:
            assert ar.Read() == array.array(
                structtype,
                [
                    1,
                    11,
                    2,
                    0,
                    3,
                    0,
                    4,
                    0,
                    5,
                    0,
                    6,
                    0,
                    7,
                    0,
                    8,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                    nv,
                    0,
                ],
            )
        # Read with negative steps
        assert ar.Read(
            array_start_idx=[2, 1],
            count=[2, 2],
            array_step=[-1, -1],
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64),
        ) == struct.pack("d" * 4, nv, nv, 6, 5)

        # array_step > 2
        assert ar.Read(
            array_start_idx=[0, 0],
            count=[1, 2],
            array_step=[0, 2],
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64),
        ) == struct.pack("d" * 2, 1, 3)

        assert ar.Read(
            array_start_idx=[0, 0],
            count=[3, 1],
            array_step=[2, 0],
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64),
        ) == struct.pack("d" * 3, 1, nv, nv)

        assert ar.Read(
            array_start_idx=[0, 1],
            count=[1, 2],
            array_step=[0, 2],
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64),
        ) == struct.pack("d" * 2, 2, 4)

        assert ar.Read(
            array_start_idx=[0, 0],
            count=[1, 2],
            array_step=[0, 3],
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64),
        ) == struct.pack("d" * 2, 1, 4)

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize(
    "fill_value,expected_read_data",
    [[base64.b64encode(b"xyz").decode("utf-8"), ["abc", "xyz"]], [None, ["abc", None]]],
)
def test_zarr_string(fill_value, expected_read_data):

    j = {
        "chunks": [1],
        "compressor": None,
        "dtype": "|S3",
        "fill_value": fill_value,
        "filters": [],
        "order": "C",
        "shape": [2],
        "zarr_format": 2,
    }

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))
        gdal.FileFromMemBuffer("/vsimem/test.zarr/0", b"abc")
        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        assert ds
        rg = ds.GetRootGroup()
        assert rg
        ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
        assert ar
        assert ar.Read() == expected_read_data

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


# Check that all required elements are present in .zarray
@pytest.mark.parametrize(
    "member",
    [
        None,
        "zarr_format",
        "chunks",
        "compressor",
        "dtype",
        "filters",
        "order",
        "shape",
        "fill_value",
    ],
)
def test_zarr_invalid_json_remove_member(member):

    j = {
        "chunks": [2, 3],
        "compressor": None,
        "dtype": "!b1",
        "fill_value": None,
        "filters": None,
        "order": "C",
        "shape": [5, 4],
        "zarr_format": 2,
    }

    if member:
        del j[member]

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))
        with gdaltest.error_handler():
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        if member == "fill_value":
            assert ds is not None
            assert gdal.GetLastErrorMsg() != ""
        elif member is None:
            assert ds
        else:
            assert ds is None
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


# Check bad values of members in .zarray
@pytest.mark.parametrize(
    "dict_update",
    [
        {"chunks": None},
        {"chunks": "invalid"},
        {"chunks": [2]},
        {"chunks": [2, 0]},
        {"shape": None},
        {"shape": "invalid"},
        {"shape": [5]},
        {"shape": [5, 0]},
        {"chunks": [1 << 40, 1 << 40], "shape": [1 << 40, 1 << 40]},
        {"shape": [1 << 30, 1 << 30, 1 << 30], "chunks": [1, 1, 1]},
        {"dtype": None},
        {"dtype": 1},
        {"dtype": ""},
        {"dtype": "!"},
        {"dtype": "!b"},
        {"dtype": "<u16"},
        {"dtype": "<u0"},
        {"dtype": "<u10000"},
        {"fill_value": []},
        {"fill_value": "x"},
        {"fill_value": "NaN"},
        {"dtype": "!S1", "fill_value": 0},
        {"order": None},
        {"order": "invalid"},
        {"compressor": "invalid"},
        {"compressor": {}},
        {"compressor": {"id": "invalid"}},
        {"filters": "invalid"},
        {"filters": {}},
        {"filters": [{"missing_id": True}]},
        {"zarr_format": None},
        {"zarr_format": 1},
    ],
)
def test_zarr_invalid_json_wrong_values(dict_update):

    j = {
        "chunks": [2, 3],
        "compressor": None,
        "dtype": "!b1",
        "fill_value": None,
        "filters": None,
        "order": "C",
        "shape": [5, 4],
        "zarr_format": 2,
    }

    j.update(dict_update)

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))
        with gdaltest.error_handler():
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        assert ds is None
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


# Check reading different compression methods
@pytest.mark.parametrize(
    "datasetname,compressor",
    [
        ("blosc.zarr", "blosc"),
        ("gzip.zarr", "gzip"),
        ("lz4.zarr", "lz4"),
        ("lzma.zarr", "lzma"),
        ("lzma_with_filters.zarr", "lzma"),
        ("zlib.zarr", "zlib"),
        ("zstd.zarr", "zstd"),
    ],
)
def test_zarr_read_compression_methods(datasetname, compressor):

    compressors = gdal.GetDriverByName("Zarr").GetMetadataItem("COMPRESSORS")
    filename = "data/zarr/" + datasetname

    if compressor not in compressors:
        with gdaltest.error_handler():
            ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
        assert ds is None
    else:
        ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
        rg = ds.GetRootGroup()
        assert rg
        ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
        assert ar
        assert ar.Read() == array.array("b", [1, 2])


@pytest.mark.parametrize("name", ["u1", "u2", "u4", "u8"])
def test_zarr_read_fortran_order(name):

    filename = "data/zarr/order_f_" + name + ".zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    rg = ds.GetRootGroup()
    assert rg
    ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
    assert ar
    assert ar.Read(
        buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Byte)
    ) == array.array("b", [i for i in range(16)])


def test_zarr_read_fortran_order_string():

    filename = "data/zarr/order_f_s3.zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    rg = ds.GetRootGroup()
    assert rg
    ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
    assert ar
    assert ar.Read() == [
        "000",
        "111",
        "222",
        "333",
        "444",
        "555",
        "666",
        "777",
        "888",
        "999",
        "AAA",
        "BBB",
        "CCC",
        "DDD",
        "EEE",
        "FFF",
    ]


def test_zarr_read_fortran_order_3d():

    filename = "data/zarr/order_f_u1_3d.zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    rg = ds.GetRootGroup()
    assert rg
    ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
    assert ar
    assert ar.Read(
        buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Byte)
    ) == array.array("b", [i for i in range(2 * 3 * 4)])


def test_zarr_read_compound_well_aligned():

    filename = "data/zarr/compound_well_aligned.zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    assert ds is not None

    rg = ds.GetRootGroup()
    assert rg
    ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
    assert ar

    dt = ar.GetDataType()
    assert dt.GetSize() == 4
    comps = dt.GetComponents()
    assert len(comps) == 2
    assert comps[0].GetName() == "a"
    assert comps[0].GetOffset() == 0
    assert comps[0].GetType().GetNumericDataType() == gdal.GDT_UInt16
    assert comps[1].GetName() == "b"
    assert comps[1].GetOffset() == 2
    assert comps[1].GetType().GetNumericDataType() == gdal.GDT_UInt16

    assert ar["a"].Read() == array.array("H", [1000, 4000, 0])
    assert ar["b"].Read() == array.array("H", [3000, 5000, 0])

    j = gdal.MultiDimInfo(ds, detailed=True)
    assert j["arrays"]["compound_well_aligned"]["values"] == [
        {"a": 1000, "b": 3000},
        {"a": 4000, "b": 5000},
        {"a": 0, "b": 0},
    ]


def test_zarr_read_compound_not_aligned():

    filename = "data/zarr/compound_not_aligned.zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    assert ds is not None

    rg = ds.GetRootGroup()
    assert rg
    ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
    assert ar

    dt = ar.GetDataType()
    assert dt.GetSize() == 6
    comps = dt.GetComponents()
    assert len(comps) == 3
    assert comps[0].GetName() == "a"
    assert comps[0].GetOffset() == 0
    assert comps[0].GetType().GetNumericDataType() == gdal.GDT_UInt16
    assert comps[1].GetName() == "b"
    assert comps[1].GetOffset() == 2
    assert comps[1].GetType().GetNumericDataType() == gdal.GDT_Byte
    assert comps[2].GetName() == "c"
    assert comps[2].GetOffset() == 4
    assert comps[2].GetType().GetNumericDataType() == gdal.GDT_UInt16

    assert ar["a"].Read() == array.array("H", [1000, 4000, 0])
    assert ar["b"].Read() == array.array("B", [2, 4, 0])
    assert ar["c"].Read() == array.array("H", [3000, 5000, 0])

    j = gdal.MultiDimInfo(ds, detailed=True)
    assert j["arrays"]["compound_not_aligned"]["values"] == [
        {"a": 1000, "b": 2, "c": 3000},
        {"a": 4000, "b": 4, "c": 5000},
        {"a": 0, "b": 0, "c": 0},
    ]


def test_zarr_read_compound_complex():

    filename = "data/zarr/compound_complex.zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    assert ds is not None

    rg = ds.GetRootGroup()
    assert rg
    ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
    assert ar

    is_64bit = sys.maxsize > 2**32

    dt = ar.GetDataType()
    assert dt.GetSize() == 24 if is_64bit else 16
    comps = dt.GetComponents()
    assert len(comps) == 4
    assert comps[0].GetName() == "a"
    assert comps[0].GetOffset() == 0
    assert comps[0].GetType().GetNumericDataType() == gdal.GDT_Byte
    assert comps[1].GetName() == "b"
    assert comps[1].GetOffset() == 2
    assert comps[1].GetType().GetClass() == gdal.GEDTC_COMPOUND
    assert comps[1].GetType().GetSize() == 1 + 1 + 2 + 1 + 1  # last one is padding

    subcomps = comps[1].GetType().GetComponents()
    assert len(subcomps) == 4

    assert comps[2].GetName() == "c"
    assert comps[2].GetOffset() == 8
    assert comps[2].GetType().GetClass() == gdal.GEDTC_STRING
    assert comps[3].GetName() == "d"
    assert comps[3].GetOffset() == 16 if is_64bit else 12
    assert comps[3].GetType().GetNumericDataType() == gdal.GDT_Int16

    j = gdal.MultiDimInfo(ds, detailed=True)
    assert j["arrays"]["compound_complex"]["values"] == [
        {"a": 1, "b": {"b1": 2, "b2": 3, "b3": 1000, "b5": 4}, "c": "AAA", "d": -1},
        {
            "a": 2,
            "b": {"b1": 255, "b2": 254, "b3": 65534, "b5": 253},
            "c": "ZZ",
            "d": -2,
        },
    ]


def test_zarr_read_array_attributes():

    filename = "data/zarr/array_attrs.zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    assert ds is not None

    j = gdal.MultiDimInfo(ds)
    assert j["arrays"]["array_attrs"]["attributes"] == {
        "bool": True,
        "double": 1.5,
        "doublearray": [1.5, 2.5],
        "int": 1,
        "int64": 1234567890123,
        "int64array": [1234567890123, -1234567890123],
        "intarray": [1, 2],
        "intdoublearray": [1, 2.5],
        "mixedstrintarray": ["foo", 1],
        "null": "",
        "obj": {},
        "str": "foo",
        "strarray": ["foo", "bar"],
    }


@pytest.mark.parametrize("crs_member", ["projjson", "wkt", "url"])
def test_zarr_read_crs(crs_member):

    zarray = {
        "chunks": [2, 3],
        "compressor": None,
        "dtype": "!b1",
        "fill_value": None,
        "filters": None,
        "order": "C",
        "shape": [5, 4],
        "zarr_format": 2,
    }

    zattrs_all = {
        "_CRS": {
            "projjson": {
                "$schema": "https://proj.org/schemas/v0.2/projjson.schema.json",
                "type": "GeographicCRS",
                "name": "WGS 84",
                "datum_ensemble": {
                    "name": "World Geodetic System 1984 ensemble",
                    "members": [
                        {
                            "name": "World Geodetic System 1984 (Transit)",
                            "id": {"authority": "EPSG", "code": 1166},
                        },
                        {
                            "name": "World Geodetic System 1984 (G730)",
                            "id": {"authority": "EPSG", "code": 1152},
                        },
                        {
                            "name": "World Geodetic System 1984 (G873)",
                            "id": {"authority": "EPSG", "code": 1153},
                        },
                        {
                            "name": "World Geodetic System 1984 (G1150)",
                            "id": {"authority": "EPSG", "code": 1154},
                        },
                        {
                            "name": "World Geodetic System 1984 (G1674)",
                            "id": {"authority": "EPSG", "code": 1155},
                        },
                        {
                            "name": "World Geodetic System 1984 (G1762)",
                            "id": {"authority": "EPSG", "code": 1156},
                        },
                    ],
                    "ellipsoid": {
                        "name": "WGS 84",
                        "semi_major_axis": 6378137,
                        "inverse_flattening": 298.257223563,
                    },
                    "accuracy": "2.0",
                    "id": {"authority": "EPSG", "code": 6326},
                },
                "coordinate_system": {
                    "subtype": "ellipsoidal",
                    "axis": [
                        {
                            "name": "Geodetic latitude",
                            "abbreviation": "Lat",
                            "direction": "north",
                            "unit": "degree",
                        },
                        {
                            "name": "Geodetic longitude",
                            "abbreviation": "Lon",
                            "direction": "east",
                            "unit": "degree",
                        },
                    ],
                },
                "scope": "Horizontal component of 3D system.",
                "area": "World.",
                "bbox": {
                    "south_latitude": -90,
                    "west_longitude": -180,
                    "north_latitude": 90,
                    "east_longitude": 180,
                },
                "id": {"authority": "EPSG", "code": 4326},
            },
            "wkt": 'GEOGCRS["WGS 84",ENSEMBLE["World Geodetic System 1984 ensemble",MEMBER["World Geodetic System 1984 (Transit)"],MEMBER["World Geodetic System 1984 (G730)"],MEMBER["World Geodetic System 1984 (G873)"],MEMBER["World Geodetic System 1984 (G1150)"],MEMBER["World Geodetic System 1984 (G1674)"],MEMBER["World Geodetic System 1984 (G1762)"],ELLIPSOID["WGS 84",6378137,298.257223563,LENGTHUNIT["metre",1]],ENSEMBLEACCURACY[2.0]],PRIMEM["Greenwich",0,ANGLEUNIT["degree",0.0174532925199433]],CS[ellipsoidal,2],AXIS["geodetic latitude (Lat)",north,ORDER[1],ANGLEUNIT["degree",0.0174532925199433]],AXIS["geodetic longitude (Lon)",east,ORDER[2],ANGLEUNIT["degree",0.0174532925199433]],USAGE[SCOPE["Horizontal component of 3D system."],AREA["World."],BBOX[-90,-180,90,180]],ID["EPSG",4326]]',
            "url": "http://www.opengis.net/def/crs/EPSG/0/4326",
        }
    }

    zattrs = {"_CRS": {crs_member: zattrs_all["_CRS"][crs_member]}}

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(zarray))
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zattrs", json.dumps(zattrs))
        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        rg = ds.GetRootGroup()
        assert rg
        ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
        srs = ar.GetSpatialRef()
        if (
            not (osr.GetPROJVersionMajor() > 6 or osr.GetPROJVersionMinor() >= 2)
            and crs_member == "projjson"
        ):
            assert srs is None
        else:
            assert srs is not None
            assert srs.GetAuthorityCode(None) == "4326"
            assert len(ar.GetAttributes()) == 0
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize("use_get_names", [True, False])
def test_zarr_read_group(use_get_names):

    filename = "data/zarr/group.zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    assert ds is not None
    rg = ds.GetRootGroup()
    assert rg.GetName() == "/"
    assert rg.GetFullName() == "/"
    if use_get_names:
        assert rg.GetGroupNames() == ["foo"]
    assert len(rg.GetAttributes()) == 1
    assert rg.GetAttribute("key") is not None
    subgroup = rg.OpenGroup("foo")
    assert subgroup is not None
    assert rg.OpenGroup("not_existing") is None
    assert subgroup.GetName() == "foo"
    assert subgroup.GetFullName() == "/foo"
    assert rg.GetMDArrayNames() is None
    if use_get_names:
        assert subgroup.GetGroupNames() == ["bar"]
    assert subgroup.GetAttributes() == []
    subsubgroup = subgroup.OpenGroup("bar")
    assert subsubgroup.GetName() == "bar"
    assert subsubgroup.GetFullName() == "/foo/bar"
    if use_get_names:
        assert subsubgroup.GetMDArrayNames() == ["baz"]
    ar = subsubgroup.OpenMDArray("baz")
    assert ar is not None
    assert ar.Read() == array.array("i", [1])
    assert subsubgroup.OpenMDArray("not_existing") is None


def test_zarr_read_group_with_zmetadata():

    filename = "data/zarr/group_with_zmetadata.zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    assert ds is not None
    rg = ds.GetRootGroup()
    assert rg.GetName() == "/"
    assert rg.GetFullName() == "/"
    assert rg.GetGroupNames() == ["foo"]
    assert len(rg.GetAttributes()) == 1
    assert rg.GetAttribute("key") is not None
    subgroup = rg.OpenGroup("foo")
    assert subgroup is not None
    assert rg.OpenGroup("not_existing") is None
    assert subgroup.GetName() == "foo"
    assert subgroup.GetFullName() == "/foo"
    assert rg.GetMDArrayNames() is None
    assert subgroup.GetGroupNames() == ["bar"]
    assert subgroup.GetAttributes() == []
    subsubgroup = subgroup.OpenGroup("bar")
    assert subsubgroup.GetName() == "bar"
    assert subsubgroup.GetFullName() == "/foo/bar"
    assert subsubgroup.GetMDArrayNames() == ["baz"]
    assert subsubgroup.GetAttribute("foo") is not None
    ar = subsubgroup.OpenMDArray("baz")
    assert ar is not None
    assert ar.Read() == array.array("i", [1])
    assert ar.GetAttribute("bar") is not None
    assert subsubgroup.OpenMDArray("not_existing") is None


@pytest.mark.parametrize(
    "use_zmetadata, filename",
    [
        (True, "data/zarr/array_dimensions.zarr"),
        (False, "data/zarr/array_dimensions.zarr"),
        (True, "data/zarr/array_dimensions_upper_level.zarr"),
        (False, "data/zarr/array_dimensions_upper_level.zarr"),
        (False, "data/zarr/array_dimensions_upper_level.zarr/subgroup/var"),
    ],
)
def test_zarr_read_ARRAY_DIMENSIONS(use_zmetadata, filename):

    ds = gdal.OpenEx(
        filename,
        gdal.OF_MULTIDIM_RASTER,
        open_options=["USE_ZMETADATA=" + str(use_zmetadata)],
    )
    assert ds is not None
    rg = ds.GetRootGroup()
    if filename != "data/zarr/array_dimensions_upper_level.zarr":
        ar = rg.OpenMDArray("var")
    else:
        ar = rg.OpenGroup("subgroup").OpenMDArray("var")
    assert ar
    dims = ar.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetName() == "lat"
    assert dims[0].GetIndexingVariable() is not None
    assert dims[0].GetIndexingVariable().GetName() == "lat"
    assert dims[0].GetType() == gdal.DIM_TYPE_HORIZONTAL_Y
    assert dims[0].GetDirection() == "NORTH"
    assert dims[1].GetName() == "lon"
    assert dims[1].GetIndexingVariable() is not None
    assert dims[1].GetIndexingVariable().GetName() == "lon"
    assert dims[1].GetType() == gdal.DIM_TYPE_HORIZONTAL_X
    assert dims[1].GetDirection() == "EAST"
    assert len(rg.GetDimensions()) == 2

    ds = gdal.OpenEx(
        filename,
        gdal.OF_MULTIDIM_RASTER,
        open_options=["USE_ZMETADATA=" + str(use_zmetadata)],
    )
    assert ds is not None
    rg = ds.GetRootGroup()
    ar = rg.OpenMDArray("lat")
    assert ar
    dims = ar.GetDimensions()
    assert len(dims) == 1
    assert dims[0].GetName() == "lat"
    assert dims[0].GetIndexingVariable() is not None
    assert dims[0].GetIndexingVariable().GetName() == "lat"
    assert dims[0].GetType() == gdal.DIM_TYPE_HORIZONTAL_Y
    assert len(rg.GetDimensions()) == 2

    ds = gdal.OpenEx(
        filename,
        gdal.OF_MULTIDIM_RASTER,
        open_options=["USE_ZMETADATA=" + str(use_zmetadata)],
    )
    assert ds is not None
    rg = ds.GetRootGroup()
    assert len(rg.GetDimensions()) == 2


@pytest.mark.parametrize("use_get_names", [True, False])
def test_zarr_read_v3(use_get_names):

    filename = "data/zarr/v3/test.zr3"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    assert ds is not None
    rg = ds.GetRootGroup()
    assert rg.GetName() == "/"
    assert rg.GetFullName() == "/"
    if use_get_names:
        assert rg.GetGroupNames() == ["marvin"]
    assert len(rg.GetAttributes()) == 1
    assert rg.GetAttribute("root_foo") is not None
    subgroup = rg.OpenGroup("marvin")
    assert subgroup is not None
    assert rg.OpenGroup("not_existing") is None
    assert subgroup.GetName() == "marvin"
    assert subgroup.GetFullName() == "/marvin"
    if use_get_names:
        assert rg.GetMDArrayNames() == ["/", "ar"]

    ar = rg.OpenMDArray("/")
    assert ar
    assert ar.Read() == array.array("i", [2] + ([1] * (5 * 10 - 1)))

    ar = rg.OpenMDArray("ar")
    assert ar
    assert ar.Read() == array.array("b", [1, 2])

    if use_get_names:
        assert subgroup.GetGroupNames() == ["paranoid"]
    assert len(subgroup.GetAttributes()) == 1

    subsubgroup = subgroup.OpenGroup("paranoid")
    assert subsubgroup.GetName() == "paranoid"
    assert subsubgroup.GetFullName() == "/marvin/paranoid"

    if use_get_names:
        assert subgroup.GetMDArrayNames() == ["android"]
    ar = subgroup.OpenMDArray("android")
    assert ar is not None
    assert ar.Read() == array.array("b", [1] * 4 * 5)
    assert subgroup.OpenMDArray("not_existing") is None


@pytest.mark.parametrize("endianness", ["le", "be"])
def test_zarr_read_half_float(endianness):

    filename = "data/zarr/f2_" + endianness + ".zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    assert ds is not None
    rg = ds.GetRootGroup()
    ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
    assert ar.Read() == array.array("f", [1.5, float("nan")])


def test_zarr_read_mdim_zarr_non_existing():

    with gdaltest.error_handler():
        assert (
            gdal.OpenEx('ZARR:"data/zarr/not_existing.zarr"', gdal.OF_MULTIDIM_RASTER)
            is None
        )

    with gdaltest.error_handler():
        assert (
            gdal.OpenEx(
                'ZARR:"https://example.org/not_existing.zarr"', gdal.OF_MULTIDIM_RASTER
            )
            is None
        )
    assert (
        "The filename should likely be prefixed with /vsicurl/"
        in gdal.GetLastErrorMsg()
    )

    with gdaltest.error_handler():
        assert (
            gdal.OpenEx(
                "ZARR:https://example.org/not_existing.zarr", gdal.OF_MULTIDIM_RASTER
            )
            is None
        )
    assert (
        "There is likely a quoting error of the whole connection string, and the filename should likely be prefixed with /vsicurl/"
        in gdal.GetLastErrorMsg()
    )

    with gdaltest.error_handler():
        assert (
            gdal.OpenEx(
                "ZARR:/vsicurl/https://example.org/not_existing.zarr",
                gdal.OF_MULTIDIM_RASTER,
            )
            is None
        )
    assert (
        "There is likely a quoting error of the whole connection string."
        in gdal.GetLastErrorMsg()
    )


def test_zarr_read_classic():

    ds = gdal.Open("data/zarr/zlib.zarr")
    assert ds
    assert not ds.GetSubDatasets()
    assert ds.ReadRaster() == array.array("b", [1, 2])

    ds = gdal.Open("ZARR:data/zarr/zlib.zarr")
    assert ds
    assert not ds.GetSubDatasets()
    assert ds.ReadRaster() == array.array("b", [1, 2])

    with gdaltest.error_handler():
        assert gdal.Open('ZARR:"data/zarr/not_existing.zarr"') is None
        assert gdal.Open('ZARR:"data/zarr/zlib.zarr":/not_existing') is None
        assert gdal.Open('ZARR:"data/zarr/zlib.zarr":/zlib:0') is None

    ds = gdal.Open('ZARR:"data/zarr/zlib.zarr":/zlib')
    assert ds
    assert not ds.GetSubDatasets()
    assert ds.ReadRaster() == array.array("b", [1, 2])

    ds = gdal.Open("data/zarr/order_f_u1_3d.zarr")
    assert ds
    subds = ds.GetSubDatasets()
    assert len(subds) == 2
    ds = gdal.Open(subds[0][0])
    assert ds
    assert ds.ReadRaster() == array.array("b", [i for i in range(12)])
    ds = gdal.Open(subds[1][0])
    assert ds
    assert ds.ReadRaster() == array.array("b", [12 + i for i in range(12)])

    with gdaltest.error_handler():
        assert gdal.Open("ZARR:data/zarr/order_f_u1_3d.zarr:/order_f_u1_3d") is None
        assert gdal.Open("ZARR:data/zarr/order_f_u1_3d.zarr:/order_f_u1_3d:2") is None
        assert gdal.Open(subds[0][0] + ":0") is None

    ds = gdal.Open("data/zarr/v3/test.zr3")
    assert ds
    subds = ds.GetSubDatasets()
    assert set(subds) == set(
        [
            ('ZARR:"data/zarr/v3/test.zr3":/', "Array /"),
            ('ZARR:"data/zarr/v3/test.zr3":/ar', "Array /ar"),
            ('ZARR:"data/zarr/v3/test.zr3":/marvin/android', "Array /marvin/android"),
        ]
    )
    ds = gdal.Open('ZARR:"data/zarr/v3/test.zr3":/')
    assert ds
    assert ds.ReadRaster() == array.array("i", [2] + ([1] * (10 * 5 - 1)))
    ds = gdal.Open('ZARR:"data/zarr/v3/test.zr3":/ar')
    assert ds
    assert ds.ReadRaster() == array.array("b", [1, 2])


def test_zarr_read_classic_2d():

    try:
        src_ds = gdal.Open("data/byte.tif")
        gdal.GetDriverByName("ZARR").CreateCopy(
            "/vsimem/test.zarr", src_ds, strict=False
        )
        ds = gdal.Open("/vsimem/test.zarr")
        assert ds is not None
        assert len(ds.GetSubDatasets()) == 0
        ds = None
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_classic_2d_with_unrelated_auxiliary_1D_arrays():

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr"
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg
            assert rg.GetName() == "/"

            dim0 = rg.CreateDimension("dim0", None, None, 2)
            dim1 = rg.CreateDimension("dim1", None, None, 3)

            dt = gdal.ExtendedDataType.Create(gdal.GDT_Float64)
            rg.CreateMDArray("main_array", [dim0, dim1], dt)
            rg.CreateMDArray("x", [dim0], dt)
            rg.CreateMDArray("y", [dim1], dt)

        create()
        ds = gdal.Open("/vsimem/test.zarr")
        assert ds is not None
        assert ds.RasterYSize == 2
        assert ds.RasterXSize == 3
        assert set(ds.GetSubDatasets()) == set(
            [
                ('ZARR:"/vsimem/test.zarr":/main_array', "Array /main_array"),
                ('ZARR:"/vsimem/test.zarr":/x', "Array /x"),
                ('ZARR:"/vsimem/test.zarr":/y', "Array /y"),
            ]
        )
        ds = None
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_classic_too_many_samples_3d():

    j = {
        "chunks": [65536, 2, 1],
        "compressor": None,
        "dtype": "!u1",
        "fill_value": None,
        "filters": None,
        "order": "C",
        "shape": [65536, 2, 1],
        "zarr_format": 2,
    }

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))
        gdal.ErrorReset()
        with gdaltest.error_handler():
            ds = gdal.Open("/vsimem/test.zarr")
        assert gdal.GetLastErrorMsg() != ""
        assert len(ds.GetSubDatasets()) == 0
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_classic_4d():

    j = {
        "chunks": [3, 2, 1, 1],
        "compressor": None,
        "dtype": "!u1",
        "fill_value": None,
        "filters": None,
        "order": "C",
        "shape": [3, 2, 1, 1],
        "zarr_format": 2,
    }

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))
        ds = gdal.Open("/vsimem/test.zarr")
        subds = ds.GetSubDatasets()
        assert len(subds) == 6
        for i in range(len(subds)):
            assert gdal.Open(subds[i][0]) is not None
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_classic_too_many_samples_4d():

    j = {
        "chunks": [256, 256, 1, 1],
        "compressor": None,
        "dtype": "!u1",
        "fill_value": None,
        "filters": None,
        "order": "C",
        "shape": [256, 256, 1, 1],
        "zarr_format": 2,
    }

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))
        gdal.ErrorReset()
        with gdaltest.error_handler():
            ds = gdal.Open("/vsimem/test.zarr")
        assert gdal.GetLastErrorMsg() != ""
        assert len(ds.GetSubDatasets()) == 0
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_empty_shape():

    ds = gdal.OpenEx("data/zarr/empty.zarr", gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg
    ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
    assert ar
    assert ar.Read() == array.array("b", [120])


def test_zarr_read_BLOSC_COMPRESSORS():

    if "blosc" not in gdal.GetDriverByName("Zarr").GetMetadataItem("COMPRESSORS"):
        pytest.skip("blosc not available")
    assert "lz4" in gdal.GetDriverByName("Zarr").GetMetadataItem("BLOSC_COMPRESSORS")


@pytest.mark.parametrize(
    "format,create_z_metadata",
    [("ZARR_V2", "YES"), ("ZARR_V2", "NO"), ("ZARR_V3", "NO")],
)
def test_zarr_create_group(format, create_z_metadata):

    filename = "tmp/test.zarr"
    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                filename,
                options=["FORMAT=" + format, "CREATE_ZMETADATA=" + create_z_metadata],
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg
            assert rg.GetName() == "/"

            attr = rg.CreateAttribute(
                "str_attr", [], gdal.ExtendedDataType.CreateString()
            )
            assert attr
            assert attr.Write("my_string") == gdal.CE_None

            attr = rg.CreateAttribute(
                "json_attr", [], gdal.ExtendedDataType.CreateString(0, gdal.GEDTST_JSON)
            )
            assert attr
            assert attr.Write({"foo": "bar"}) == gdal.CE_None

            attr = rg.CreateAttribute(
                "str_array_attr", [2], gdal.ExtendedDataType.CreateString()
            )
            assert attr
            assert attr.Write(["first_string", "second_string"]) == gdal.CE_None

            with gdaltest.error_handler():
                attr = rg.CreateAttribute(
                    "dim_2_not_supported", [2, 2], gdal.ExtendedDataType.CreateString()
                )
                assert attr is None

            attr = rg.CreateAttribute(
                "int_attr", [], gdal.ExtendedDataType.Create(gdal.GDT_Int32)
            )
            assert attr
            assert attr.Write(12345678) == gdal.CE_None

            attr = rg.CreateAttribute(
                "uint_attr", [], gdal.ExtendedDataType.Create(gdal.GDT_UInt32)
            )
            assert attr
            assert attr.Write(4000000000) == gdal.CE_None

            attr = rg.CreateAttribute(
                "int_array_attr", [2], gdal.ExtendedDataType.Create(gdal.GDT_Int32)
            )
            assert attr
            assert attr.Write([12345678, -12345678]) == gdal.CE_None

            attr = rg.CreateAttribute(
                "double_attr", [], gdal.ExtendedDataType.Create(gdal.GDT_Float64)
            )
            assert attr
            assert attr.Write(12345678.5) == gdal.CE_None

            attr = rg.CreateAttribute(
                "double_array_attr", [2], gdal.ExtendedDataType.Create(gdal.GDT_Float64)
            )
            assert attr
            assert attr.Write([12345678.5, -12345678.5]) == gdal.CE_None

            subgroup = rg.CreateGroup("foo")
            assert subgroup
            assert subgroup.GetName() == "foo"
            assert subgroup.GetFullName() == "/foo"
            assert rg.GetGroupNames() == ["foo"]
            subgroup = rg.OpenGroup("foo")
            assert subgroup

        create()

        if create_z_metadata == "YES":
            f = gdal.VSIFOpenL(filename + "/.zmetadata", "rb")
            assert f
            data = gdal.VSIFReadL(1, 10000, f)
            gdal.VSIFCloseL(f)
            j = json.loads(data)
            assert "foo/.zgroup" in j["metadata"]

        def update():
            ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE)
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            assert rg.GetGroupNames() == ["foo"]

            attr = rg.GetAttribute("str_attr")
            assert attr
            assert attr.Read() == "my_string"
            assert attr.Write("my_string_modified") == gdal.CE_None

            subgroup = rg.OpenGroup("foo")
            assert subgroup
            subgroup = rg.CreateGroup("bar")
            assert subgroup
            assert set(rg.GetGroupNames()) == set(["foo", "bar"])
            subgroup = rg.OpenGroup("foo")
            assert subgroup
            subsubgroup = subgroup.CreateGroup("baz")
            assert subsubgroup
            ds = None

        update()

        ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
        assert ds
        rg = ds.GetRootGroup()
        assert rg

        attr = rg.GetAttribute("str_attr")
        assert attr
        assert attr.Read() == "my_string_modified"

        attr = rg.GetAttribute("json_attr")
        assert attr
        assert attr.GetDataType().GetSubType() == gdal.GEDTST_JSON
        assert attr.Read() == {"foo": "bar"}

        attr = rg.GetAttribute("str_array_attr")
        assert attr
        assert attr.Read() == ["first_string", "second_string"]

        attr = rg.GetAttribute("int_attr")
        assert attr
        assert attr.GetDataType().GetNumericDataType() == gdal.GDT_Int32
        assert attr.ReadAsDouble() == 12345678

        attr = rg.GetAttribute("uint_attr")
        assert attr
        assert attr.GetDataType().GetNumericDataType() == gdal.GDT_Float64
        assert attr.ReadAsDouble() == 4000000000

        attr = rg.GetAttribute("int_array_attr")
        assert attr
        assert attr.GetDataType().GetNumericDataType() == gdal.GDT_Int32
        assert attr.ReadAsIntArray() == (12345678, -12345678)

        attr = rg.GetAttribute("double_attr")
        assert attr
        assert attr.GetDataType().GetNumericDataType() == gdal.GDT_Float64
        assert attr.ReadAsDouble() == 12345678.5

        attr = rg.GetAttribute("double_array_attr")
        assert attr
        assert attr.GetDataType().GetNumericDataType() == gdal.GDT_Float64
        assert attr.Read() == (12345678.5, -12345678.5)

        assert set(rg.GetGroupNames()) == set(["foo", "bar"])
        with gdaltest.error_handler():
            assert rg.CreateGroup("not_opened_in_update_mode") is None
            assert (
                rg.CreateAttribute(
                    "not_opened_in_update_mode",
                    [],
                    gdal.ExtendedDataType.CreateString(),
                )
                is None
            )
        subgroup = rg.OpenGroup("foo")
        assert subgroup
        subsubgroup = subgroup.OpenGroup("baz")
        assert subsubgroup
        ds = None

    finally:
        gdal.RmdirRecursive(filename)


@pytest.mark.parametrize(
    "group_name",
    [
        "foo",  # already existing
        "directory_with_that_name",
        "",
        ".",
        "..",
        "a/b",
        "a\\n",
        "a:b",
        ".zarray",
    ],
)
@pytest.mark.parametrize("format", ["ZARR_V2", "ZARR_V3"])
def test_zarr_create_group_errors(group_name, format):

    try:
        ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
            "/vsimem/test.zarr", options=["FORMAT=" + format]
        )
        assert ds is not None
        rg = ds.GetRootGroup()
        assert rg
        subgroup = rg.CreateGroup("foo")
        assert subgroup
        if format == "ZARR_V2":
            gdal.Mkdir("/vsimem/test.zarr/directory_with_that_name", 0)
        else:
            gdal.Mkdir("/vsimem/test.zarr/meta/root/directory_with_that_name", 0)
        with gdaltest.error_handler():
            assert rg.CreateGroup(group_name) is None

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def getCompoundDT():
    x = gdal.EDTComponent.Create("x", 0, gdal.ExtendedDataType.Create(gdal.GDT_Int16))
    y = gdal.EDTComponent.Create("y", 0, gdal.ExtendedDataType.Create(gdal.GDT_Int32))
    subcompound = gdal.ExtendedDataType.CreateCompound("", 4, [y])
    subcompound_component = gdal.EDTComponent.Create("y", 4, subcompound)
    return gdal.ExtendedDataType.CreateCompound("", 8, [x, subcompound_component])


@pytest.mark.parametrize(
    "datatype,nodata",
    [
        [gdal.ExtendedDataType.Create(gdal.GDT_Byte), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_Byte), 1],
        [gdal.ExtendedDataType.Create(gdal.GDT_UInt16), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_Int16), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_UInt32), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_Int32), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_Float32), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_Float64), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_Float64), 1.5],
        [gdal.ExtendedDataType.Create(gdal.GDT_Float64), float("nan")],
        [gdal.ExtendedDataType.Create(gdal.GDT_Float64), float("infinity")],
        [gdal.ExtendedDataType.Create(gdal.GDT_Float64), float("-infinity")],
        [gdal.ExtendedDataType.Create(gdal.GDT_CInt16), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_CInt32), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_CFloat32), None],
        [gdal.ExtendedDataType.Create(gdal.GDT_CFloat64), None],
        [gdal.ExtendedDataType.CreateString(10), None],
        [gdal.ExtendedDataType.CreateString(10), "ab"],
        [getCompoundDT(), None],
        [
            getCompoundDT(),
            bytes(array.array("h", [12]))
            + bytes(array.array("h", [0]))
            + bytes(array.array("i", [2345678])),  # padding
        ],
    ],
)
@pytest.mark.parametrize("format", ["ZARR_V2", "ZARR_V3"])
def test_zarr_create_array(datatype, nodata, format):

    error_expected = False
    if format == "ZARR_V3":
        if datatype.GetClass() != gdal.GEDTC_NUMERIC or gdal.DataTypeIsComplex(
            datatype.GetNumericDataType()
        ):
            error_expected = True
    elif datatype.GetNumericDataType() in (gdal.GDT_CInt16, gdal.GDT_CInt32):
        error_expected = True

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr", options=["FORMAT=" + format]
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg
            assert rg.GetName() == "/"

            dim0 = rg.CreateDimension("dim0", None, None, 2)
            dim1 = rg.CreateDimension("dim1", None, None, 3)

            if error_expected:
                with gdaltest.error_handler():
                    ar = rg.CreateMDArray("my_ar", [dim0, dim1], datatype)
                assert ar is None
                return False
            else:
                ar = rg.CreateMDArray("my_ar", [dim0, dim1], datatype)
                assert ar
                if nodata:
                    if datatype.GetClass() == gdal.GEDTC_STRING:
                        assert ar.SetNoDataValueString(nodata) == gdal.CE_None
                    elif datatype.GetClass() == gdal.GEDTC_NUMERIC:
                        assert ar.SetNoDataValueDouble(nodata) == gdal.CE_None
                    else:
                        assert ar.SetNoDataValueRaw(nodata) == gdal.CE_None
                return True

        if create():
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray("my_ar")
            assert ar
            got_dt = ar.GetDataType()
            if got_dt.GetClass() == gdal.GEDTC_COMPOUND:
                comps = got_dt.GetComponents()
                assert len(comps) == 2
                assert comps[1].GetType().GetClass() == gdal.GEDTC_COMPOUND
                comps[1] = gdal.EDTComponent.Create(
                    comps[1].GetName(),
                    comps[1].GetType().GetSize(),
                    gdal.ExtendedDataType.CreateCompound(
                        "",
                        comps[1].GetType().GetSize(),
                        comps[1].GetType().GetComponents(),
                    ),
                )
                got_dt = gdal.ExtendedDataType.CreateCompound(
                    "", got_dt.GetSize(), comps
                )
            assert got_dt == datatype
            assert len(ar.GetDimensions()) == 2
            assert [ar.GetDimensions()[i].GetSize() for i in range(2)] == [2, 3]
            if nodata:
                if datatype.GetClass() == gdal.GEDTC_STRING:
                    got_nodata = ar.GetNoDataValueAsString()
                    assert got_nodata == nodata
                elif datatype.GetClass() == gdal.GEDTC_NUMERIC:
                    got_nodata = ar.GetNoDataValueAsDouble()
                    if math.isnan(nodata):
                        assert math.isnan(got_nodata)
                    else:
                        assert got_nodata == nodata
                else:
                    got_nodata = ar.GetNoDataValueAsRaw()
                    assert got_nodata == nodata
            else:
                assert ar.GetNoDataValueAsRaw() is None

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize(
    "array_name",
    [
        "foo",  # already existing
        "directory_with_that_name",
        "",
        ".",
        "..",
        "a/b",
        "a\\n",
        "a:b",
        ".zarray",
    ],
)
def test_zarr_create_array_errors(array_name):

    try:
        ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional("/vsimem/test.zarr")
        assert ds is not None
        rg = ds.GetRootGroup()
        assert rg
        assert (
            rg.CreateMDArray("foo", [], gdal.ExtendedDataType.Create(gdal.GDT_Byte))
            is not None
        )
        gdal.Mkdir("/vsimem/test.zarr/directory_with_that_name", 0)
        with gdaltest.error_handler():
            assert (
                rg.CreateMDArray(
                    array_name, [], gdal.ExtendedDataType.Create(gdal.GDT_Byte)
                )
                is None
            )

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize(
    "compressor,options,expected_json",
    [
        ["NONE", [], None],
        ["zlib", [], {"id": "zlib", "level": 6}],
        ["zlib", ["ZLIB_LEVEL=1"], {"id": "zlib", "level": 1}],
        [
            "blosc",
            [],
            {"blocksize": 0, "clevel": 5, "cname": "lz4", "id": "blosc", "shuffle": 1},
        ],
    ],
)
def test_zarr_create_array_compressor(compressor, options, expected_json):

    compressors = gdal.GetDriverByName("Zarr").GetMetadataItem("COMPRESSORS")
    if compressor != "NONE" and compressor not in compressors:
        pytest.skip("compressor %s not available" % compressor)

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr"
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg
            assert (
                rg.CreateMDArray(
                    "test",
                    [],
                    gdal.ExtendedDataType.Create(gdal.GDT_Byte),
                    ["COMPRESS=" + compressor] + options,
                )
                is not None
            )

        create()

        f = gdal.VSIFOpenL("/vsimem/test.zarr/test/.zarray", "rb")
        assert f
        data = gdal.VSIFReadL(1, 1000, f)
        gdal.VSIFCloseL(f)
        j = json.loads(data)
        assert j["compressor"] == expected_json

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize(
    "compressor,options,expected_json",
    [
        ["NONE", [], None],
        [
            "gzip",
            [],
            {
                "codec": "https://purl.org/zarr/spec/codec/gzip/1.0",
                "configuration": {"level": 6},
            },
        ],
    ],
)
def test_zarr_create_array_compressor_v3(compressor, options, expected_json):

    compressors = gdal.GetDriverByName("Zarr").GetMetadataItem("COMPRESSORS")
    if compressor != "NONE" and compressor not in compressors:
        pytest.skip("compressor %s not available" % compressor)

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr", options=["FORMAT=ZARR_V3"]
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg
            assert (
                rg.CreateMDArray(
                    "test",
                    [],
                    gdal.ExtendedDataType.Create(gdal.GDT_Byte),
                    ["COMPRESS=" + compressor] + options,
                )
                is not None
            )

        create()

        f = gdal.VSIFOpenL("/vsimem/test.zarr/meta/root/test.array.json", "rb")
        assert f
        data = gdal.VSIFReadL(1, 1000, f)
        gdal.VSIFCloseL(f)
        j = json.loads(data)
        if expected_json is None:
            assert "compressor" not in j
        else:
            assert j["compressor"] == expected_json

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize("format", ["ZARR_V2", "ZARR_V3"])
def test_zarr_create_array_bad_compressor(format):

    try:
        ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
            "/vsimem/test.zarr", options=["FORMAT=" + format]
        )
        assert ds is not None
        rg = ds.GetRootGroup()
        assert rg
        with gdaltest.error_handler():
            assert (
                rg.CreateMDArray(
                    "test",
                    [],
                    gdal.ExtendedDataType.Create(gdal.GDT_Byte),
                    ["COMPRESS=invalid"],
                )
                is None
            )
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize("format", ["ZARR_V2", "ZARR_V3"])
def test_zarr_create_array_attributes(format):

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr", options=["FORMAT=" + format]
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.CreateMDArray(
                "test", [], gdal.ExtendedDataType.Create(gdal.GDT_Byte)
            )
            assert ar

            attr = ar.CreateAttribute(
                "str_attr", [], gdal.ExtendedDataType.CreateString()
            )
            assert attr
            assert attr.Write("my_string") == gdal.CE_None

            with gdaltest.error_handler():
                assert (
                    ar.CreateAttribute(
                        "invalid_2d", [2, 3], gdal.ExtendedDataType.CreateString()
                    )
                    is None
                )

        create()

        def update():
            ds = gdal.OpenEx(
                "/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE
            )
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray("test")
            assert ar

            attr = ar.GetAttribute("str_attr")
            assert attr
            assert attr.Read() == "my_string"
            assert attr.Write("my_string_modified") == gdal.CE_None

        update()

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        assert ds
        rg = ds.GetRootGroup()
        assert rg
        ar = rg.OpenMDArray("test")
        assert ar

        attr = ar.GetAttribute("str_attr")
        assert attr
        assert attr.Read() == "my_string_modified"
        with gdaltest.error_handler():
            assert attr.Write("foo") == gdal.CE_Failure

        with gdaltest.error_handler():
            assert (
                ar.CreateAttribute(
                    "another_attr", [], gdal.ExtendedDataType.CreateString()
                )
                is None
            )
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_create_array_set_crs():

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr"
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.CreateMDArray(
                "test", [], gdal.ExtendedDataType.Create(gdal.GDT_Byte)
            )
            assert ar
            crs = osr.SpatialReference()
            crs.ImportFromEPSG(4326)
            assert ar.SetSpatialRef(crs) == gdal.CE_None

        create()

        f = gdal.VSIFOpenL("/vsimem/test.zarr/test/.zattrs", "rb")
        assert f
        data = gdal.VSIFReadL(1, 10000, f)
        gdal.VSIFCloseL(f)
        j = json.loads(data)
        assert "_CRS" in j
        crs = j["_CRS"]
        assert "wkt" in crs
        assert "url" in crs
        if "projjson" in crs:
            assert crs["projjson"]["type"] == "GeographicCRS"

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_create_array_set_dimension_name():

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr"
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg

            dim0 = rg.CreateDimension("dim0", None, None, 2)
            dim0_ar = rg.CreateMDArray(
                "dim0", [dim0], gdal.ExtendedDataType.Create(gdal.GDT_Byte)
            )
            dim0.SetIndexingVariable(dim0_ar)

            rg.CreateMDArray(
                "test", [dim0], gdal.ExtendedDataType.Create(gdal.GDT_Byte)
            )

        create()

        f = gdal.VSIFOpenL("/vsimem/test.zarr/test/.zattrs", "rb")
        assert f
        data = gdal.VSIFReadL(1, 10000, f)
        gdal.VSIFCloseL(f)
        j = json.loads(data)
        assert "_ARRAY_DIMENSIONS" in j
        assert j["_ARRAY_DIMENSIONS"] == ["dim0"]

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize(
    "dtype,structtype,gdaltype,fill_value,nodata_value",
    [
        ["!b1", "B", gdal.GDT_Byte, None, None],
        ["!i1", "b", gdal.GDT_Int16, None, None],
        ["!i1", "b", gdal.GDT_Int16, -1, -1],
        ["!u1", "B", gdal.GDT_Byte, None, None],
        ["!u1", "B", gdal.GDT_Byte, "1", 1],
        ["<i2", "h", gdal.GDT_Int16, None, None],
        [">i2", "h", gdal.GDT_Int16, None, None],
        ["<i4", "i", gdal.GDT_Int32, None, None],
        [">i4", "i", gdal.GDT_Int32, None, None],
        ["<i8", "q", gdal.GDT_Float64, None, None],
        [">i8", "q", gdal.GDT_Float64, None, None],
        ["<u2", "H", gdal.GDT_UInt16, None, None],
        [">u2", "H", gdal.GDT_UInt16, None, None],
        ["<u4", "I", gdal.GDT_UInt32, None, None],
        [">u4", "I", gdal.GDT_UInt32, None, None],
        ["<u4", "I", gdal.GDT_UInt32, 4000000000, 4000000000],
        ["<u8", "Q", gdal.GDT_Float64, 4000000000, 4000000000],
        [">u8", "Q", gdal.GDT_Float64, None, None],
        ["<f4", "f", gdal.GDT_Float32, None, None],
        [">f4", "f", gdal.GDT_Float32, None, None],
        ["<f4", "f", gdal.GDT_Float32, 1.5, 1.5],
        ["<f4", "f", gdal.GDT_Float32, "NaN", float("nan")],
        ["<f4", "f", gdal.GDT_Float32, "Infinity", float("infinity")],
        ["<f4", "f", gdal.GDT_Float32, "-Infinity", float("-infinity")],
        ["<f8", "d", gdal.GDT_Float64, None, None],
        [">f8", "d", gdal.GDT_Float64, None, None],
        ["<f8", "d", gdal.GDT_Float64, "NaN", float("nan")],
        ["<f8", "d", gdal.GDT_Float64, "Infinity", float("infinity")],
        ["<f8", "d", gdal.GDT_Float64, "-Infinity", float("-infinity")],
        ["<c8", "f", gdal.GDT_CFloat32, None, None],
        [">c8", "f", gdal.GDT_CFloat32, None, None],
        ["<c16", "d", gdal.GDT_CFloat64, None, None],
        [">c16", "d", gdal.GDT_CFloat64, None, None],
    ],
)
@pytest.mark.parametrize("use_optimized_code_paths", [True, False])
def test_zarr_write_array_content(
    dtype, structtype, gdaltype, fill_value, nodata_value, use_optimized_code_paths
):

    j = {
        "chunks": [2, 3],
        "compressor": None,
        "dtype": dtype,
        "fill_value": fill_value,
        "filters": None,
        "order": "C",
        "shape": [5, 4],
        "zarr_format": 2,
    }

    filename = (
        "/vsimem/test"
        + dtype.replace("<", "lt").replace(">", "gt").replace("!", "not")
        + structtype
        + ".zarr"
    )
    try:
        gdal.Mkdir(filename, 0o755)
        f = gdal.VSIFOpenL(filename + "/.zarray", "wb")
        assert f
        data = json.dumps(j)
        gdal.VSIFWriteL(data, 1, len(data), f)
        gdal.VSIFCloseL(f)

        if gdaltype not in (gdal.GDT_CFloat32, gdal.GDT_CFloat64):
            tile_0_0_data = struct.pack(dtype[0] + (structtype * 6), 1, 2, 3, 5, 6, 7)
            tile_0_1_data = struct.pack(dtype[0] + (structtype * 6), 4, 0, 0, 8, 0, 0)
        else:
            tile_0_0_data = struct.pack(
                dtype[0] + (structtype * 12), 1, 11, 2, 0, 3, 0, 5, 0, 6, 0, 7, 0
            )
            tile_0_1_data = struct.pack(
                dtype[0] + (structtype * 12), 4, 0, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0
            )
        gdal.FileFromMemBuffer(filename + "/0.0", tile_0_0_data)
        gdal.FileFromMemBuffer(filename + "/0.1", tile_0_1_data)

        with gdaltest.config_option(
            "GDAL_ZARR_USE_OPTIMIZED_CODE_PATHS",
            "YES" if use_optimized_code_paths else "NO",
        ):
            ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE)
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
        assert ar

        dt = gdal.ExtendedDataType.Create(
            gdal.GDT_CFloat64
            if gdaltype in (gdal.GDT_CFloat32, gdal.GDT_CFloat64)
            else gdal.GDT_Float64
        )

        # Write all nodataset. That should cause tiles to be removed.
        nv = nodata_value if nodata_value else 0
        buf_nodata = array.array(
            "d",
            [nv]
            * (
                5 * 4 * (2 if gdaltype in (gdal.GDT_CFloat32, gdal.GDT_CFloat64) else 1)
            ),
        )
        assert ar.Write(buf_nodata, buffer_datatype=dt) == gdal.CE_None
        assert ar.Read(buffer_datatype=dt) == bytearray(buf_nodata)

        if (
            fill_value is None
            or fill_value == 0
            or not gdal.DataTypeIsComplex(gdaltype)
        ):
            assert gdal.VSIStatL(filename + "/0.0") is None

        # Write all ones
        ones = array.array(
            "d",
            [0]
            * (
                5 * 4 * (2 if gdaltype in (gdal.GDT_CFloat32, gdal.GDT_CFloat64) else 1)
            ),
        )
        assert ar.Write(ones, buffer_datatype=dt) == gdal.CE_None
        assert ar.Read(buffer_datatype=dt) == bytearray(ones)

        # Write with odd array_step
        assert (
            ar.Write(
                struct.pack("d" * 4, nv, nv, 6, 5),
                array_start_idx=[2, 1],
                count=[2, 2],
                array_step=[-1, -1],
                buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64),
            )
            == gdal.CE_None
        )

        # Check back
        assert ar.Read(
            array_start_idx=[2, 1],
            count=[2, 2],
            array_step=[-1, -1],
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64),
        ) == struct.pack("d" * 4, nv, nv, 6, 5)

        # Force dirty block eviction
        ar.Read(buffer_datatype=dt)

        # Check back again
        assert ar.Read(
            array_start_idx=[2, 1],
            count=[2, 2],
            array_step=[-1, -1],
            buffer_datatype=gdal.ExtendedDataType.Create(gdal.GDT_Float64),
        ) == struct.pack("d" * 4, nv, nv, 6, 5)

    finally:
        gdal.RmdirRecursive(filename)


@pytest.mark.parametrize(
    "string_format,input_str,output_str",
    [
        ("ASCII", "0123456789truncated", "0123456789"),
        ("UNICODE", "\u00E9" + "123456789truncated", "\u00E9" + "123456789"),
    ],
    ids=("ASCII", "UNICODE"),
)
def test_zarr_create_array_string(string_format, input_str, output_str):

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr"
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg

            dim0 = rg.CreateDimension("dim0", None, None, 2)

            ar = rg.CreateMDArray(
                "test",
                [dim0],
                gdal.ExtendedDataType.CreateString(10),
                ["STRING_FORMAT=" + string_format, "COMPRESS=ZLIB"],
            )
            assert ar.Write(["ab", input_str]) == gdal.CE_None

        create()

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE)
        assert ds
        rg = ds.GetRootGroup()
        assert rg
        ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
        assert ar.Read() == ["ab", output_str]

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize(
    "srcfilename", ["data/zarr/unicode_le.zarr", "data/zarr/unicode_be.zarr"]
)
def test_zarr_update_array_string(srcfilename):

    filename = "/vsimem/test.zarr"

    try:
        gdal.Mkdir(filename, 0)
        gdal.FileFromMemBuffer(
            filename + "/.zarray", open(srcfilename + "/.zarray", "rb").read()
        )
        gdal.FileFromMemBuffer(filename + "/0", open(srcfilename + "/0", "rb").read())

        eta = "\u03B7"

        def update():
            ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE)
            rg = ds.GetRootGroup()
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
            assert ar.Read() == ["\u00E9"]
            assert ar.Write([eta]) == gdal.CE_None
            assert gdal.GetLastErrorMsg() == ""

        update()

        def check():
            ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
            rg = ds.GetRootGroup()
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
            assert ar.Read() == [eta]

        check()

    finally:
        gdal.RmdirRecursive(filename)


@pytest.mark.parametrize("format", ["ZARR_V2", "ZARR_V3"])
def test_zarr_create_fortran_order_3d_and_compression_and_dim_separator(format):

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr", options=["FORMAT=" + format]
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg

            dim0 = rg.CreateDimension("dim0", None, None, 2)
            dim1 = rg.CreateDimension("dim1", None, None, 3)
            dim2 = rg.CreateDimension("dim2", None, None, 4)

            ar = rg.CreateMDArray(
                "test",
                [dim0, dim1, dim2],
                gdal.ExtendedDataType.Create(gdal.GDT_Byte),
                ["CHUNK_MEMORY_LAYOUT=F", "COMPRESS=zlib", "DIM_SEPARATOR=/"],
            )
            assert (
                ar.Write(array.array("b", [i for i in range(2 * 3 * 4)]))
                == gdal.CE_None
            )

        create()

        if format == "ZARR_V2":
            f = gdal.VSIFOpenL("/vsimem/test.zarr/test/.zarray", "rb")
        else:
            f = gdal.VSIFOpenL("/vsimem/test.zarr/meta/root/test.array.json", "rb")
        assert f
        data = gdal.VSIFReadL(1, 10000, f)
        gdal.VSIFCloseL(f)
        j = json.loads(data)
        if format == "ZARR_V2":
            assert "order" in j
            assert j["order"] == "F"
        else:
            assert "chunk_memory_layout" in j
            assert j["chunk_memory_layout"] == "F"

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE)
        assert ds
        rg = ds.GetRootGroup()
        assert rg
        ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
        assert ar.Read() == array.array("b", [i for i in range(2 * 3 * 4)])

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_create_unit_offset_scale():

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr"
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg

            ar = rg.CreateMDArray(
                "test", [], gdal.ExtendedDataType.Create(gdal.GDT_Byte)
            )
            assert ar.SetOffset(1.5) == gdal.CE_None
            assert ar.SetScale(2.5) == gdal.CE_None
            assert ar.SetUnit("my unit") == gdal.CE_None

        create()

        f = gdal.VSIFOpenL("/vsimem/test.zarr/test/.zattrs", "rb")
        assert f
        data = gdal.VSIFReadL(1, 10000, f)
        gdal.VSIFCloseL(f)
        j = json.loads(data)
        assert "add_offset" in j
        assert j["add_offset"] == 1.5
        assert "scale_factor" in j
        assert j["scale_factor"] == 2.5
        assert "units" in j
        assert j["units"] == "my unit"

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE)
        assert ds
        rg = ds.GetRootGroup()
        assert rg
        ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
        assert ar.GetOffset() == 1.5
        assert ar.GetScale() == 2.5
        assert ar.GetUnit() == "my unit"

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_getcoordinatevariables():

    src_ds = gdal.OpenEx(
        "data/netcdf/expanded_form_of_grid_mapping.nc", gdal.OF_MULTIDIM_RASTER
    )
    if src_ds is None:
        pytest.skip()

    try:

        def create(src_ds):
            ds = gdal.MultiDimTranslate("/vsimem/test.zarr", src_ds, format="Zarr")
            src_ds = None
            assert ds
            rg = ds.GetRootGroup()

            ar = rg.OpenMDArray("temp")
            coordinate_vars = ar.GetCoordinateVariables()
            assert len(coordinate_vars) == 2
            assert coordinate_vars[0].GetName() == "lat"
            assert coordinate_vars[1].GetName() == "lon"

            assert len(coordinate_vars[0].GetCoordinateVariables()) == 0

        create(src_ds)

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_create_copy():

    tst = gdaltest.GDALTest("Zarr", "../../gcore/data/uint16.tif", 1, 4672)

    try:
        return tst.testCreate(vsimem=1, new_filename="/vsimem/test.zarr")
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize("format", ["ZARR_V2", "ZARR_V3"])
def test_zarr_create(format):

    try:
        ds = gdal.GetDriverByName("Zarr").Create(
            "/vsimem/test.zarr", 1, 1, 3, options=["ARRAY_NAME=foo", "FORMAT=" + format]
        )
        assert ds.GetGeoTransform(can_return_null=True) is None
        assert ds.GetSpatialRef() is None
        assert ds.GetRasterBand(1).GetNoDataValue() is None
        assert ds.GetRasterBand(1).SetNoDataValue(10) == gdal.CE_None
        assert ds.GetRasterBand(1).GetOffset() is None
        assert ds.GetRasterBand(1).SetOffset(1.5) == gdal.CE_None
        assert ds.GetRasterBand(1).GetScale() is None
        assert ds.GetRasterBand(1).SetScale(2.5) == gdal.CE_None
        assert ds.GetRasterBand(1).GetUnitType() == ""
        assert ds.GetRasterBand(1).SetUnitType("my_unit") == gdal.CE_None
        assert ds.SetMetadata({"FOO": "BAR"}) == gdal.CE_None
        ds = None

        ds = gdal.Open("ZARR:/vsimem/test.zarr:/foo_band1")
        assert ds
        assert ds.GetMetadata() == {"FOO": "BAR"}
        assert ds.GetRasterBand(1).GetNoDataValue() == 10.0
        assert ds.GetRasterBand(1).GetOffset() == 1.5

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_create_append_subdataset():

    try:

        def create():
            ds = gdal.GetDriverByName("Zarr").Create(
                "/vsimem/test.zarr", 3, 2, 1, options=["ARRAY_NAME=foo"]
            )
            assert ds
            ds.SetGeoTransform([2, 1, 0, 49, 0, -1])
            ds = None

            # Same dimensions. Will reuse the ones of foo
            ds = gdal.GetDriverByName("Zarr").Create(
                "/vsimem/test.zarr",
                3,
                2,
                1,
                options=["APPEND_SUBDATASET=YES", "ARRAY_NAME=bar"],
            )
            assert ds
            ds.SetGeoTransform([2, 1, 0, 49, 0, -1])
            ds = None

            # Different dimensions.
            ds = gdal.GetDriverByName("Zarr").Create(
                "/vsimem/test.zarr",
                30,
                20,
                1,
                options=["APPEND_SUBDATASET=YES", "ARRAY_NAME=baz"],
            )
            assert ds
            ds.SetGeoTransform([2, 0.1, 0, 49, 0, -0.1])
            ds = None

        create()

        def check():
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            rg = ds.GetRootGroup()

            foo = rg.OpenMDArray("foo")
            assert foo
            assert foo.GetDimensions()[0].GetName() == "Y"
            assert foo.GetDimensions()[1].GetName() == "X"

            bar = rg.OpenMDArray("bar")
            assert bar
            assert bar.GetDimensions()[0].GetName() == "Y"
            assert bar.GetDimensions()[1].GetName() == "X"

            baz = rg.OpenMDArray("baz")
            assert baz
            assert baz.GetDimensions()[0].GetName() == "baz_Y"
            assert baz.GetDimensions()[1].GetName() == "baz_X"

        check()

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize(
    "blocksize", ["1,2", "2,2,0", "4000000000,4000000000,4000000000"]
)
def test_zarr_create_array_invalid_blocksize(blocksize):

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr"
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg

            dim0 = rg.CreateDimension("dim0", None, None, 2)
            dim1 = rg.CreateDimension("dim1", None, None, 2)
            dim2 = rg.CreateDimension("dim2", None, None, 2)

            with gdaltest.error_handler():
                ar = rg.CreateMDArray(
                    "test",
                    [dim0, dim1, dim2],
                    gdal.ExtendedDataType.Create(gdal.GDT_Byte),
                    ["BLOCKSIZE=" + blocksize],
                )
                assert ar is None

        create()

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_filters():

    filename = "data/zarr/delta_filter_i4.zarr"
    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    rg = ds.GetRootGroup()
    assert rg
    ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
    assert ar
    assert ar.Read() == array.array("i", [i for i in range(10)])


def test_zarr_update_with_filters():

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer(
            "/vsimem/test.zarr/.zarray",
            open("data/zarr/delta_filter_i4.zarr/.zarray", "rb").read(),
        )
        gdal.FileFromMemBuffer(
            "/vsimem/test.zarr/0", open("data/zarr/delta_filter_i4.zarr/0", "rb").read()
        )

        def update():
            ds = gdal.OpenEx(
                "/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE
            )
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
            assert ar
            assert ar.Read() == array.array("i", [i for i in range(10)])
            assert (
                ar.Write(array.array("i", [10 - i for i in range(10)])) == gdal.CE_None
            )

        update()

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE)
        assert ds
        rg = ds.GetRootGroup()
        assert rg
        ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
        assert ar
        assert ar.Read() == array.array("i", [10 - i for i in range(10)])

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_create_with_filter():

    tst = gdaltest.GDALTest(
        "Zarr", "../../gcore/data/uint16.tif", 1, 4672, options=["FILTER=delta"]
    )

    try:
        ret = tst.testCreate(vsimem=1, new_filename="/vsimem/test.zarr")

        f = gdal.VSIFOpenL("/vsimem/test.zarr/test/.zarray", "rb")
        assert f
        data = gdal.VSIFReadL(1, 10000, f)
        gdal.VSIFCloseL(f)
        j = json.loads(data)
        assert "filters" in j
        assert j["filters"] == [{"id": "delta", "dtype": "<u2"}]

        return ret
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_pam_spatial_ref():

    try:

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                "/vsimem/test.zarr"
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg

            dim0 = rg.CreateDimension("dim0", None, None, 2)
            dim1 = rg.CreateDimension("dim1", None, None, 2)
            rg.CreateMDArray(
                "test", [dim0, dim1], gdal.ExtendedDataType.Create(gdal.GDT_Byte)
            )

        create()

        assert gdal.VSIStatL("/vsimem/test.zarr/pam.aux.xml") is None

        def check_crs_before():
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
            assert ar
            crs = ar.GetSpatialRef()
            assert crs is None

        check_crs_before()

        assert gdal.VSIStatL("/vsimem/test.zarr/pam.aux.xml") is None

        def set_crs():
            # Open in read-only
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
            assert ar
            crs = osr.SpatialReference()
            crs.ImportFromEPSG(4326)
            crs.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)
            crs.SetCoordinateEpoch(2021.2)
            assert ar.SetSpatialRef(crs) == gdal.CE_None

        set_crs()

        assert gdal.VSIStatL("/vsimem/test.zarr/pam.aux.xml") is not None

        f = gdal.VSIFOpenL("/vsimem/test.zarr/pam.aux.xml", "rb+")
        assert f
        data = gdal.VSIFReadL(1, 1000, f).decode("utf-8")
        assert data.endswith("</PAMDataset>\n")
        data = data[0 : -len("</PAMDataset>\n")] + "<Other/>" + "</PAMDataset>\n"
        gdal.VSIFSeekL(f, 0, 0)
        gdal.VSIFWriteL(data, 1, len(data), f)
        gdal.VSIFCloseL(f)

        def check_crs():
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
            assert ar
            crs = ar.GetSpatialRef()
            assert crs is not None
            assert crs.GetAuthorityCode(None) == "4326"
            assert crs.GetDataAxisToSRSAxisMapping() == [2, 1]
            assert crs.GetCoordinateEpoch() == 2021.2

        check_crs()

        def check_crs_classic_dataset():
            ds = gdal.Open("/vsimem/test.zarr")
            crs = ds.GetSpatialRef()
            assert crs is not None

        check_crs_classic_dataset()

        def unset_crs():
            # Open in read-only
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
            assert ar
            assert ar.SetSpatialRef(None) == gdal.CE_None

        unset_crs()

        f = gdal.VSIFOpenL("/vsimem/test.zarr/pam.aux.xml", "rb")
        assert f
        data = gdal.VSIFReadL(1, 1000, f).decode("utf-8")
        gdal.VSIFCloseL(f)
        assert "<Other />" in data

        def check_unset_crs():
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            assert ds
            rg = ds.GetRootGroup()
            assert rg
            ar = rg.OpenMDArray(rg.GetMDArrayNames()[0])
            assert ar
            crs = ar.GetSpatialRef()
            assert crs is None

        check_unset_crs()

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_too_large_tile_size():

    j = {
        "chunks": [1000000, 2000],
        "compressor": None,
        "dtype": "!b1",
        "fill_value": None,
        "filters": None,
        "order": "C",
        "shape": [5, 4],
        "zarr_format": 2,
    }

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))
        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        assert ds is not None
        with gdaltest.error_handler():
            assert ds.GetRootGroup().OpenMDArray("test").Read() is None
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_recursive_array_loading():

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)

        j = {"zarr_format": 2}
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zgroup", json.dumps(j))

        j = {
            "chunks": [1],
            "compressor": None,
            "dtype": "!b1",
            "fill_value": None,
            "filters": None,
            "order": "C",
            "shape": [1],
            "zarr_format": 2,
        }
        gdal.FileFromMemBuffer("/vsimem/test.zarr/a/.zarray", json.dumps(j))
        gdal.FileFromMemBuffer("/vsimem/test.zarr/b/.zarray", json.dumps(j))

        j = {"_ARRAY_DIMENSIONS": ["b"]}
        gdal.FileFromMemBuffer("/vsimem/test.zarr/a/.zattrs", json.dumps(j))

        j = {"_ARRAY_DIMENSIONS": ["a"]}
        gdal.FileFromMemBuffer("/vsimem/test.zarr/b/.zattrs", json.dumps(j))

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        assert ds is not None
        with gdaltest.error_handler():
            ar = ds.GetRootGroup().OpenMDArray("a")
            assert ar
            assert (
                gdal.GetLastErrorMsg()
                == "Attempt at recursively loading /vsimem/test.zarr/a/.zarray"
            )
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_too_deep_array_loading():

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)

        j = {"zarr_format": 2}
        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zgroup", json.dumps(j))

        j = {
            "chunks": [1],
            "compressor": None,
            "dtype": "!b1",
            "fill_value": None,
            "filters": None,
            "order": "C",
            "shape": [1],
            "zarr_format": 2,
        }

        N = 33
        for i in range(N):
            gdal.FileFromMemBuffer("/vsimem/test.zarr/%d/.zarray" % i, json.dumps(j))

        for i in range(N - 1):
            j = {"_ARRAY_DIMENSIONS": ["%d" % (i + 1)]}
            gdal.FileFromMemBuffer("/vsimem/test.zarr/%d/.zattrs" % i, json.dumps(j))

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        assert ds is not None
        with gdaltest.error_handler():
            ar = ds.GetRootGroup().OpenMDArray("0")
            assert ar
            assert gdal.GetLastErrorMsg() == "Too deep call stack in LoadArray()"
    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


@pytest.mark.parametrize(
    "filename,path",
    [
        ("data/zarr/nczarr_v2.zarr", "/MyGroup/Group_A"),
        ("data/zarr/nczarr_v2.zarr/MyGroup", "/Group_A"),
        ("data/zarr/nczarr_v2.zarr/MyGroup/Group_A", ""),
        ("data/zarr/nczarr_v2.zarr/MyGroup/Group_A/dset2", None),
    ],
)
def test_zarr_read_nczarr_v2(filename, path):

    with gdaltest.error_handler():
        assert gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE) is None

    ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
    assert ds is not None
    rg = ds.GetRootGroup()

    ar = rg.OpenMDArrayFromFullname((path if path else "") + "/dset2")
    assert ar
    dims = ar.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetSize() == 3
    assert dims[0].GetName() == "lat"
    assert dims[0].GetFullName() == "/MyGroup/lat"
    assert dims[0].GetIndexingVariable() is not None
    assert dims[0].GetIndexingVariable().GetName() == "lat"
    assert dims[0].GetType() == gdal.DIM_TYPE_HORIZONTAL_Y
    assert dims[0].GetDirection() == "NORTH"

    assert dims[1].GetSize() == 3
    assert dims[1].GetName() == "lon"
    assert dims[1].GetFullName() == "/MyGroup/lon"
    assert dims[1].GetIndexingVariable() is not None
    assert dims[1].GetIndexingVariable().GetName() == "lon"
    assert dims[1].GetType() == gdal.DIM_TYPE_HORIZONTAL_X
    assert dims[1].GetDirection() == "EAST"

    if path:
        ar = rg.OpenMDArrayFromFullname(path + "/dset3")
        assert ar
        dims = ar.GetDimensions()
        assert len(dims) == 2
        assert dims[0].GetSize() == 2
        assert dims[0].GetName() == "lat"
        assert dims[0].GetFullName() == "/MyGroup/Group_A/lat"

        assert dims[1].GetSize() == 2
        assert dims[1].GetName() == "lon"
        assert dims[1].GetFullName() == "/MyGroup/Group_A/lon"

    if filename == "data/zarr/nczarr_v2.zarr":
        mygroup = rg.OpenGroup("MyGroup")
        assert mygroup.GetMDArrayNames() == ["lon", "lat", "dset1"]


@pytest.mark.parametrize("format", ["ZARR_V2", "ZARR_V3"])
def test_zarr_cache_tile_presence(format):

    if gdal.GetDriverByName("netCDF") is None:
        pytest.skip("netCDF driver missing")

    filename = "tmp/test.zarr"
    try:
        # Create a Zarr array with sparse tiles
        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(
                filename, options=["FORMAT=" + format]
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg

            dim0 = rg.CreateDimension("dim0", None, None, 2)
            dim1 = rg.CreateDimension("dim1", None, None, 5)
            ar = rg.CreateMDArray(
                "test",
                [dim0, dim1],
                gdal.ExtendedDataType.Create(gdal.GDT_Byte),
                ["BLOCKSIZE=1,2"],
            )
            assert ar
            assert (
                ar.Write(struct.pack("B" * 1, 10), array_start_idx=[0, 0], count=[1, 1])
                == gdal.CE_None
            )
            assert (
                ar.Write(
                    struct.pack("B" * 1, 100), array_start_idx=[1, 3], count=[1, 1]
                )
                == gdal.CE_None
            )

        create()

        # Create the tile presence cache
        def open_with_cache_tile_presence_option():
            ds = gdal.OpenEx(
                filename,
                gdal.OF_MULTIDIM_RASTER,
                open_options=["CACHE_TILE_PRESENCE=YES"],
            )
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg.OpenMDArray("test") is not None

        open_with_cache_tile_presence_option()

        # Check that the cache exists
        if format == "ZARR_V2":
            cache_filename = filename + "/test/.zarray.gmac"
        else:
            cache_filename = filename + "/meta/root/test.array.json.gmac"
        assert gdal.VSIStatL(cache_filename) is not None

        # Read content of the array
        def read_content():
            ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
            assert ds is not None
            rg = ds.GetRootGroup()
            ar = rg.OpenMDArray("test")
            assert ar is not None
            assert struct.unpack("B" * 2 * 5, ar.Read()) == (
                10,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                100,
                0,
            )

        read_content()

        # again
        open_with_cache_tile_presence_option()

        read_content()

        # Now alter the cache to mark a present tile as missing
        def alter_cache():
            ds = gdal.OpenEx(cache_filename, gdal.OF_MULTIDIM_RASTER | gdal.OF_UPDATE)
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg.GetMDArrayNames() == ["_test_tile_presence"]
            ar = rg.OpenMDArray("_test_tile_presence")
            assert struct.unpack("B" * 2 * 3, ar.Read()) == (1, 0, 0, 0, 1, 0)
            assert (
                ar.Write(struct.pack("B" * 1, 0), array_start_idx=[1, 1], count=[1, 1])
                == gdal.CE_None
            )

        alter_cache()

        # Check that reading the array reflects the above modification
        def read_content_altered():
            ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
            assert ds is not None
            rg = ds.GetRootGroup()
            ar = rg.OpenMDArray("test")
            assert ar is not None
            assert struct.unpack("B" * 2 * 5, ar.Read()) == (
                10,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
            )

        read_content_altered()

    finally:
        gdal.RmdirRecursive(filename)


@pytest.mark.parametrize("compression", ["NONE", "ZLIB"])
def test_zarr_advise_read(compression):

    filename = "tmp/test.zarr"
    try:
        dim0_size = 1230
        dim1_size = 2570
        dim0_blocksize = 20
        dim1_blocksize = 30
        data_ar = [(i % 256) for i in range(dim0_size * dim1_size)]

        # Create empty block
        y_offset = dim0_blocksize
        x_offset = dim1_blocksize
        for y in range(dim0_blocksize):
            for x in range(dim1_blocksize):
                data_ar[dim1_size * (y + y_offset) + x + x_offset] = 0

        data = array.array("B", data_ar)

        def create():
            ds = gdal.GetDriverByName("ZARR").CreateMultiDimensional(filename)
            assert ds is not None
            rg = ds.GetRootGroup()
            assert rg

            dim0 = rg.CreateDimension("dim0", None, None, dim0_size)
            dim1 = rg.CreateDimension("dim1", None, None, dim1_size)
            ar = rg.CreateMDArray(
                "test",
                [dim0, dim1],
                gdal.ExtendedDataType.Create(gdal.GDT_Byte),
                [
                    "COMPRESS=" + compression,
                    "BLOCKSIZE=%d,%d" % (dim0_blocksize, dim1_blocksize),
                ],
            )
            assert ar
            ar.SetNoDataValueDouble(0)
            assert ar.Write(data) == gdal.CE_None

        create()

        def read():
            ds = gdal.OpenEx(filename, gdal.OF_MULTIDIM_RASTER)
            assert ds is not None
            rg = ds.GetRootGroup()
            ar = rg.OpenMDArray("test")

            with gdaltest.error_handler():
                assert ar.AdviseRead(options=["CACHE_SIZE=1"]) == gdal.CE_Failure

            got_data_before_advise_read = ar.Read(
                array_start_idx=[40, 51], count=[2 * dim0_blocksize, 2 * dim1_blocksize]
            )

            assert ar.AdviseRead() == gdal.CE_None
            assert ar.Read() == data

            assert (
                ar.AdviseRead(
                    array_start_idx=[40, 51], count=[2 * dim0_blocksize, dim1_blocksize]
                )
                == gdal.CE_None
            )
            # Read more than AdviseRead() window
            got_data = ar.Read(
                array_start_idx=[40, 51], count=[2 * dim0_blocksize, 2 * dim1_blocksize]
            )
            assert got_data == got_data_before_advise_read

        read()

    finally:
        gdal.RmdirRecursive(filename)


def test_zarr_read_invalid_nczarr_dim():

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)

        j = {
            "chunks": [1, 1],
            "compressor": None,
            "dtype": "!b1",
            "fill_value": None,
            "filters": None,
            "order": "C",
            "shape": [1, 1],
            "zarr_format": 2,
            "_NCZARR_ARRAY": {"dimrefs": ["/MyGroup/lon", "/OtherGroup/lat"]},
        }

        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))

        j = {
            "chunks": [1],
            "compressor": None,
            "dtype": "!b1",
            "fill_value": None,
            "filters": None,
            "order": "C",
            "shape": [1],
            "zarr_format": 2,
        }

        gdal.FileFromMemBuffer("/vsimem/test.zarr/MyGroup/lon/.zarray", json.dumps(j))

        j = {"_NCZARR_GROUP": {"dims": {"lon": 0}}}

        gdal.FileFromMemBuffer("/vsimem/test.zarr/MyGroup/.zgroup", json.dumps(j))

        j = {
            "chunks": [2],
            "compressor": None,
            "dtype": "!b1",
            "fill_value": None,
            "filters": None,
            "order": "C",
            "shape": [2],
            "zarr_format": 2,
        }

        gdal.FileFromMemBuffer(
            "/vsimem/test.zarr/OtherGroup/lat/.zarray", json.dumps(j)
        )

        j = {"_NCZARR_GROUP": {"dims": {"lat": 2, "invalid.name": 2}}}

        gdal.FileFromMemBuffer("/vsimem/test.zarr/OtherGroup/.zgroup", json.dumps(j))

        with gdaltest.error_handler():
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            assert ds
            rg = ds.GetRootGroup()
            ar = rg.OpenMDArray("test")
            assert ar

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_nczar_repeated_array_names():

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)

        j = {
            "_NCZARR_GROUP": {
                "dims": {"lon": 1},
                "vars": ["a", "a", "lon", "lon"],
                "groups": ["g", "g"],
            }
        }

        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zgroup", json.dumps(j))

        j = {
            "chunks": [1, 1],
            "compressor": None,
            "dtype": "!b1",
            "fill_value": None,
            "filters": None,
            "order": "C",
            "shape": [1, 1],
            "zarr_format": 2,
        }

        gdal.FileFromMemBuffer("/vsimem/test.zarr/a/.zarray", json.dumps(j))

        j = {
            "chunks": [1],
            "compressor": None,
            "dtype": "!b1",
            "fill_value": None,
            "filters": None,
            "order": "C",
            "shape": [1],
            "zarr_format": 2,
        }

        gdal.FileFromMemBuffer("/vsimem/test.zarr/lon/.zarray", json.dumps(j))

        with gdaltest.error_handler():
            ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
            assert ds
            rg = ds.GetRootGroup()
            assert rg.GetMDArrayNames() == ["lon", "a"]
            ar = rg.OpenMDArray("a")
            assert ar
            assert rg.GetGroupNames() == ["g"]

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_test_overflow_in_AllocateWorkingBuffers_due_to_fortran():

    if sys.maxsize < (1 << 32):
        pytest.skip()

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)

        j = {
            "chunks": [(1 << 32) - 1, (1 << 32) - 1],
            "compressor": None,
            "dtype": "!b1",
            "fill_value": None,
            "filters": None,
            "order": "F",
            "shape": [1, 1],
            "zarr_format": 2,
        }

        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        assert ds
        rg = ds.GetRootGroup()
        ar = rg.OpenMDArray("test")
        with gdaltest.error_handler():
            assert ar.Read(count=[1, 1]) is None

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_test_overflow_in_AllocateWorkingBuffers_due_to_type_change():

    if sys.maxsize < (1 << 32):
        pytest.skip()

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)

        j = {
            "chunks": [(1 << 32) - 1, ((1 << 32) - 1) / 8],
            "compressor": None,
            "dtype": "<u8",
            "fill_value": None,
            "filters": None,
            "order": "C",
            "shape": [1, 1],
            "zarr_format": 2,
        }

        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        assert ds
        rg = ds.GetRootGroup()
        ar = rg.OpenMDArray("test")
        with gdaltest.error_handler():
            assert ar.Read(count=[1, 1]) is None

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")


def test_zarr_read_do_not_crash_on_invalid_byteswap_on_ascii_string():

    try:
        gdal.Mkdir("/vsimem/test.zarr", 0)

        j = {
            "chunks": [1],
            "compressor": None,
            "dtype": [["x", ">S2"]],  # byteswap here is not really valid...
            "fill_value": base64.b64encode(b"XX").decode("utf-8"),
            "filters": None,
            "order": "C",
            "shape": [1],
            "zarr_format": 2,
        }

        gdal.FileFromMemBuffer("/vsimem/test.zarr/.zarray", json.dumps(j))

        ds = gdal.OpenEx("/vsimem/test.zarr", gdal.OF_MULTIDIM_RASTER)
        assert ds
        rg = ds.GetRootGroup()
        rg.OpenMDArray("test")

    finally:
        gdal.RmdirRecursive("/vsimem/test.zarr")
