import json
import os

import pytest

from assonant._bluesky.data_parser import DocumentsDataParser
from assonant.naming_standards import BeamlineName


def assert_query_config_by_pv_name_structure(config_dict: dict) -> None:
    """
    Validate structure returned by query_config_by_pv_name() from bluesky data parser.
    """
    assert isinstance(config_dict, dict), "Configs must be a dict"

    for pv_name, pv_info in config_dict.items():
        assert isinstance(pv_name, str), "Config PV name must be str"
        assert isinstance(pv_info, dict), f"{pv_name} value must be dict"

        required_keys = {"value", "timestamps"}
        assert required_keys == set(pv_info.keys()), f"{pv_name} must contain keys {required_keys}"


def assert_query_data_by_stream_structure(data_dict: dict, config_may_exist: bool = True) -> None:
    """
    Validate structure returned by query_data_by_stream_name() from bluesky data parser.
    """
    assert isinstance(data_dict, dict), "Stream data must be a dict"

    for seq_num, seq_dict in data_dict.items():
        assert isinstance(seq_num, int), f"{seq_num} must be an int"
        assert isinstance(seq_dict, dict), f"{seq_num} must map to dict"

        for pv_name, pv_info in seq_dict.items():
            assert isinstance(pv_name, str), f"{pv_name} must be an str"
            assert isinstance(pv_info, dict), f"{pv_name} must map to dict"

            required_keys = {"data"}
            assert required_keys.issubset(pv_info.keys()), f"{pv_name} must contain keys {required_keys}"

            # validate data block
            data_block = pv_info["data"]
            assert isinstance(data_block, dict), f"{pv_name} data must be dict"

            data_required = {"value", "timestamps"}
            assert data_required == set(data_block.keys()), f"{pv_name} data must contain {data_required}"

            # validate configs block
            if config_may_exist is True:
                if "configs" in pv_info.keys():
                    assert_query_config_by_pv_name_structure(pv_info["configs"])
                # else: Nothing happens as "configs" is optional
            else:
                assert "configs" not in pv_info.keys(), f"{pv_name} must not contain configs block in this test case"


def assert_query_all_data_structure(all_data: dict, config_may_exist: bool = True) -> None:
    """
    Validate structure returned by query_all_data() from bluesky data parser.
    """
    assert isinstance(all_data, dict), "All data must be dict"

    for stream_name, stream_data in all_data.items():
        assert isinstance(stream_name, str), "Stream name must be str"
        assert isinstance(stream_data, dict), "Stream data must map to dict"

        assert_query_data_by_stream_structure(stream_data, config_may_exist)


@pytest.fixture(scope="session")
def example_bluesky_data_parser(example_bluesky_all_streams_acquisition_file_path) -> DocumentsDataParser:
    """
    Provides a functional DocumentsDataParser for the example_bluesky_all_streams_acquisition_file.json
    """

    with open(example_bluesky_all_streams_acquisition_file_path) as f:
        example_bluesky_docs_list = json.load(f)

    return DocumentsDataParser(bluesky_documents_list=example_bluesky_docs_list)


def test_query_all_data(example_bluesky_data_parser):
    """
    Validate that query_all_data method of DocumentsDataParser works correctly and don't break when using
    it to retrieve data from the example_bluesky_all_streams_acquisition_file.json file.

    PS: Only types and structure is being checked, not values!

    Expected Behavior:
        - The returned value must be a dictionary with the structure described in the method documentation.
        - The returned value must contain data related to all Acquisition Types (Primary, Baseline and Monitor)
    """
    retrieved_data = example_bluesky_data_parser.query_all_data(with_config=True)
    # Assert that the returned structure is correct
    assert_query_all_data_structure(retrieved_data, config_may_exist=True)


def test_query_data_by_stream_name(example_bluesky_data_parser):
    """
    Validate that query_data_by_stream_name method of DocumentsDataParser works correctly and don't break when using
    it to retrieve data from the example_bluesky_all_streams_acquisition_file.json file.

    PS: Only types and structure is being checked, not values!

    Expected Behavior:
        - The returned value must be a dictionary with the structure described in the method documentation.
        - The returned value must contain data related to the specified Acquisition Type
    """
    for stream_name in example_bluesky_data_parser.get_stream_names():

        retrieved_data = example_bluesky_data_parser.query_data_by_stream_name(event_stream_name=stream_name)

        # Assert that the returned structure is correct
        assert_query_data_by_stream_structure(retrieved_data, config_may_exist=True)


def test_query_config_by_pv_name(example_bluesky_data_parser):
    """
    Validate that query_config_by_pv_name method of DocumentsDataParser works correctly and don't break when using
    it to retrieve data from the example_bluesky_all_streams_acquisition_file.json file.

    PS: Only types and structure is being checked, not values!

    Expected Behavior:
        - The returned value must be a dictionary with the structure described in the
        method documentation. (Check if assertion method is correctly reflecting it!)
        - The returned value must contain config information for all PVs related to all Acquisition Types
    """
    for stream_name in example_bluesky_data_parser.get_stream_names():

        stream_data = example_bluesky_data_parser.query_data_by_stream_name(event_stream_name=stream_name)

        for pv_name in stream_data.keys():
            retrieved_pv_config = example_bluesky_data_parser.query_config_by_pv_name(
                pv_name=pv_name, stream_name=stream_name
            )
            # Assert that the returned structure is correct
            assert_query_config_by_pv_name_structure(retrieved_pv_config)


