from __future__ import annotations

import hypothesis as h
import hypothesis.strategies as st
import pytest
from packaging.version import parse as vparse

import ibis.expr.datatypes as dt
import ibis.tests.strategies as ibst

np = pytest.importorskip("numpy")
npst = pytest.importorskip("hypothesis.extra.numpy")

from ibis.formats.numpy import NumpySchema, NumpyType  # noqa: E402

roundtripable_types = st.deferred(
    lambda: (
        npst.integer_dtypes(endianness="=")
        | npst.floating_dtypes(endianness="=")
        | npst.datetime64_dtypes(max_period="ns", endianness="=")
        | npst.timedelta64_dtypes(max_period="s", endianness="=")
    )
)


@st.composite
def numpy_schema(draw, item_strategy=roundtripable_types, max_size=10):
    num_fields = draw(st.integers(min_value=0, max_value=max_size))
    names = draw(
        st.lists(st.text(), min_size=num_fields, max_size=num_fields, unique=True)
    )
    types = draw(st.lists(item_strategy, min_size=num_fields, max_size=num_fields))
    return list(zip(names, types))


def assert_dtype_roundtrip(
    numpy_type, ibis_type=None, restored_type=None, nullable=True
):
    dtype = NumpyType.to_ibis(numpy_type, nullable=nullable)
    if ibis_type is not None:
        assert dtype == ibis_type

    nptyp = NumpyType.from_ibis(dtype)
    if restored_type is None:
        restored_type = numpy_type
    assert nptyp == restored_type


@h.given(roundtripable_types)
def test_roundtripable_types(numpy_type):
    assert_dtype_roundtrip(numpy_type, nullable=False)
    assert_dtype_roundtrip(numpy_type, nullable=True)


@h.given(npst.unicode_string_dtypes(endianness="="))
def test_non_roundtripable_str_type(numpy_type):
    assert_dtype_roundtrip(
        numpy_type, dt.String(nullable=False), np.dtype("object"), nullable=False
    )
    assert_dtype_roundtrip(
        numpy_type, dt.String(nullable=True), np.dtype("object"), nullable=True
    )


@h.given(npst.byte_string_dtypes(endianness="="))
def test_non_roundtripable_bytes_type(numpy_type):
    assert_dtype_roundtrip(
        numpy_type, dt.Binary(nullable=False), np.dtype("object"), nullable=False
    )
    assert_dtype_roundtrip(
        numpy_type, dt.Binary(nullable=True), np.dtype("object"), nullable=True
    )


@h.given(
    ibst.null_dtype
    | ibst.variadic_dtypes()
    | ibst.decimal_dtypes()
    | ibst.struct_dtypes()
)
def test_variadic_to_numpy(ibis_type):
    assert NumpyType.from_ibis(ibis_type) == np.dtype("object")


@h.given(ibst.timestamp_dtype())
def test_timestamp_to_numpy(ibis_type):
    assert NumpyType.from_ibis(ibis_type) == np.dtype("datetime64[ns]")


@h.given(ibst.date_dtype())
def test_date_to_numpy(ibis_type):
    assert NumpyType.from_ibis(ibis_type) == np.dtype("datetime64[D]")


@h.given(ibst.time_dtype())
def test_time_to_numpy(ibis_type):
    assert NumpyType.from_ibis(ibis_type) == np.dtype("timedelta64[ns]")


@h.given(ibst.schema())
def test_schema_to_numpy(ibis_schema):
    numpy_schema = NumpySchema.from_ibis(ibis_schema)
    assert len(numpy_schema) == len(ibis_schema)

    for name, numpy_type in numpy_schema:
        assert numpy_type == NumpyType.from_ibis(ibis_schema[name])


@h.given(numpy_schema())
def test_schema_from_numpy(numpy_schema):
    ibis_schema = NumpySchema.to_ibis(numpy_schema)
    assert len(numpy_schema) == len(ibis_schema)

    for name, numpy_type in numpy_schema:
        assert NumpyType.from_ibis(ibis_schema[name]) == numpy_type


@pytest.mark.parametrize(
    ("numpy_dtype", "ibis_dtype"),
    [
        (np.bool_, dt.boolean),
        (np.int8, dt.int8),
        (np.int16, dt.int16),
        (np.int32, dt.int32),
        (np.int64, dt.int64),
        (np.uint8, dt.uint8),
        (np.uint16, dt.uint16),
        (np.uint32, dt.uint32),
        (np.uint64, dt.uint64),
        (np.float16, dt.float16),
        (np.float32, dt.float32),
        (np.float64, dt.float64),
        (np.double, dt.double),
        (np.str_, dt.string),
        (np.datetime64, dt.timestamp),
    ],
)
def test_dtype_from_numpy(numpy_dtype, ibis_dtype):
    assert NumpyType.to_ibis(np.dtype(numpy_dtype)) == ibis_dtype


def test_dtype_from_numpy_dtype_timedelta():
    if vparse(pytest.importorskip("pyarrow").__version__) < vparse("9"):
        pytest.skip("pyarrow < 9 globally mutates the timedelta64 numpy dtype")

    assert NumpyType.to_ibis(np.dtype(np.timedelta64)) == dt.Interval(unit="s")
