# from lionagi.core.branch.base_branch import BaseBranch
# from lionagi.core.branch.util import MessageUtil
# from lionagi.core.messages.schema import System
# from lionagi.core.schema import DataLogger


# import unittest
# from unittest.mock import patch, MagicMock
# import pandas as pd
# from datetime import datetime
# import json


# class TestBaseBranch(unittest.TestCase):

#     def setUp(self):
#         # Patching DataLogger to avoid filesystem interactions during tests
#         self.patcher = patch(
#             "lionagi.core.branch.base_branch.DataLogger", autospec=True
#         )
#         self.MockDataLogger = self.patcher.start()
#         self.addCleanup(self.patcher.stop)

#         self.branch = BaseBranch()
#         # Example messages to populate the DataFrame for testing
#         self.test_messages = [
#             {
#                 "node_id": "1",
#                 "timestamp": "2021-01-01 00:00:00",
#                 "role": "system",
#                 "sender": "system",
#                 "content": json.dumps({"system_info": "System message"}),
#             },
#             {
#                 "node_id": "2",
#                 "timestamp": "2021-01-01 00:01:00",
#                 "role": "user",
#                 "sender": "user1",
#                 "content": json.dumps({"instruction": "User message"}),
#             },
#             {
#                 "node_id": "3",
#                 "timestamp": "2021-01-01 00:02:00",
#                 "role": "assistant",
#                 "sender": "assistant",
#                 "content": json.dumps({"response": "Assistant response"}),
#             },
#             {
#                 "node_id": "4",
#                 "timestamp": "2021-01-01 00:03:00",
#                 "role": "assistant",
#                 "sender": "action_request",
#                 "content": json.dumps({"action_request": "Action request"}),
#             },
#             {
#                 "node_id": "5",
#                 "timestamp": "2021-01-01 00:04:00",
#                 "role": "assistant",
#                 "sender": "action_response",
#                 "content": json.dumps({"action_response": "Action response"}),
#             },
#         ]
#         self.branch = BaseBranch(messages=pd.DataFrame(self.test_messages))

#     def test_init_with_empty_messages(self):
#         """Test __init__ method with no messages provided."""
#         branch = BaseBranch()
#         self.assertTrue(branch.messages.empty)

#     def test_init_with_given_messages(self):
#         """Test __init__ method with provided messages DataFrame."""
#         messages = pd.DataFrame(
#             [
#                 [
#                     "0",
#                     datetime(2024, 1, 1),
#                     "system",
#                     "system",
#                     json.dumps({"system_info": "Hi"}),
#                 ]
#             ],
#             columns=["node_id", "timestamp", "role", "sender", "content"],
#         )
#         branch = BaseBranch(messages=messages)
#         self.assertFalse(branch.messages.empty)

#     def test_add_message(self):
#         """Test adding a message."""
#         message_info = {"info": "Test Message"}
#         with patch.object(
#             MessageUtil, "create_message", return_value=System(system=message_info)
#         ) as mocked_create_message:
#             self.branch.add_message(system=message_info)
#             mocked_create_message.assert_called_once()
#             self.assertEqual(len(self.branch.messages), 6)

#     def test_chat_messages_without_sender(self):
#         """Test chat_messages property without including sender information."""
#         chat_messages = self.branch.chat_messages
#         self.assertEqual(len(chat_messages), 5)
#         self.assertNotIn("Sender", chat_messages[0]["content"])

#     def test_chat_messages_with_sender(self):
#         """Test retrieving chat messages with sender information included."""
#         chat_messages_with_sender = self.branch.chat_messages_with_sender
#         expected_prefixes = [
#             "Sender system: ",
#             "Sender user1: ",
#             "Sender assistant: ",
#             "Sender action_request: ",
#             "Sender action_response: ",
#         ]

#         # Verify the number of messages returned matches the number added
#         self.assertEqual(len(chat_messages_with_sender), len(self.test_messages))

#         # Check each message for correct sender prefix in content
#         for i, message_dict in enumerate(chat_messages_with_sender):
#             self.assertTrue(
#                 message_dict["content"].startswith(expected_prefixes[i]),
#                 msg=f"Message content does not start with expected sender prefix. Found: {message_dict['content']}",
#             )

