from __future__ import annotations
from pathlib import Path
import sys

from typing import Any, cast, TYPE_CHECKING
import numpy as np
import zarr  # type: ignore
import pytest
from hpcflow.sdk.core.test_utils import (
    make_schemas,
    make_test_data_YAML_workflow,
    make_workflow,
)
from hpcflow.sdk.persistence.json import (
    JSONPersistentStore,
    JsonStoreElement,
    JsonStoreElementIter,
    JsonStoreEAR,
)
from hpcflow.sdk.persistence.zarr import ZarrPersistentStore
from hpcflow.sdk.core.parameters import NullDefault

from hpcflow.app import app as hf

if TYPE_CHECKING:
    from hpcflow.sdk.persistence.zarr import ZarrPersistentStore


@pytest.mark.skip("need to refactor `make_test_store_from_spec`")
def test_store_pending_add_task(tmp_path: Path):
    """Check expected pending state after adding a task."""

    # make store: 0 tasks:
    store = JSONPersistentStore.make_test_store_from_spec(hf, [], dir=tmp_path)
    task_ID = store.add_task()
    assert store._pending.add_tasks == {task_ID: []}


@pytest.mark.skip("need to refactor `make_test_store_from_spec`")
def test_store_pending_add_element(tmp_path: Path):
    """Check expected pending state after adding an element."""

    # make store: 1 task with 0 elements:
    store = JSONPersistentStore.make_test_store_from_spec(app=hf, spec=[{}], dir=tmp_path)
    elem_ID = store.add_element(task_ID=0)
    assert store._pending.add_elements == {
        elem_ID: JsonStoreElement(
            id_=elem_ID,
            is_pending=True,
            es_idx=0,
            task_ID=0,
            iteration_IDs=[],
            index=0,
            seq_idx={},
            src_idx={},
        )
    } and store._pending.add_task_element_IDs == {0: [0]}


@pytest.mark.skip("need to refactor `make_test_store_from_spec`")
@pytest.mark.parametrize("elem_ID", [0, 1])
def test_store_pending_add_element_iter(tmp_path: Path, elem_ID: int):
    """Check expected pending state after adding an element iteration."""

    # make store: 1 task with 2 elements and 0 iterations:
    store = JSONPersistentStore.make_test_store_from_spec(
        hf,
        [{"elements": [{}, {}]}],
        dir=tmp_path,
    )
    iter_ID = store.add_element_iteration(
        element_ID=elem_ID,
        data_idx={},
        schema_parameters=[],
    )
    assert store._pending.add_elem_iters == {
        iter_ID: JsonStoreElementIter(
            id_=iter_ID,
            is_pending=True,
            element_ID=elem_ID,
            EAR_IDs={},
            data_idx={},
            schema_parameters=[],
            EARs_initialised=False,
        )
    } and store._pending.add_elem_iter_IDs == {elem_ID: [iter_ID]}


@pytest.mark.skip("need to refactor `make_test_store_from_spec`")
def test_store_pending_add_EAR(tmp_path: Path):
    """Check expected pending state after adding an EAR."""

    # make store: 1 task with 1 element and 1 iteration:
    store = JSONPersistentStore.make_test_store_from_spec(
        hf,
        [{"elements": [{"iterations": [{}]}]}],
        dir=tmp_path,
    )
    EAR_ID = store.add_EAR(
        elem_iter_ID=0,
        action_idx=0,
        commands_idx=[],
        data_idx={},
        metadata={},
    )
    assert store._pending.add_EARs == {
        EAR_ID: JsonStoreEAR(
            id_=EAR_ID,
            is_pending=True,
            elem_iter_ID=0,
            action_idx=0,
            commands_idx=[],
            data_idx={},
            metadata={},
        )
    }


@pytest.mark.skip("need to refactor `make_test_store_from_spec`")
def test_get_task_elements_task_is_pending(tmp_path: Path):
    """Check we get an empty list when getting all task elements of a pending task to
    which no elements have been added."""
    # make store: 0 tasks:
    store = JSONPersistentStore.make_test_store_from_spec(hf, [], dir=tmp_path)
    task_ID = store.add_task()
    assert store.get_task_elements(task_ID, slice(0, None)) == []


