# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.

import json
import os
import tempfile
from dataclasses import dataclass
from unittest.mock import patch

import numpy as np
import pytest

from braket.jobs.data_persistence import (
    load_job_checkpoint,
    load_job_result,
    save_job_checkpoint,
    save_job_result,
)
from braket.jobs_data import PersistedJobDataFormat


@pytest.mark.parametrize(
    "job_name, file_suffix, data_format, checkpoint_data, expected_saved_data",
    [
        (
            "job_plaintext_simple_dict",
            "",
            PersistedJobDataFormat.PLAINTEXT,
            {"converged": True, "energy": -0.2},
            json.dumps({
                "braketSchemaHeader": {
                    "name": "braket.jobs_data.persisted_job_data",
                    "version": "1",
                },
                "dataDictionary": {"converged": True, "energy": -0.2},
                "dataFormat": "plaintext",
            }),
        ),
        (
            "job_pickled_simple_dict",
            "suffix1",
            PersistedJobDataFormat.PICKLED_V4,
            {"converged": True, "energy": -0.2},
            json.dumps({
                "braketSchemaHeader": {
                    "name": "braket.jobs_data.persisted_job_data",
                    "version": "1",
                },
                "dataDictionary": {
                    "converged": "gASILg==\n",
                    "energy": "gASVCgAAAAAAAABHv8mZmZmZmZou\n",
                },
                "dataFormat": "pickled_v4",
            }),
        ),
    ],
)
def test_save_job_checkpoint(
    job_name, file_suffix, data_format, checkpoint_data, expected_saved_data
):
    with tempfile.TemporaryDirectory() as tmp_dir:
        with patch.dict(
            os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
        ):
            save_job_checkpoint(checkpoint_data, file_suffix, data_format)

        expected_file_location = (
            f"{tmp_dir}/{job_name}_{file_suffix}.json"
            if file_suffix
            else f"{tmp_dir}/{job_name}.json"
        )
        with open(expected_file_location) as expected_file:
            assert expected_file.read() == expected_saved_data


@pytest.mark.xfail(raises=ValueError)
@pytest.mark.parametrize("checkpoint_data", [{}, None])
def test_save_job_checkpoint_raises_error_empty_data(checkpoint_data):
    job_name = "foo"
    with tempfile.TemporaryDirectory() as tmp_dir:
        with patch.dict(
            os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
        ):
            save_job_checkpoint(checkpoint_data)


@pytest.mark.parametrize(
    "job_name, file_suffix, data_format, saved_data, expected_checkpoint_data",
    [
        (
            "job_plaintext_simple_dict",
            "",
            PersistedJobDataFormat.PLAINTEXT,
            json.dumps({
                "braketSchemaHeader": {
                    "name": "braket.jobs_data.persisted_job_data",
                    "version": "1",
                },
                "dataDictionary": {"converged": True, "energy": -0.2},
                "dataFormat": "plaintext",
            }),
            {"converged": True, "energy": -0.2},
        ),
        (
            "job_pickled_simple_dict",
            "",
            PersistedJobDataFormat.PICKLED_V4,
            json.dumps({
                "braketSchemaHeader": {
                    "name": "braket.jobs_data.persisted_job_data",
                    "version": "1",
                },
                "dataDictionary": {
                    "converged": "gASILg==\n",
                    "energy": "gASVCgAAAAAAAABHv8mZmZmZmZou\n",
                },
                "dataFormat": "pickled_v4",
            }),
            {"converged": True, "energy": -0.2},
        ),
    ],
)
def test_load_job_checkpoint(
    job_name, file_suffix, data_format, saved_data, expected_checkpoint_data
):
    with tempfile.TemporaryDirectory() as tmp_dir:
        file_path = (
            f"{tmp_dir}/{job_name}_{file_suffix}.json"
            if file_suffix
            else f"{tmp_dir}/{job_name}.json"
        )
        with open(file_path, "w") as f:
            f.write(saved_data)

        with patch.dict(
            os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
        ):
            loaded_data = load_job_checkpoint(job_name, file_suffix)
            assert loaded_data == expected_checkpoint_data


@pytest.mark.xfail(raises=FileNotFoundError)
def test_load_job_checkpoint_raises_error_file_not_exists():
    job_name = "old_job"
    file_suffix = "correct_suffix"
    with tempfile.TemporaryDirectory() as tmp_dir:
        file_path = f"{tmp_dir}/{job_name}_{file_suffix}.json"
        with open(file_path, "w") as _:
            pass

        with patch.dict(
            os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
        ):
            load_job_checkpoint(job_name, "wrong_suffix")


