import base64
import copy
import json

import numpy as np
import pytest

from zarr.codecs import Blosc, Delta, Pickle, Zlib, Zstd
from zarr.errors import MetadataError
from zarr.meta import (
    ZARR_FORMAT,
    decode_array_metadata,
    decode_dtype,
    decode_group_metadata,
    encode_array_metadata,
    encode_dtype,
    encode_fill_value,
    decode_fill_value,
    get_extended_dtype_info,
    _v3_complex_types,
    _v3_datetime_types,
    _default_entry_point_metadata_v3,
    Metadata3,
)
from zarr.util import normalize_dtype, normalize_fill_value


def assert_json_equal(expect, actual):
    if isinstance(actual, bytes):
        actual = str(actual, "ascii")
    ej = json.loads(expect)
    aj = json.loads(actual)
    assert ej == aj


def test_encode_decode_array_1():
    meta = dict(
        shape=(100,),
        chunks=(10,),
        dtype=np.dtype("<f8"),
        compressor=Zlib(1).get_config(),
        fill_value=None,
        filters=None,
        order="C",
    )

    meta_json = f"""{{
        "chunks": [10],
        "compressor": {{"id": "zlib", "level": 1}},
        "dtype": "<f8",
        "fill_value": null,
        "filters": null,
        "order": "C",
        "shape": [100],
        "zarr_format": {ZARR_FORMAT}
    }}"""

    # test encoding
    meta_enc = encode_array_metadata(meta)
    assert_json_equal(meta_json, meta_enc)

    # test decoding
    meta_dec = decode_array_metadata(meta_enc)
    assert ZARR_FORMAT == meta_dec["zarr_format"]
    assert meta["shape"] == meta_dec["shape"]
    assert meta["chunks"] == meta_dec["chunks"]
    assert meta["dtype"] == meta_dec["dtype"]
    assert meta["compressor"] == meta_dec["compressor"]
    assert meta["order"] == meta_dec["order"]
    assert meta_dec["fill_value"] is None
    assert meta_dec["filters"] is None


def test_encode_decode_array_2():
    # some variations
    df = Delta(astype="<u2", dtype="V14")
    compressor = Blosc(cname="lz4", clevel=3, shuffle=2)
    dtype = np.dtype([("a", "<i4"), ("b", "S10")])
    fill_value = np.zeros((), dtype=dtype)[()]
    meta = dict(
        shape=(100, 100),
        chunks=(10, 10),
        dtype=dtype,
        compressor=compressor.get_config(),
        fill_value=fill_value,
        order="F",
        filters=[df.get_config()],
    )

    meta_json = f"""{{
        "chunks": [10, 10],
        "compressor": {{
            "id": "blosc",
            "clevel": 3,
            "cname": "lz4",
            "shuffle": 2,
            "blocksize": 0
        }},
        "dtype": [["a", "<i4"], ["b", "|S10"]],
        "fill_value": "AAAAAAAAAAAAAAAAAAA=",
        "filters": [
            {{"id": "delta", "astype": "<u2", "dtype": "|V14"}}
        ],
        "order": "F",
        "shape": [100, 100],
        "zarr_format": {ZARR_FORMAT}
    }}"""

    # test encoding
    meta_enc = encode_array_metadata(meta)
    assert_json_equal(meta_json, meta_enc)

    # test decoding
    meta_dec = decode_array_metadata(meta_enc)
    assert ZARR_FORMAT == meta_dec["zarr_format"]
    assert meta["shape"] == meta_dec["shape"]
    assert meta["chunks"] == meta_dec["chunks"]
    assert meta["dtype"] == meta_dec["dtype"]
    assert meta["compressor"] == meta_dec["compressor"]
    assert meta["order"] == meta_dec["order"]
    assert fill_value == meta_dec["fill_value"]
    assert [df.get_config()] == meta_dec["filters"]