#         # Optionally, verify the content matches after removing the prefix
#         for i, message_dict in enumerate(chat_messages_with_sender):
#             prefix, content = message_dict["content"].split(": ", 1)
#             self.assertEqual(
#                 content,
#                 self.test_messages[i]["content"],
#                 msg=f"Message content does not match after removing sender prefix. Found: {content}, Expected: {self.test_messages[i]['content']}",
#             )

#     def test_last_message(self):
#         """Test retrieving the last message."""
#         last_message = self.branch.last_message
#         self.assertEqual(
#             json.loads(last_message.iloc[0]["content"]),
#             {"action_response": "Action response"},
#         )

#     def test_last_message_content(self):
#         """Test retrieving the content of the last message."""
#         content = self.branch.last_message_content
#         self.assertEqual(content, {"action_response": "Action response"})

#     def test_first_system_message(self):
#         """Test retrieving the first 'system' message."""
#         first_system_message = self.branch.first_system
#         self.assertEqual(
#             json.loads(first_system_message.iloc[0]["content"]),
#             {"system_info": "System message"},
#         )

#     def test_last_response(self):
#         """Test retrieving the last 'assistant' message."""
#         last_response = self.branch.last_response
#         self.assertEqual(last_response.iloc[0]["sender"], "action_response")

#     def test_last_response_content(self):
#         """Test extracting content of the last 'assistant' message."""
#         content = self.branch.last_response_content
#         self.assertEqual(content, {"action_response": "Action response"})

#     def test_action_request_messages(self):
#         """Test filtering messages sent by 'action_request'."""
#         action_requests = self.branch.action_request
#         self.assertTrue(all(action_requests["sender"] == "action_request"))

#     def test_action_response_messages(self):
#         """Test filtering messages sent by 'action_response'."""
#         action_responses = self.branch.action_response
#         self.assertTrue(all(action_responses["sender"] == "action_response"))

#     def test_responses(self):
#         """Test filtering of 'assistant' role messages."""
#         responses = self.branch.responses
#         # Verify that all returned messages have the 'assistant' role
#         self.assertTrue(all(responses["role"] == "assistant"))
#         # Optionally, check the count matches the expected number of 'assistant' messages
#         expected_count = sum(
#             1 for msg in self.test_messages if msg["role"] == "assistant"
#         )
#         self.assertEqual(len(responses), expected_count)

#     def test_assistant_responses(self):
#         """Test filtering of 'assistant' messages excluding action requests/responses."""
#         assistant_responses = self.branch.assistant_responses
#         # Verify that no returned messages are from 'action_request' or 'action_response' senders
#         self.assertTrue(all(assistant_responses["sender"] != "action_request"))
#         self.assertTrue(all(assistant_responses["sender"] != "action_response"))
#         # Verify that all returned messages have the 'assistant' role
#         self.assertTrue(all(assistant_responses["role"] == "assistant"))

#     def test_info(self):
#         """Test summarization of message counts by role."""
#         info = self.branch.info
#         # Verify that the dictionary contains keys for each role and 'total'
#         self.assertIn("assistant", info)
#         self.assertIn("user", info)
#         self.assertIn("system", info)
#         self.assertIn("total", info)
#         # Optionally, verify the counts match expected values
#         for role in ["assistant", "user", "system"]:
#             expected_count = sum(1 for msg in self.test_messages if msg["role"] == role)
#             self.assertEqual(info[role], expected_count)
#         self.assertEqual(info["total"], len(self.test_messages))

#     def test_sender_info(self):
#         """Test summarization of message counts by sender."""
#         sender_info = self.branch.sender_info
#         # Verify that the dictionary contains keys for each sender and the counts match
#         for sender in set(msg["sender"] for msg in self.test_messages):
#             expected_count = sum(
#                 1 for msg in self.test_messages if msg["sender"] == sender
#             )
#             self.assertEqual(sender_info.get(sender, 0), expected_count)