@pytest.mark.skip("need to refactor `make_test_store_from_spec`")
def test_get_task_elements_single_element_is_pending(tmp_path: Path):
    """Check expected return when getting all task elements of a persistent task that has
    a single pending element."""
    # make store: 1 task
    store = JSONPersistentStore.make_test_store_from_spec(hf, [{}], dir=tmp_path)
    store.add_element(task_ID=0)
    assert store.get_task_elements(0, slice(0, None)) == [
        {
            "id": 0,
            "is_pending": True,
            "element_idx": 0,
            "iteration_IDs": [],
            "task_ID": 0,
            "iterations": [],
        }
    ]


@pytest.mark.skip("need to refactor `make_test_store_from_spec`")
def test_get_task_elements_multi_element_one_pending(tmp_path: Path):
    """Check expected return when getting all task elements of a persistent task that has
    a persistent element and a pending element."""
    # make store: 1 task with 1 element:
    store = JSONPersistentStore.make_test_store_from_spec(
        hf, [{"elements": [{}]}], dir=tmp_path
    )
    store.add_element(task_ID=0)
    assert store.get_task_elements(0, slice(0, None)) == [
        {
            "id": 0,
            "is_pending": False,
            "element_idx": 0,
            "iteration_IDs": [],
            "task_ID": 0,
            "iterations": [],
        },
        {
            "id": 1,
            "is_pending": True,
            "element_idx": 1,
            "iteration_IDs": [],
            "task_ID": 0,
            "iterations": [],
        },
    ]


@pytest.mark.skip("need to refactor `make_test_store_from_spec`")
def test_get_task_elements_single_element_iter_pending(tmp_path: Path):
    """Check expected return when getting all task elements of a persistent task that has
    a persistent element with a pending iteration."""
    # make store: 1 task with 1 element:
    store = JSONPersistentStore.make_test_store_from_spec(
        hf, [{"elements": [{}]}], dir=tmp_path
    )
    store.add_element_iteration(element_ID=0, data_idx={}, schema_parameters=[])
    assert store.get_task_elements(0, slice(0, None)) == [
        {
            "id": 0,
            "is_pending": False,
            "element_idx": 0,
            "iteration_IDs": [0],
            "task_ID": 0,
            "iterations": [
                {
                    "id": 0,
                    "is_pending": True,
                    "element_ID": 0,
                    "EAR_IDs": {},
                    "data_idx": {},
                    "schema_parameters": [],
                    "EARs": {},
                }
            ],
        },
    ]


@pytest.mark.skip("need to refactor `make_test_store_from_spec`")
def test_get_task_elements_single_element_iter_EAR_pending(tmp_path: Path):
    """Check expected return when getting all task elements of a persistent task that has
    a persistent element with a persistent iteration and a pending EAR"""
    # make store: 1 task with 1 element with 1 iteration:
    store = JSONPersistentStore.make_test_store_from_spec(
        hf, [{"elements": [{"iterations": [{}]}]}], dir=tmp_path
    )
    store.add_EAR(elem_iter_ID=0, action_idx=0, commands_idx=[], data_idx={}, metadata={})
    assert store.get_task_elements(0, slice(0, None)) == [
        {
            "id": 0,
            "is_pending": False,
            "element_idx": 0,
            "iteration_IDs": [0],
            "task_ID": 0,
            "iterations": [
                {
                    "id": 0,
                    "is_pending": False,
                    "element_ID": 0,
                    "EAR_IDs": {0: [0]},
                    "data_idx": {},
                    "schema_parameters": [],
                    "EARs": {
                        0: [
                            {
                                "id_": 0,
                                "is_pending": True,
                                "elem_iter_ID": 0,
                                "action_idx": 0,
                                "commands_idx": [],
                                "data_idx": {},
                                "metadata": {},
                            }
                        ]
                    },
                },
            ],
        },
    ]


def test_make_zarr_store_zstd_compressor(null_config, tmp_path: Path):
    wk = make_test_data_YAML_workflow(
        workflow_name="workflow_1.yaml",
        path=tmp_path,
        store="zarr",
        store_kwargs={"compressor": "zstd"},
    )