def test_encode_decode_array_complex():
    # some variations
    for k in ["c8", "c16"]:
        compressor = Blosc(cname="lz4", clevel=3, shuffle=2)
        dtype = np.dtype(k)
        fill_value = dtype.type(np.nan - 1j)
        meta = dict(
            shape=(100, 100),
            chunks=(10, 10),
            dtype=dtype,
            compressor=compressor.get_config(),
            fill_value=fill_value,
            order=dtype.char,
            filters=[],
        )

        meta_json = f"""{{
            "chunks": [10, 10],
            "compressor": {{
                "id": "blosc",
                "clevel": 3,
                "cname": "lz4",
                "shuffle": 2,
                "blocksize": 0
            }},
            "dtype": "{dtype.str}",
            "fill_value": ["NaN", -1.0],
            "filters": [],
            "order": "{dtype.char}",
            "shape": [100, 100],
            "zarr_format": {ZARR_FORMAT}
        }}"""

        # test encoding
        meta_enc = encode_array_metadata(meta)
        assert_json_equal(meta_json, meta_enc)

        # test decoding
        meta_dec = decode_array_metadata(meta_enc)
        assert ZARR_FORMAT == meta_dec["zarr_format"]
        assert meta["shape"] == meta_dec["shape"]
        assert meta["chunks"] == meta_dec["chunks"]
        assert meta["dtype"] == meta_dec["dtype"]
        assert meta["compressor"] == meta_dec["compressor"]
        assert meta["order"] == meta_dec["order"]
        # Based off of this SO answer: https://stackoverflow.com/a/49972198
        assert np.all(
            fill_value.view((np.uint8, fill_value.itemsize))
            == meta_dec["fill_value"].view((np.uint8, meta_dec["fill_value"].itemsize))
        )
        assert [] == meta_dec["filters"]


def test_encode_decode_array_datetime_timedelta():
    # some variations
    for k in ["m8[s]", "M8[s]"]:
        compressor = Blosc(cname="lz4", clevel=3, shuffle=2)
        dtype = np.dtype(k)
        fill_value = dtype.type("NaT")
        meta = dict(
            shape=(100, 100),
            chunks=(10, 10),
            dtype=dtype,
            compressor=compressor.get_config(),
            fill_value=fill_value,
            order=dtype.char,
            filters=[],
        )

        meta_json = f"""{{
            "chunks": [10, 10],
            "compressor": {{
                "id": "blosc",
                "clevel": 3,
                "cname": "lz4",
                "shuffle": 2,
                "blocksize": 0
            }},
            "dtype": "{dtype.str}",
            "fill_value": -9223372036854775808,
            "filters": [],
            "order": "{dtype.char}",
            "shape": [100, 100],
            "zarr_format": {ZARR_FORMAT}
        }}"""

        # test encoding
        meta_enc = encode_array_metadata(meta)
        assert_json_equal(meta_json, meta_enc)

        # test decoding
        meta_dec = decode_array_metadata(meta_enc)
        assert ZARR_FORMAT == meta_dec["zarr_format"]
        assert meta["shape"] == meta_dec["shape"]
        assert meta["chunks"] == meta_dec["chunks"]
        assert meta["dtype"] == meta_dec["dtype"]
        assert meta["compressor"] == meta_dec["compressor"]
        assert meta["order"] == meta_dec["order"]
        # Based off of this SO answer: https://stackoverflow.com/a/49972198
        assert np.all(
            fill_value.view((np.uint8, fill_value.itemsize))
            == meta_dec["fill_value"].view((np.uint8, meta_dec["fill_value"].itemsize))
        )
        assert [] == meta_dec["filters"]


def test_encode_decode_array_dtype_shape():
    meta = dict(
        shape=(100,),
        chunks=(10,),
        dtype=np.dtype("(10, 10)<f8"),
        compressor=Zlib(1).get_config(),
        fill_value=None,
        filters=None,
        order="C",
    )

    meta_json = f"""{{
        "chunks": [10],
        "compressor": {{"id": "zlib", "level": 1}},
        "dtype": "<f8",
        "fill_value": null,
        "filters": null,
        "order": "C",
        "shape": [100, 10, 10],
        "zarr_format": {ZARR_FORMAT}
    }}"""

    # test encoding
    meta_enc = encode_array_metadata(meta)
    assert_json_equal(meta_json, meta_enc)

    # test decoding
    meta_dec = decode_array_metadata(meta_enc)
    assert ZARR_FORMAT == meta_dec["zarr_format"]
    # to maintain consistency with numpy unstructured arrays, unpack dimensions into shape
    assert meta["shape"] + meta["dtype"].shape == meta_dec["shape"]
    assert meta["chunks"] == meta_dec["chunks"]
    # to maintain consistency with numpy unstructured arrays, unpack dtypes
    assert meta["dtype"].base == meta_dec["dtype"]
    assert meta["compressor"] == meta_dec["compressor"]
    assert meta["order"] == meta_dec["order"]
    assert meta_dec["fill_value"] is None
    assert meta_dec["filters"] is None