def test_query_all_data_without_config(example_bluesky_data_parser):
    """
    Validate that query_all_data method of DocumentsDataParser works correctly and don't break when using
    it to retrieve data from the example_bluesky_all_streams_acquisition_file.json without config metadata.

    PS: Only types and structure is being checked, not values!

    Expected Behavior:
        - Returned structure must be the same as query_all_data with config but without the "configs" block in
    the innermost dicts.
    """
    retrieved_data = example_bluesky_data_parser.query_all_data(with_config=False)
    # Assert that the returned structure is correct
    assert_query_all_data_structure(retrieved_data, config_may_exist=False)


def test_query_data_by_stream_name_without_config(example_bluesky_data_parser):
    """
    Validate that query_data_by_stream_name method of DocumentsDataParser works correctly and don't break when using
    it to retrieve data from the example_bluesky_all_streams_acquisition_file.json without config metadata.

    PS: Only types and structure is being checked, not values!

    Expected Behavior:
        - Returned structure must be the same as query_data_by_stream_name with config but without the
        "configs" block in the innermost dicts.
    """
    for stream_name in example_bluesky_data_parser.get_stream_names():

        retrieved_data = example_bluesky_data_parser.query_data_by_stream_name(
            event_stream_name=stream_name, with_config=False
        )

        # Assert that the returned structure is correct
        assert_query_data_by_stream_structure(retrieved_data, config_may_exist=False)


@pytest.mark.parametrize("beamline_name", [beamline_name for beamline_name in BeamlineName])
@pytest.mark.forked
def test_parse_real_data_with_config_from_beamline_example_bluesky_files(
    beamline_name, tmp_path, real_bluesky_data_test_cases_dir_path_map
):
    """
    Validate that DataParser works correctly and don't break when using
    it to retrieve data from beamline real bluesky example JSON files.

    PS: Only types and structure is being checked, not values!

    Expected Behavior:
        - Returned structure for each case must be the same no matter from which beamline it comes.
        - Returned structure must be the same as docummented on method docstring for
        query_all_data() method. (Check if assertion method is up-to-date!)
        - Returned structure MUST contain config blocks which also must be checked by the
        assertion method.
    """
    if beamline_name.value not in real_bluesky_data_test_cases_dir_path_map.keys():
        pytest.skip(reason=f"There is no test case with real data for {beamline_name} beamline yet!")
    else:
        for i in range(0, len(real_bluesky_data_test_cases_dir_path_map[beamline_name.value])):
            # Execute test for each test case dir that exist for that beamline
            test_case_dir = real_bluesky_data_test_cases_dir_path_map[beamline_name.value][i]

            for test_case_file in os.scandir(test_case_dir):
                if test_case_file.is_file() and test_case_file.name.endswith(".json"):

                    with open(test_case_file) as f:
                        test_case_bluesky_docs_list = json.load(f)

                    doc_parser = DocumentsDataParser(bluesky_documents_list=test_case_bluesky_docs_list)

                    # Assert that Data Parser can retrieve all data with config without error
                    all_data = doc_parser.query_all_data(with_config=True)
                    assert_query_all_data_structure(all_data, config_may_exist=True)


@pytest.mark.parametrize("beamline_name", [beamline_name for beamline_name in BeamlineName])
@pytest.mark.forked
def test_parse_real_data_without_config_from_beamline_example_bluesky_files(
    beamline_name, tmp_path, real_bluesky_data_test_cases_dir_path_map
):
    """
    Validate that DataParser works correctly and don't break when using it to retrieve
    data from beamline real bluesky example JSON files without config metadata.

    PS: Only types and structure is being checked, not values!

    Expected Behavior:
        - Returned structure for each case must be the same no matter from which beamline it comes.
        - Returned structure must be the same as docummented on method docstring for
        query_all_data() method. (Check if assertion method is up-to-date!)
        - Returned structure must NOT contain config blocks which also must be checked by the
        assertion method.
    """
    if beamline_name.value not in real_bluesky_data_test_cases_dir_path_map.keys():
        pytest.skip(reason=f"There is no test case with real data for {beamline_name} beamline yet!")
    else:

        for i in range(0, len(real_bluesky_data_test_cases_dir_path_map[beamline_name.value])):
            # Execute test for each test case dir that exist for that beamline
            test_case_dir = real_bluesky_data_test_cases_dir_path_map[beamline_name.value][i]

            for test_case_file in os.scandir(test_case_dir):
                if test_case_file.is_file() and test_case_file.name.endswith(".json"):

                    with open(test_case_file) as f:
                        test_case_bluesky_docs_list = json.load(f)

                    doc_parser = DocumentsDataParser(bluesky_documents_list=test_case_bluesky_docs_list)

                    # Assert that Data Parser can retrieve all data with config without error
                    all_data = doc_parser.query_all_data(with_config=False)
                    assert_query_all_data_structure(all_data, config_may_exist=False)