def test_make_zarr_store_no_compressor(null_config, tmp_path: Path):
    wk = make_test_data_YAML_workflow(
        workflow_name="workflow_1.yaml",
        path=tmp_path,
        store="zarr",
        store_kwargs={"compressor": None},
    )


@pytest.mark.integration
@pytest.mark.skipif(
    sys.version_info < (3, 9), reason="Python 3.8 support is being removed anyway."
)
def test_zarr_rechunk_data_equivalent(null_config, tmp_path: Path):
    t1 = hf.Task(
        schema=hf.task_schemas.test_t1_conditional_OS,
        inputs={"p1": 101},
        repeats=3,
    )
    wk = hf.Workflow.from_template_data(
        tasks=[t1],
        template_name="test_run_rechunk",
        workflow_name="test_run_rechunk",
        path=tmp_path,
    )
    wk.submit(wait=True, status=False, add_to_known=False)
    wk.rechunk_runs(backup=True, status=False, chunk_size=None)  # None -> one chunk

    arr = cast("ZarrPersistentStore", wk._store)._get_EARs_arr()
    assert arr.chunks == arr.shape

    bak_path = (Path(arr.store.path) / arr.path).with_suffix(".bak")
    arr_bak = zarr.open(bak_path)

    assert arr_bak.chunks == (1, 1)  # runs array is 2D

    # check backup and new runs data are equal:
    assert np.all(arr[:] == arr_bak[:])

    # check attributes are equal:
    assert arr.attrs.asdict() == arr_bak.attrs.asdict()


@pytest.mark.integration
@pytest.mark.skipif(
    sys.version_info < (3, 9), reason="Python 3.8 support is being removed anyway."
)
def test_zarr_rechunk_data_equivalent_custom_chunk_size(null_config, tmp_path: Path):
    t1 = hf.Task(
        schema=hf.task_schemas.test_t1_conditional_OS,
        inputs={"p1": 101},
        repeats=3,
    )
    wk = hf.Workflow.from_template_data(
        tasks=[t1],
        template_name="test_run_rechunk",
        workflow_name="test_run_rechunk",
        path=tmp_path,
    )
    wk.submit(wait=True, status=False, add_to_known=False)
    wk.rechunk_runs(backup=True, status=False, chunk_size=2)

    arr = cast("ZarrPersistentStore", wk._store)._get_EARs_arr()
    assert arr.chunks == (2, 2)  # runs array is 2D

    bak_path = (Path(arr.store.path) / arr.path).with_suffix(".bak")
    arr_bak = zarr.open(bak_path)

    assert arr_bak.chunks == (1, 1)  # runs array is 2D

    # check backup and new runs data are equal:
    assert np.all(arr[:] == arr_bak[:])


@pytest.mark.integration
def test_zarr_rechunk_data_no_backup_load_runs(null_config, tmp_path: Path):
    t1 = hf.Task(
        schema=hf.task_schemas.test_t1_conditional_OS,
        inputs={"p1": 101},
        repeats=3,
    )
    wk = hf.Workflow.from_template_data(
        tasks=[t1],
        template_name="test_run_rechunk",
        workflow_name="test_run_rechunk",
        path=tmp_path,
    )
    wk.submit(wait=True, status=False, add_to_known=False)
    wk.rechunk_runs(backup=False, status=False)

    arr = cast("ZarrPersistentStore", wk._store)._get_EARs_arr()

    bak_path = (Path(arr.store.path) / arr.path).with_suffix(".bak")
    assert not bak_path.is_file()

    # check we can load runs:
    runs = wk._store._get_persistent_EARs(id_lst=list(range(wk.num_EARs)))
    run_ID = []
    for i in runs.values():
        run_ID.append(i.id_)


@pytest.mark.integration
def test_zarr_rechunk_data_no_backup_load_parameter_base(null_config, tmp_path: Path):
    t1 = hf.Task(
        schema=hf.task_schemas.test_t1_conditional_OS,
        inputs={"p1": 101},
        repeats=3,
    )
    wk = hf.Workflow.from_template_data(
        tasks=[t1],
        template_name="test_run_rechunk",
        workflow_name="test_run_rechunk",
        path=tmp_path,
    )
    wk.submit(wait=True, status=False, add_to_known=False)

    params_old = wk.get_all_parameter_data()
    wk.rechunk_parameter_base(backup=False, status=False)

    wk = wk.reload()
    params_new = wk.get_all_parameter_data()
    assert params_new == params_old

    arr = cast("ZarrPersistentStore", wk._store)._get_parameter_base_array()

    bak_path = (Path(arr.store.path) / arr.path).with_suffix(".bak")
    assert not bak_path.is_file()

    # check we can load parameters:
    param_IDs = []
    for i in wk.get_all_parameters():
        param_IDs.append(i.id_)