@pytest.mark.parametrize("cname", ["zlib", "zstd"])
def test_encode_decode_array_dtype_shape_v3(cname):
    if cname == "zlib":
        compressor = Zlib(1)
    elif cname == "zstd":
        compressor = Zstd(1)
    meta = dict(
        shape=(100,),
        chunk_grid=dict(type="regular", chunk_shape=(10,), separator=("/")),
        data_type=np.dtype("(10, 10)<f8"),
        compressor=compressor,
        fill_value=None,
        chunk_memory_layout="C",
    )
    meta_expected = {
        "attributes": {},
        "chunk_grid": {"chunk_shape": [10], "separator": "/", "type": "regular"},
        "chunk_memory_layout": "C",
        "compressor": {
            "codec": f"https://purl.org/zarr/spec/codec/{cname}/1.0",
            "configuration": {"level": 1},
        },
        "data_type": "<f8",
        "extensions": [],
        "fill_value": None,
        "shape": [100, 10, 10],
    }

    if cname == "zstd":
        meta_expected["compressor"]["configuration"]["checksum"] = False

    meta_expected_json = json.dumps(meta_expected)

    # test encoding
    meta_enc = Metadata3.encode_array_metadata(meta)
    assert_json_equal(meta_enc, meta_expected_json)

    # test decoding
    meta_dec = Metadata3.decode_array_metadata(meta_enc)
    # to maintain consistency with numpy unstructured arrays, unpack dimensions into shape
    assert meta["shape"] + meta["data_type"].shape == meta_dec["shape"]
    assert meta["chunk_grid"] == meta_dec["chunk_grid"]
    # to maintain consistency with numpy unstructured arrays, unpack dtypes
    assert meta["data_type"].base == meta_dec["data_type"]
    assert meta["compressor"] == meta_dec["compressor"]
    assert meta["chunk_memory_layout"] == meta_dec["chunk_memory_layout"]
    assert meta_dec["fill_value"] is None
    assert "filters" not in meta_dec


@pytest.mark.parametrize("comp_id", ["gzip", "zlib", "blosc", "bz2", "lz4", "lzma", "zstd"])
def test_decode_metadata_implicit_compressor_config_v3(comp_id):
    meta = {
        "attributes": {},
        "chunk_grid": {"chunk_shape": [10], "separator": "/", "type": "regular"},
        "chunk_memory_layout": "C",
        "compressor": {
            "codec": f"https://purl.org/zarr/spec/codec/{comp_id}/1.0",
            "configuration": {
                # intentionally left empty
            },
        },
        "data_type": "<f8",
        "extensions": [],
        "fill_value": None,
        "shape": [100, 10, 10],
    }
    meta_json = json.dumps(meta)

    # test decoding
    meta_dec = Metadata3.decode_array_metadata(meta_json)
    assert meta_dec["compressor"].codec_id == comp_id


def test_encode_decode_array_structured():
    meta = dict(
        shape=(100,),
        chunks=(10,),
        dtype=np.dtype("<i8, (10, 10)<f8, (5, 10, 15)u1"),
        compressor=Zlib(1).get_config(),
        fill_value=None,
        filters=None,
        order="C",
    )

    meta_json = f"""{{
        "chunks": [10],
        "compressor": {{"id": "zlib", "level": 1}},
        "dtype": [["f0", "<i8"], ["f1", "<f8", [10, 10]], ["f2", "|u1", [5, 10, 15]]],
        "fill_value": null,
        "filters": null,
        "order": "C",
        "shape": [100],
        "zarr_format": {ZARR_FORMAT}
    }}"""

    # test encoding
    meta_enc = encode_array_metadata(meta)
    assert_json_equal(meta_json, meta_enc)

    # test decoding
    meta_dec = decode_array_metadata(meta_enc)
    assert ZARR_FORMAT == meta_dec["zarr_format"]
    # to maintain consistency with numpy unstructured arrays, unpack dimensions into shape
    assert meta["shape"] + meta["dtype"].shape == meta_dec["shape"]
    assert meta["chunks"] == meta_dec["chunks"]
    # to maintain consistency with numpy unstructured arrays, unpack dimensions into shape
    assert meta["dtype"].base == meta_dec["dtype"]
    assert meta["compressor"] == meta_dec["compressor"]
    assert meta["order"] == meta_dec["order"]
    assert meta_dec["fill_value"] is None
    assert meta_dec["filters"] is None


