from lionagi.core.session.base.base_branch import BaseBranch
from lionagi.core.session.base.util import MessageUtil, System
from lionagi.core.schema import DataLogger

import unittest
from unittest.mock import patch, MagicMock, create_autospec
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.session.base.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.core.session.base.base_branch.MessageUtil.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.core.session.base.base_branch.MessageUtil.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(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.core.session.base.base_branch.PathUtil.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.core.session.base.base_branch.PathUtil.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()