#     def test_describe(self):
#         """Test detailed description of the branch."""
#         description = self.branch.describe
#         # Verify that the description contains expected keys
#         self.assertIn("total_messages", description)
#         self.assertIn("summary_by_role", description)
#         self.assertIn("messages", description)
#         # Optionally, verify the accuracy of the values
#         self.assertEqual(description["total_messages"], len(self.test_messages))
#         self.assertEqual(len(description["messages"]), min(5, len(self.test_messages)))

#     @patch("lionagi.libs.ln_dataframe.read_csv")
#     def test_from_csv(cls, mock_read_csv):
#         # Define a mock return value for read_csv
#         mock_messages_df = pd.DataFrame(
#             {
#                 "node_id": ["1", "2"],
#                 "timestamp": [datetime(2021, 1, 1), datetime(2021, 1, 1)],
#                 "role": ["system", "user"],
#                 "sender": ["system", "user1"],
#                 "content": [
#                     json.dumps({"system_info": "System message"}),
#                     json.dumps({"instruction": "User message"}),
#                 ],
#             }
#         )
#         mock_read_csv.return_value = mock_messages_df

#         # Call the from_csv class method
#         branch = BaseBranch.from_csv(filename="dummy.csv")

#         # Verify that read_csv was called correctly
#         mock_read_csv.assert_called_once_with("dummy.csv")

#         # Verify that the branch instance contains the correct messages
#         pd.testing.assert_frame_equal(branch.messages, mock_messages_df)

#     @patch("lionagi.libs.ln_dataframe.read_json")
#     def test_from_json(cls, mock_read_csv):
#         # Define a mock return value for read_csv
#         mock_messages_df = pd.DataFrame(
#             {
#                 "node_id": ["1", "2"],
#                 "timestamp": [datetime(2021, 1, 1), datetime(2021, 1, 1)],
#                 "role": ["system", "user"],
#                 "sender": ["system", "user1"],
#                 "content": [
#                     json.dumps({"system_info": "System message"}),
#                     json.dumps({"instruction": "User message"}),
#                 ],
#             }
#         )
#         mock_read_csv.return_value = mock_messages_df

#         # Call the from_csv class method
#         branch = BaseBranch.from_json_string(filename="dummy.json")

#         # Verify that read_csv was called correctly
#         mock_read_csv.assert_called_once_with("dummy.json")

#         # Verify that the branch instance contains the correct messages
#         pd.testing.assert_frame_equal(branch.messages, mock_messages_df)

#     @patch(
#         "lionagi.libs.sys_util.SysUtil.create_path",
#         return_value="path/to/messages.csv",
#     )
#     @patch.object(pd.DataFrame, "to_csv")
#     def test_to_csv(self, mock_to_csv, mock_create_path):
#         self.branch.datalogger = MagicMock()
#         self.branch.datalogger.persist_path = "data/logs/"

#         self.branch.to_csv_file("messages.csv", verbose=False, clear=False)

#         mock_create_path.assert_called_once_with(
#             self.branch.datalogger.persist_path,
#             "messages.csv",
#             timestamp=True,
#             dir_exist_ok=True,
#             time_prefix=False,
#         )
#         mock_to_csv.assert_called_once_with("path/to/messages.csv")

#         # Verify that messages are not cleared after exporting
#         assert not self.branch.messages.empty

#     @patch(
#         "lionagi.libs.sys_util.SysUtil.create_path",
#         return_value="path/to/messages.json",
#     )
#     @patch.object(pd.DataFrame, "to_json")
#     def test_to_json(self, mock_to_json, mock_create_path):
#         self.branch.datalogger = MagicMock()
#         self.branch.datalogger.persist_path = "data/logs/"

#         self.branch.to_json_file("messages.json", verbose=False, clear=False)

#         mock_create_path.assert_called_once_with(
#             self.branch.datalogger.persist_path,
#             "messages.json",
#             timestamp=True,
#             dir_exist_ok=True,
#             time_prefix=False,
#         )
#         mock_to_json.assert_called_once_with(
#             "path/to/messages.json", orient="records", lines=True, date_format="iso"
#         )