def test_encode_decode_fill_values_nan():
    fills = (
        (np.nan, "NaN", np.isnan),
        (-np.inf, "-Infinity", np.isneginf),
        (np.inf, "Infinity", np.isposinf),
    )

    for v, s, f in fills:
        meta = dict(
            shape=(100,),
            chunks=(10,),
            dtype=np.dtype("<f8"),
            compressor=Zlib(1).get_config(),
            fill_value=v,
            filters=None,
            order="C",
        )

        meta_json = f"""{{
            "chunks": [10],
            "compressor": {{"id": "zlib", "level": 1}},
            "dtype": "<f8",
            "fill_value": "{s}",
            "filters": null,
            "order": "C",
            "shape": [100],
            "zarr_format": {ZARR_FORMAT}
        }}"""

        # test encoding
        meta_enc = encode_array_metadata(meta)
        assert_json_equal(meta_json, meta_enc)

        # test decoding
        meta_dec = decode_array_metadata(meta_enc)
        actual = meta_dec["fill_value"]
        assert f(actual)


def test_encode_decode_fill_values_bytes():
    dtype = np.dtype("S10")
    fills = b"foo", bytes(10)

    for v in fills:
        # setup and encode metadata
        meta = dict(
            shape=(100,),
            chunks=(10,),
            dtype=dtype,
            compressor=Zlib(1).get_config(),
            fill_value=v,
            filters=None,
            order="C",
        )
        meta_enc = encode_array_metadata(meta)

        # define expected metadata encoded as JSON
        s = base64.standard_b64encode(v)
        s = s.decode()
        meta_json = f"""{{
            "chunks": [10],
            "compressor": {{"id": "zlib", "level": 1}},
            "dtype": "|S10",
            "fill_value": "{s}",
            "filters": null,
            "order": "C",
            "shape": [100],
            "zarr_format": {ZARR_FORMAT}
        }}"""

        # test encoding
        assert_json_equal(meta_json, meta_enc)

        # test decoding
        meta_dec = decode_array_metadata(meta_enc)
        actual = meta_dec["fill_value"]
        expect = np.array(v, dtype=dtype)[()]
        assert expect == actual


def test_decode_array_unsupported_format():
    # unsupported format
    meta_json = f"""{{
        "zarr_format": {ZARR_FORMAT - 1},
        "shape": [100],
        "chunks": [10],
        "dtype": "<f8",
        "compressor": {{"id": "zlib", "level": 1}},
        "fill_value": null,
        "order": "C"
    }}"""
    with pytest.raises(MetadataError):
        decode_array_metadata(meta_json)


def test_decode_array_missing_fields():
    # missing fields
    meta_json = f"""{{
        "zarr_format": {ZARR_FORMAT}
    }}"""
    with pytest.raises(MetadataError):
        decode_array_metadata(meta_json)


def test_encode_decode_dtype():
    for dt in ["f8", [("a", "f8")], [("a", "f8"), ("b", "i1")]]:
        e = encode_dtype(np.dtype(dt))
        s = json.dumps(e)  # check JSON serializable
        o = json.loads(s)
        d = decode_dtype(o)
        assert np.dtype(dt) == d


def test_decode_group():
    # typical
    b = f"""{{
        "zarr_format": {ZARR_FORMAT}
    }}"""
    meta = decode_group_metadata(b)
    assert ZARR_FORMAT == meta["zarr_format"]

    # unsupported format
    b = f"""{{
        "zarr_format": {ZARR_FORMAT - 1}
    }}"""
    with pytest.raises(MetadataError):
        decode_group_metadata(b)


