import base64
import json

import numpy as np
import pytest

from zarr.v2.codecs import Blosc, Delta, Pickle, Zlib
from zarr.v2.errors import MetadataError
from zarr.v2.meta import (
    ZARR_FORMAT,
    decode_array_metadata,
    decode_dtype,
    decode_group_metadata,
    encode_array_metadata,
    encode_dtype,
    encode_fill_value,
    decode_fill_value,
)
from zarr.v2.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 = (
        """{
        "chunks": [10],
        "compressor": {"id": "zlib", "level": 1},
        "dtype": "<f8",
        "fill_value": null,
        "filters": null,
        "order": "C",
        "shape": [100],
        "zarr_format": %s
    }"""
        % 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 = (
        """{
        "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": %s
    }"""
        % 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 = """{
            "chunks": [10, 10],
            "compressor": {
                "id": "blosc",
                "clevel": 3,
                "cname": "lz4",
                "shuffle": 2,
                "blocksize": 0
            },
            "dtype": "%s",
            "fill_value": ["NaN", -1.0],
            "filters": [],
            "order": "%s",
            "shape": [100, 100],
            "zarr_format": %s
        }""" % (
            dtype.str,
            dtype.char,
            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 = """{
            "chunks": [10, 10],
            "compressor": {
                "id": "blosc",
                "clevel": 3,
                "cname": "lz4",
                "shuffle": 2,
                "blocksize": 0
            },
            "dtype": "%s",
            "fill_value": -9223372036854775808,
            "filters": [],
            "order": "%s",
            "shape": [100, 100],
            "zarr_format": %s
        }""" % (
            dtype.str,
            dtype.char,
            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 = (
        """{
        "chunks": [10],
        "compressor": {"id": "zlib", "level": 1},
        "dtype": "<f8",
        "fill_value": null,
        "filters": null,
        "order": "C",
        "shape": [100, 10, 10],
        "zarr_format": %s
    }"""
        % 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


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 = (
        """{
        "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": %s
    }"""
        % 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 = """{
            "chunks": [10],
            "compressor": {"id": "zlib", "level": 1},
            "dtype": "<f8",
            "fill_value": "%s",
            "filters": null,
            "order": "C",
            "shape": [100],
            "zarr_format": %s
        }""" % (
            s,
            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 = """{
            "chunks": [10],
            "compressor": {"id": "zlib", "level": 1},
            "dtype": "|S10",
            "fill_value": "%s",
            "filters": null,
            "order": "C",
            "shape": [100],
            "zarr_format": %s
        }""" % (
            s,
            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 = """{
        "zarr_format": %s,
        "shape": [100],
        "chunks": [10],
        "dtype": "<f8",
        "compressor": {"id": "zlib", "level": 1},
        "fill_value": null,
        "order": "C"
    }""" % (ZARR_FORMAT - 1)
    with pytest.raises(MetadataError):
        decode_array_metadata(meta_json)


def test_decode_array_missing_fields():
    # missing fields
    meta_json = (
        """{
        "zarr_format": %s
    }"""
        % 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 = (
        """{
        "zarr_format": %s
    }"""
        % ZARR_FORMAT
    )
    meta = decode_group_metadata(b)
    assert ZARR_FORMAT == meta["zarr_format"]

    # unsupported format
    b = """{
        "zarr_format": %s
    }""" % (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)
