# Copyright 2022 Soul Machines Ltd

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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 unittest
from typing import List

from smskillsdk.models.api import Memory
from smskillsdk.utils import memory


class TestMemoryUtils(unittest.TestCase):
    def test_serialize(self):
        memory_dict = {
            "key1": "value1",
            "key2": "value2",
            "key3": "value3",
        }

        memory_array = memory.serialize_memory(memory_dict)
        self.assertEqual(len(memory_array), 3)

        for i, memory_element in enumerate(memory_array):
            self.assertTrue(
                hasattr(memory_element, "name"),
                f"Memory element {i} is missing 'name' attribute",
            )
            self.assertTrue(
                hasattr(memory_element, "value"),
                f"Memory element {i} is missing 'value' attribute",
            )
            self.assertFalse(hasattr(memory_element, "session_id"))
            self.assertIn(memory_element.name, memory_dict)
            self.assertEqual(memory_element.value, memory_dict[memory_element.name])

    def test_serialize_with_session_id(self):
        memory_dict = {
            "key1": "value1",
            "key2": "value2",
            "key3": "value3",
        }
        session_id = "123"

        memory_array = memory.serialize_memory(memory_dict, session_id)
        self.assertEqual(len(memory_array), 3)

        for i, memory_element in enumerate(memory_array):
            self.assertTrue(
                hasattr(memory_element, "name"),
                f"Memory element {i} is missing 'name' attribute",
            )
            self.assertTrue(
                hasattr(memory_element, "value"),
                f"Memory element {i} is missing 'value' attribute",
            )
            self.assertTrue(hasattr(memory_element, "session_id"))
            self.assertIn(memory_element.name, memory_dict)
            self.assertEqual(memory_element.value, memory_dict[memory_element.name])
            self.assertEqual(memory_element.session_id, session_id)

    def test_deserialize(self):
        memory_array = [
            Memory(**{"name": "key1", "value": "value1"}),
            Memory(**{"name": "key2", "value": 1}),
            Memory(**{"name": "key3", "value": {"subkey2": "subvalue2"}}),
            Memory(**{"key": "key4", "value": "value"}),
        ]

        memory_dict = memory.deserialize_memory(memory_array)
        self.assertEqual(len(memory_dict), 3)

        for memory_element in memory_array:
            if hasattr(memory_element, "name") and hasattr(memory_element, "value"):
                self.assertIn(memory_element.name, memory_dict)
                self.assertEqual(memory_element.value, memory_dict[memory_element.name])

    def test_deserialize_with_session_id(self):
        memory_array = [
            Memory(**{"name": "key1", "value": "value1", "session_id": "123"}),
            Memory(**{"name": "key2", "value": 1}),
            Memory(
                **{
                    "name": "key3",
                    "value": {"subkey2": "subvalue2"},
                    "session_id": "123",
                }
            ),
        ]

        memory_dict = memory.deserialize_memory(memory_array, "123")
        self.assertEqual(set(memory_dict.keys()), {"key1", "key3"})

    def test_set_memory_value(self):
        memories: List[Memory] = list()

        memory.set_memory_value(memories, "key1", "value1")
        self.assertEqual(len(memories), 1)
        self.assertEqual(memories[0].name, "key1")
        self.assertEqual(memories[0].value, "value1")

        memory.set_memory_value(memories, "key1", "value1", "sessionid1")
        self.assertEqual(len(memories), 2)
        self.assertEqual(memories[-1].name, "key1")
        self.assertEqual(memories[-1].value, "value1")
        self.assertFalse(hasattr(memories[0], "session_id"))
        self.assertTrue(hasattr(memories[-1], "session_id"))

        memory.set_memory_value(memories, "key1", "value2")
        self.assertEqual(len(memories), 2)
        self.assertEqual(memories[0].name, "key1")
        self.assertEqual(memories[0].value, "value2")

        memory.set_memory_value(memories, "key1", "value3", "sessionid1")
        self.assertEqual(len(memories), 2)
        self.assertEqual(memories[0].name, "key1")
        self.assertEqual(memories[0].value, "value2")
        self.assertEqual(memories[-1].name, "key1")
        self.assertEqual(memories[-1].value, "value3")

        memory.set_memory_value(memories, "key2", "value2")
        self.assertEqual(len(memories), 3)
        self.assertEqual(memories[-1].name, "key2")
        self.assertEqual(memories[-1].value, "value2")

    def test_get_memory_value(self):
        memories: List[Memory] = list()

        self.assertEqual(memory.get_memory_value(memories, "key1")[0], False)

        memories.append(Memory(**{"name": "key1", "value": "value1"}))
        self.assertEqual(memory.get_memory_value(memories, "key1"), (True, "value1"))

        memories.append(
            Memory(**{"name": "key1", "value": "value2", "session_id": "sessionid1"})
        )
        self.assertEqual(memory.get_memory_value(memories, "key1"), (True, "value1"))
        self.assertEqual(
            memory.get_memory_value(memories, "key1", "sessionid1"), (True, "value2")
        )
        self.assertEqual(
            memory.get_memory_value(memories, "key1", "sessionid2")[0], False
        )


if __name__ == "__main__":
    unittest.main()