def test_get_parameter_sources_duplicate_ids(null_config, tmp_path):
    wk = make_workflow(
        schemas_spec=[[{"p1": None}, ("p1",), "t1"]],
        local_inputs={0: ("p1",)},
        path=tmp_path,
    )
    id_lst = [0, 1, 1, 2, 0]
    src = wk._store.get_parameter_sources(id_lst)
    assert len(src) == len(id_lst)
    assert src[0] == src[4]
    assert src[1] == src[2]


def _transform_jobscript_dependencies_to_encodable(
    deps: dict[tuple[int, int], dict[tuple[int, int], dict[str, Any]]],
) -> dict[str, list[dict[str, Any]]]:
    """Transform a dict of jobscript dependencies written in a more testing-friendly/
    convenient format into the format expected by the method
    `ZarrPersistentStore._encode_jobscript_block_dependencies`.

    """
    max_js_idx = max(i[0] for i in deps)
    sub_js: dict[str, list[dict[str, Any]]] = {
        "jobscripts": [
            {"blocks": [], "index": js_idx} for js_idx in range(max_js_idx + 1)
        ]
    }
    for (js_idx, blk_idx), deps_i in deps.items():
        sub_js["jobscripts"][js_idx]["blocks"].append(
            {
                "dependencies": [[[k[0], k[1]], v] for k, v in deps_i.items()],
                "index": blk_idx,
            }
        )
    return sub_js


def test_zarr_encode_jobscript_block_dependencies_element_mapping_array_non_array_equivalence():
    deps_1 = {
        (0, 0): {},
        (1, 0): {(0, 0): {"js_element_mapping": {0: [0]}, "is_array": True}},
    }
    deps_2 = {
        (0, 0): {},
        (1, 0): {(0, 0): {"js_element_mapping": {0: np.array([0])}, "is_array": True}},
    }
    deps_1 = _transform_jobscript_dependencies_to_encodable(deps_1)
    deps_2 = _transform_jobscript_dependencies_to_encodable(deps_2)
    arr_1 = ZarrPersistentStore._encode_jobscript_block_dependencies(deps_1)
    arr_2 = ZarrPersistentStore._encode_jobscript_block_dependencies(deps_2)
    assert np.array_equal(arr_1, arr_2)


def test_zarr_encode_decode_jobscript_block_dependencies():

    deps = {
        (0, 0): {},
        (1, 0): {
            (0, 0): {
                "js_element_mapping": {0: [0], 1: [1]},
                "is_array": True,
            }
        },
        (2, 0): {
            (1, 0): {
                "js_element_mapping": {0: [0, 1], 1: [0, 1]},
                "is_array": False,
            }
        },
        (2, 1): {
            (0, 0): {"js_element_mapping": {0: [0, 1]}, "is_array": False},
            (2, 0): {"js_element_mapping": {0: [0, 1]}, "is_array": False},
        },
    }
    deps_t = _transform_jobscript_dependencies_to_encodable(deps)
    arr = ZarrPersistentStore._encode_jobscript_block_dependencies(deps_t)
    assert np.array_equal(
        arr,
        np.array(
            [
                2,
                0,
                0,
                12,
                1,
                0,
                9,
                0,
                0,
                1,
                2,
                0,
                0,
                2,
                1,
                1,
                14,
                2,
                0,
                11,
                1,
                0,
                0,
                3,
                0,
                0,
                1,
                3,
                1,
                0,
                1,
                18,
                2,
                1,
                7,
                0,
                0,
                0,
                3,
                0,
                0,
                1,
                7,
                2,
                0,
                0,
                3,
                0,
                0,
                1,
            ]
        ),
    )
    deps_rt = ZarrPersistentStore._decode_jobscript_block_dependencies(arr)
    assert deps_rt == deps