@pytest.mark.parametrize(
    "fill_value,dtype,object_codec,result",
    [
        (
            (0.0, None),
            [("x", float), ("y", object)],
            Pickle(),
            True,  # Pass
        ),
        (
            (0.0, None),
            [("x", float), ("y", object)],
            None,
            False,  # Fail
        ),
    ],
)
def test_encode_fill_value(fill_value, dtype, object_codec, result):
    # normalize metadata (copied from _init_array_metadata)
    dtype, object_codec = normalize_dtype(dtype, object_codec)
    dtype = dtype.base
    fill_value = normalize_fill_value(fill_value, dtype)

    # test
    if result:
        encode_fill_value(fill_value, dtype, object_codec)
    else:
        with pytest.raises(ValueError):
            encode_fill_value(fill_value, dtype, object_codec)


@pytest.mark.parametrize(
    "fill_value,dtype,object_codec,result",
    [
        (
            (0.0, None),
            [("x", float), ("y", object)],
            Pickle(),
            True,  # Pass
        ),
        (
            (0.0, None),
            [("x", float), ("y", object)],
            None,
            False,  # Fail
        ),
    ],
)
def test_decode_fill_value(fill_value, dtype, object_codec, result):
    # normalize metadata (copied from _init_array_metadata)
    dtype, object_codec = normalize_dtype(dtype, object_codec)
    dtype = dtype.base
    fill_value = normalize_fill_value(fill_value, dtype)

    # test
    if result:
        v = encode_fill_value(fill_value, dtype, object_codec)
        decode_fill_value(v, dtype, object_codec)
    else:
        with pytest.raises(ValueError):
            # No encoding is possible
            decode_fill_value(fill_value, dtype, object_codec)


def test_get_extended_dtype_info():
    extended_types = list(_v3_complex_types) + list(_v3_datetime_types)
    extended_types += ["|S4", "|S8", ">U4", "<U4", "|O", "|V16"]

    for dtype in extended_types:
        try:
            info = get_extended_dtype_info(np.asarray([], dtype=dtype).dtype)
        except TypeError:  # pragma: no cover
            # skip any numpy dtypes not supported by a particular architecture
            pass  # pragma: no cover
        assert "extension" in info
        assert "type" in info
        assert "fallback" in info

    class invalid_dtype:
        str = "unknown_type"

    with pytest.raises(ValueError):
        get_extended_dtype_info(invalid_dtype)


def test_metadata3_exceptions():
    with pytest.raises(KeyError):
        # dict must have a key named 'type'
        Metadata3.decode_dtype({})

    required = ["zarr_format", "metadata_encoding", "metadata_key_suffix", "extensions"]
    for key in required:
        meta = copy.copy(_default_entry_point_metadata_v3)
        meta.pop(key)
        with pytest.raises(ValueError):
            # cannot encode metadata that is missing a required key
            Metadata3.encode_hierarchy_metadata(meta)

    meta = copy.copy(_default_entry_point_metadata_v3)
    meta["extra_key"] = []
    with pytest.raises(ValueError):
        # cannot encode metadata that has unexpected keys
        Metadata3.encode_hierarchy_metadata(meta)

    json = Metadata3.encode_hierarchy_metadata(_default_entry_point_metadata_v3)
    # form a new json bytes object with one of the keys missing
    temp = json.split(b"\n")
    temp = temp[:2] + temp[3:]
    bad_json = b"\n".join(temp)
    with pytest.raises(ValueError):
        # cannot encode metadata that is missing a required key
        Metadata3.decode_hierarchy_metadata(bad_json)

    json = Metadata3.encode_hierarchy_metadata(_default_entry_point_metadata_v3)
    temp = json.split(b"\n")
    temp = temp[:2] + [b'    "unexpected": [],'] + temp[2:]
    bad_json = b"\n".join(temp)
    with pytest.raises(ValueError):
        # cannot encode metadata that has extra, unexpected keys
        Metadata3.decode_hierarchy_metadata(bad_json)

    codec_meta = dict(configuration=None, codec="unknown")
    with pytest.raises(NotImplementedError):
        Metadata3._decode_codec_metadata(codec_meta)

    with pytest.raises(MetadataError):
        Metadata3.decode_array_metadata(dict())
