import dask.array
import numpy
import orjson
import pytest
import xarray
import xarray.testing

from ..adapters.mapping import MapAdapter
from ..adapters.xarray import DatasetAdapter
from ..client import Context, from_context, record_history
from ..client import xarray as xarray_client
from ..serialization.xarray import serialize_json
from ..server.app import build_app

image = numpy.random.random((3, 5))
temp = 15 + 8 * numpy.random.randn(2, 2, 3)
precip = 10 * numpy.random.rand(2, 2, 3)
lon = [[-99.83, -99.32], [-99.79, -99.23]]
lat = [[42.25, 42.21], [42.63, 42.59]]

EXPECTED = {
    "image": xarray.Dataset(
        {
            "image": xarray.DataArray(
                xarray.Variable(
                    data=dask.array.from_array(image),
                    dims=["x", "y"],
                    attrs={"thing": "stuff"},
                ),
                coords={
                    "x": dask.array.arange(image.shape[0]) / 10,
                    "y": dask.array.arange(image.shape[1]) / 50,
                },
            ),
            "z": xarray.DataArray(data=dask.array.ones(17)),
        },
        attrs={"snow": "cold"},
    ),
    "weather": xarray.Dataset(
        {
            "temperature": (["x", "y", "time"], temp),
            "precipitation": (["x", "y", "time"], precip),
        },
        coords={
            "lon": (["x", "y"], lon),
            "lat": (["x", "y"], lat),
            "time": [1, 2, 3],
        },
    ),
    "wide": xarray.Dataset(
        {f"column_{i:03}": xarray.DataArray(i * numpy.ones(2)) for i in range(10)},
        coords={"time": numpy.arange(2)},
    ),
    "ragged": xarray.Dataset(
        {
            f"{i}": xarray.DataArray(i * numpy.ones(2 * i), dims=f"dim{i}")
            for i in range(3)
        }
    ),
}

tree = MapAdapter(
    {key: DatasetAdapter.from_dataset(ds) for key, ds in EXPECTED.items()}
)


@pytest.fixture(scope="module")
def client():
    app = build_app(tree)
    with Context.from_app(app) as context:
        client = from_context(context)
        yield client


@pytest.mark.parametrize("key", list(tree))
def test_xarray_dataset(client, key):
    expected = EXPECTED[key]
    actual = client[key].read().load()
    xarray.testing.assert_equal(actual, expected)


@pytest.mark.parametrize("key", ["image", "weather"])
def test_dataset_column_access(client, key):
    expected_dataset = EXPECTED[key]
    actual_dataset = client[key].read().load()
    for col in expected_dataset:
        actual = actual_dataset[col]
        expected = expected_dataset[col]
        xarray.testing.assert_equal(actual, expected)


def test_wide_table_optimization(client):
    wide = client["wide"]
    with record_history() as history:
        wide.read()
    # This should be just a couple requests.
    # This upper bound is somewhat arbitrary to give wiggle room for future
    # minor changes. The point is: it's much less than one request per variable.
    assert len(history.requests) < 4


def test_wide_table_optimization_off(client):
    wide = client["wide"]
    with record_history() as history:
        wide.read(optimize_wide_table=False)
    assert len(history.requests) >= 10


def test_url_limit_handling(client):
    "Check that requests and split up to stay below the URL length limit."
    expected = EXPECTED["wide"]
    dsc = client["wide"]
    dsc.read()  # Dry run to run any one-off state-initializing requests.
    # Accumulate Requests here for later inspection.
    requests = []

    def accumulate(request):
        # httpx.AsyncClient requires event hooks to be async functions.
        requests.append(request)

    client.context.http_client.event_hooks["request"].append(accumulate)
    actual = dsc.read()
    xarray.testing.assert_equal(actual, expected)
    normal_request_count = len(requests)
    original = xarray_client.URL_CHARACTER_LIMIT
    try:
        # It should never be necessary to tune this for real-world use, but we
        # use this knob as a way to test its operation.
        xarray_client.URL_CHARACTER_LIMIT = 200
        # The client will need to split this across more requests in order to
        # stay within the tighter limit.
        requests.clear()  # Empty the Request cache before the next batch of requests.
        actual = dsc.read()
        xarray.testing.assert_equal(actual, expected)
        higher_request_count = len(requests)
        # Tighten even more.
        xarray_client.URL_CHARACTER_LIMIT = 100
        requests.clear()  # Empty the Request cache before the next batch of requests.
        actual = dsc.read()
        xarray.testing.assert_equal(actual, expected)
        highest_request_count = len(requests)
    finally:
        # Restore default.
        xarray_client.URL_CHARACTER_LIMIT = original
    # The goal here is to test the *trend* not the specific values because the
    # number of requests may evolve as the library changes, but the trend should
    # hold.
    assert highest_request_count > higher_request_count > normal_request_count


@pytest.mark.parametrize("ds_node", tree.values(), ids=tree.keys())
@pytest.mark.asyncio
async def test_serialize_json(ds_node: DatasetAdapter):
    """Verify that serialized Dataset keys are a subset
    of all coordinates and variables from the Dataset.
    Index variables are removed by serialize_json().
    """
    metadata = None  # Not used
    filter_for_access = None  # Not used
    result = await serialize_json(ds_node, metadata, filter_for_access)

    result_data_keys = orjson.loads(result).keys()
    ds_coords_and_vars = set(ds_node)

    assert set(result_data_keys).issubset(ds_coords_and_vars)