@pytest.mark.xfail(raises=ValueError)
def test_load_job_checkpoint_raises_error_corrupted_data():
    job_name = "old_job_corrupted_data"
    file_suffix = "foo"
    with tempfile.TemporaryDirectory() as tmp_dir:
        file_path = f"{tmp_dir}/{job_name}_{file_suffix}.json"
        with open(file_path, "w") as corrupted_file:
            corrupted_file.write(
                json.dumps({
                    "braketSchemaHeader": {
                        "name": "braket.jobs_data.persisted_job_data",
                        "version": "1",
                    },
                    "dataDictionary": {
                        "converged": "gASILg==\n",
                        "energy": "gASVCgBHv--corrupted---\n",
                    },
                    "dataFormat": "pickled_v4",
                })
            )

        with patch.dict(
            os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
        ):
            load_job_checkpoint(job_name, file_suffix)


@dataclass
class CustomClassToPersist:
    float_val: float
    str_val: str
    bool_val: bool


def test_save_and_load_job_checkpoint():
    with tempfile.TemporaryDirectory() as tmp_dir:
        job_name = "job_name_1"
        data = {
            "np_array": np.array([1]),
            "custom_class": CustomClassToPersist(3.4, "str", True),
            "none_value": None,
            "nested_dict": {"a": {"b": False}},
        }
        with patch.dict(
            os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
        ):
            save_job_checkpoint(data, data_format=PersistedJobDataFormat.PICKLED_V4)
            retrieved = load_job_checkpoint(job_name)
            assert retrieved == data


@pytest.mark.parametrize(
    "data_format, result_data, expected_saved_data",
    [
        (
            PersistedJobDataFormat.PLAINTEXT,
            {"converged": True, "energy": -0.2},
            json.dumps({
                "braketSchemaHeader": {
                    "name": "braket.jobs_data.persisted_job_data",
                    "version": "1",
                },
                "dataDictionary": {"converged": True, "energy": -0.2},
                "dataFormat": "plaintext",
            }),
        ),
        (
            PersistedJobDataFormat.PICKLED_V4,
            {"converged": True, "energy": -0.2},
            json.dumps({
                "braketSchemaHeader": {
                    "name": "braket.jobs_data.persisted_job_data",
                    "version": "1",
                },
                "dataDictionary": {
                    "converged": "gASILg==\n",
                    "energy": "gASVCgAAAAAAAABHv8mZmZmZmZou\n",
                },
                "dataFormat": "pickled_v4",
            }),
        ),
    ],
)
def test_save_job_result(data_format, result_data, expected_saved_data):
    with tempfile.TemporaryDirectory() as tmp_dir:
        with patch.dict(os.environ, {"AMZN_BRAKET_JOB_RESULTS_DIR": tmp_dir}):
            save_job_result(result_data, data_format)

        expected_file_location = f"{tmp_dir}/results.json"
        with open(expected_file_location) as expected_file:
            assert expected_file.read() == expected_saved_data


@pytest.mark.parametrize("result_data", [{}, None])
def test_save_job_result_does_not_raise_error_empty_data(result_data):
    with tempfile.TemporaryDirectory() as tmp_dir:
        with patch.dict(os.environ, {"AMZN_BRAKET_JOB_RESULTS_DIR": tmp_dir}):
            save_job_result(result_data)


@pytest.mark.parametrize(
    "first_result_data,"
    "first_data_format,"
    "second_result_data,"
    "second_data_format,"
    "expected_result_data",
    (
        (
            "hello",
            PersistedJobDataFormat.PLAINTEXT,
            "goodbye",
            PersistedJobDataFormat.PLAINTEXT,
            {"result": "goodbye"},
        ),
        (
            "hello",
            PersistedJobDataFormat.PLAINTEXT,
            "goodbye",
            PersistedJobDataFormat.PICKLED_V4,
            {"result": "goodbye"},
        ),
        ("hello", PersistedJobDataFormat.PICKLED_V4, "goodbye", None, {"result": "goodbye"}),
        (
            # not json serializable
            PersistedJobDataFormat,
            PersistedJobDataFormat.PICKLED_V4,
            {"other_field": "value"},
            None,
            {"result": PersistedJobDataFormat, "other_field": "value"},
        ),
    ),
)
def test_update_result_data(
    first_result_data,
    first_data_format,
    second_result_data,
    second_data_format,
    expected_result_data,
):
    with tempfile.TemporaryDirectory() as tmp_dir:
        with patch.dict(os.environ, {"AMZN_BRAKET_JOB_RESULTS_DIR": tmp_dir}):
            save_job_result(first_result_data, first_data_format)
            save_job_result(second_result_data, second_data_format)

            assert load_job_result() == expected_result_data


def test_update_pickled_results_as_plaintext_error():
    with tempfile.TemporaryDirectory() as tmp_dir:
        with patch.dict(os.environ, {"AMZN_BRAKET_JOB_RESULTS_DIR": tmp_dir}):
            save_job_result(np.arange(5), PersistedJobDataFormat.PICKLED_V4)

            cannot_convert_pickled_to_plaintext = (
                "Cannot update results object serialized with "
                "pickled_v4 using data format plaintext."
            )
            with pytest.raises(TypeError, match=cannot_convert_pickled_to_plaintext):
                save_job_result("hello", PersistedJobDataFormat.PLAINTEXT)