def test_zarr_encode_decode_jobscript_block_dependencies_large_many_to_one():
    deps = {
        (0, 0): {},
        (1, 0): {
            (0, 0): {"js_element_mapping": {0: list(range(1_000_000))}, "is_array": False}
        },
    }
    deps_t = _transform_jobscript_dependencies_to_encodable(deps)
    arr = ZarrPersistentStore._encode_jobscript_block_dependencies(deps_t)
    deps_rt = ZarrPersistentStore._decode_jobscript_block_dependencies(arr)
    assert deps_rt == deps


def test_zarr_encode_decode_jobscript_block_dependencies_large_one_to_one():
    deps = {
        (0, 0): {},
        (1, 0): {
            (0, 0): {
                "js_element_mapping": {i: [i] for i in range(1_000_000)},
                "is_array": False,
            }
        },
    }
    deps_t = _transform_jobscript_dependencies_to_encodable(deps)
    arr = ZarrPersistentStore._encode_jobscript_block_dependencies(deps_t)
    deps_rt = ZarrPersistentStore._decode_jobscript_block_dependencies(arr)
    assert deps_rt == deps


@pytest.mark.parametrize(
    "array",
    (
        np.array([]),
        np.empty(0),
        np.empty((0, 1, 2)),
        np.array([1, 2, 3]),
        np.array([[1, 2, 3], [4, 5, 6]]),
    ),
)
def test_zarr_save_persistent_array_shape(null_config, tmp_path, array):
    s1 = make_schemas(({"p1": None}, ()))
    t1 = hf.Task(schema=s1, inputs={"p1": array})
    wk = hf.Workflow.from_template_data(
        template_name="test_save_empty_array",
        tasks=[t1],
        path=tmp_path,
    )
    assert array.shape == wk.tasks[0].elements[0].get("inputs.p1")[:].shape


def test_zarr_single_chunk_threshold(null_config, tmp_path):
    # test very large arrays (> ~1 GB) are saved using multiple chunks
    array = np.zeros(
        268_435_456
    )  # ~ 2.147483647 GB; greater than blosc's max buffer size
    s1 = make_schemas(({"p1": None}, ()))
    t1 = hf.Task(schema=s1, inputs={"p1": array})
    wk = hf.Workflow.from_template_data(
        template_name="test_large_array",
        tasks=[t1],
        path=tmp_path,
    )


@pytest.mark.parametrize("store", ["json", "zarr"])
def test_store_parameter_encode_decode_types(null_config, tmp_path, store):

    (s1,) = make_schemas(
        (
            {
                "p1": NullDefault.NULL,
                "p2": NullDefault.NULL,
                "p3": NullDefault.NULL,
                "p4": NullDefault.NULL,
                "p5": NullDefault.NULL,
                "p6": NullDefault.NULL,
                "p7": NullDefault.NULL,
            },
            tuple(),
        ),
    )

    p1 = [1, 2, 3]
    p2 = (1, 2, 3)
    p3 = {1, 2, 3}
    p4 = None
    p5 = np.arange(10)
    p6 = np.ma.array(np.arange(10), mask=np.random.randint(0, 2, 10))
    p7 = [[1, 2], (3, 4), {5, 6}]

    t1 = hf.Task(
        schema=s1,
        inputs={
            "p1": p1,
            "p2": p2,
            "p3": p3,
            "p4": p4,
            "p5": p5 if store == "zarr" else None,
            "p6": p6 if store == "zarr" else None,
            "p7": p7,
        },
    )

    wk = hf.Workflow.from_template_data(
        template_name="test_store_encoders",
        tasks=[t1],
        store=store,
        path=tmp_path,
    )

    assert wk.tasks[0].elements[0].get("inputs.p1") == p1
    assert wk.tasks[0].elements[0].get("inputs.p2") == p2
    assert wk.tasks[0].elements[0].get("inputs.p3") == p3
    assert wk.tasks[0].elements[0].get("inputs.p4") == p4
    assert wk.tasks[0].elements[0].get("inputs.p7") == p7

    if store == "zarr":
        assert np.allclose(wk.tasks[0].elements[0].get("inputs.p5"), p5)
        assert np.allclose(wk.tasks[0].elements[0].get("inputs.p6"), p6)
