"""
Tests for our custom array writing locking stuff.
"""

import logging
import sys
from time import sleep

import dask
import numpy as np
import pytest
import zarr
from dask.distributed import Client

from pangeo_forge.utils import chunk_bounds_and_conflicts, lock_for_conflicts


@pytest.mark.parametrize("n_tasks, conflicts", [(2, (0,)), (2, (1, 2)), (4, (0,))])
def test_locks(n_tasks, conflicts, tmp_target, dask_cluster):

    this_client = Client(dask_cluster)

    @dask.delayed
    def do_stuff(n):
        logger = logging.getLogger("pangeo_forge")
        handler = logging.StreamHandler(stream=sys.stdout)
        handler.setLevel(logging.DEBUG)
        logger.addHandler(handler)
        with lock_for_conflicts(conflicts):
            # test for concurrency by writing an object and asserting there are
            # no other objects present
            mapper = tmp_target.get_mapper()
            key = f"foo_{n}"
            mapper[key] = b"bar"
            assert list(mapper) == [key]
            sleep(0.25)
            del mapper[key]
            assert list(mapper) == []

    dask.compute([do_stuff(n) for n in range(n_tasks)])

    this_client.close()
    del this_client


# chunking is always over first dimension
@pytest.mark.parametrize(
    "shape, zarr_chunks, write_chunks",
    [
        ((100,), (10,), (15,)),  # just a few conflicts
        ((100,), (100,), (10,)),  # lots of conflicts!
    ],
)
def test_locked_array_writing(shape, zarr_chunks, write_chunks, tmp_target, dask_cluster):

    sequence_lens = (shape[0] // write_chunks[0]) * [write_chunks[0]]
    remainder = shape[0] % write_chunks[0]
    if remainder > 0:
        sequence_lens.append(remainder)

    chunk_bounds, chunk_conflicts = chunk_bounds_and_conflicts(sequence_lens, zarr_chunks[0])

    data = np.random.rand(*shape)

    mapper = tmp_target.get_mapper()
    zarr_array = zarr.open(mapper, mode="w", shape=shape, chunks=zarr_chunks, dtype=data.dtype)

    @dask.delayed
    def write_block(n):
        write_region = slice(chunk_bounds[n], chunk_bounds[n + 1])
        with lock_for_conflicts(chunk_conflicts[n]):
            zarr_array[write_region] = data[write_region]

    this_client = Client(dask_cluster)
    dask.compute([write_block(n) for n in range(len(sequence_lens))])
    np.testing.assert_equal(data, zarr_array)

    this_client.close()
    del this_client