#         # Verify that messages are not cleared after exporting
#         assert not self.branch.messages.empty

#     # def test_log_to_csv(self):
#     #     self.branch.log_to_csv('log.csv', verbose=False, clear=False)

#     #     self.branch.datalogger.to_csv_file.assert_called_once_with(filename='log.csv', dir_exist_ok=True, timestamp=True,
#     #                                             time_prefix=False, verbose=False, clear=False)

#     # def test_log_to_json(self):
#     #     branch = BaseBranch()
#     #     branch.log_to_json('log.json', verbose=False, clear=False)

#     #     self.branch.datalogger.to_json_file.assert_called_once_with(filename='log.json', dir_exist_ok=True, timestamp=True,
#     #                                              time_prefix=False, verbose=False, clear=False)

#     def test_remove_message(self):
#         """Test removing a message from the branch based on its node ID."""
#         initial_length = len(self.branch.messages)
#         node_id_to_remove = "2"
#         self.branch.remove_message(node_id_to_remove)
#         final_length = len(self.branch.messages)
#         self.assertNotIn(node_id_to_remove, self.branch.messages["node_id"].tolist())
#         self.assertEqual(final_length, initial_length - 1)

#     def test_update_message(self):
#         """Test updating a specific message's column identified by node_id with a new value."""
#         node_id_to_update = "2"
#         new_value = "Updated content"
#         self.branch.update_message(node_id_to_update, "content", new_value)
#         updated_value = self.branch.messages.loc[
#             self.branch.messages["node_id"] == node_id_to_update, "content"
#         ].values[0]
#         self.assertEqual(updated_value, new_value)

#     # def test_change_first_system_message(self):
#     #     """Test updating the first system message with new content and/or sender."""
#     #     new_system_content = {"system_info": "Updated system message"}
#     #     self.branch.change_first_system_message(new_system_content['system_info'])
#     #     first_system_message_content = self.branch.messages.loc[self.branch.messages['role'] == 'system', 'content'].iloc[0]
#     #     self.assertIn(json.dumps({"system_info": "Updated system message"}), first_system_message_content)

#     def test_rollback(self):
#         """Test removing the last 'n' messages from the branch."""
#         steps_to_remove = 2
#         initial_length = len(self.branch.messages)
#         self.branch.rollback(steps_to_remove)
#         final_length = len(self.branch.messages)
#         self.assertEqual(final_length, initial_length - steps_to_remove)

#     def test_clear_messages(self):
#         """Test clearing all messages from the branch."""
#         self.branch.clear_messages()
#         self.assertTrue(self.branch.messages.empty)

#     def test_replace_keyword(self):
#         """Test replacing occurrences of a specified keyword with a replacement string."""
#         keyword = "Assistant response"
#         replacement = "Helper feedback"
#         self.branch.replace_keyword(keyword, replacement)
#         self.assertTrue(
#             any(replacement in message for message in self.branch.messages["content"])
#         )

#     def test_search_keywords(self):
#         """Test filtering messages by a specified keyword or list of keywords."""
#         keyword_to_search = "Assistant"
#         filtered_messages = self.branch.search_keywords(keyword_to_search)
#         self.assertTrue(
#             all(
#                 keyword_to_search in content for content in filtered_messages["content"]
#             )
#         )

#     def test_extend(self):
#         """Test extending branch messages with additional messages."""
#         additional_messages = pd.DataFrame(
#             [
#                 {
#                     "node_id": "6",
#                     "timestamp": datetime(2021, 1, 1),
#                     "role": "user",
#                     "sender": "user2",
#                     "content": json.dumps({"instruction": "Another user message"}),
#                 }
#             ]
#         )
#         initial_length = len(self.branch.messages)
#         self.branch.extend(additional_messages)
#         self.assertEqual(len(self.branch.messages), initial_length + 1)

#     def test_filter_by(self):
#         """Test filtering messages by various criteria."""
#         # Filtering by role
#         filtered_messages = self.branch.filter_by(role="user")
#         self.assertTrue(
#             all(msg["role"] == "user" for _, msg in filtered_messages.iterrows())
#         )


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