import unittest
from unittest import mock
from chatmastermind.ais.openai import OpenAI
from chatmastermind.message import Message, Question, Answer
from chatmastermind.chat import Chat
from chatmastermind.ai import AIResponse, Tokens
from chatmastermind.configuration import OpenAIConfig


class OpenAITest(unittest.TestCase):

    @mock.patch('openai.ChatCompletion.create')
    def test_request(self, mock_create: mock.MagicMock) -> None:
        # Create a test instance of OpenAI
        config = OpenAIConfig()
        openai = OpenAI(config)

        # Set up the mock response from openai.ChatCompletion.create
        mock_chunk1 = {
            'choices': [
                {
                    'index': 0,
                    'delta': {
                        'content': 'Answer 1'
                    },
                    'finish_reason': None
                },
                {
                    'index': 1,
                    'delta': {
                        'content': 'Answer 2'
                    },
                    'finish_reason': None
                }
            ],
        }
        mock_chunk2 = {
            'choices': [
                {
                    'index': 0,
                    'finish_reason': 'stop'
                },
                {
                    'index': 1,
                    'finish_reason': 'stop'
                }
            ],
        }
        mock_create.return_value = iter([mock_chunk1, mock_chunk2])

        # Create test data
        question = Message(Question('Question'))
        chat = Chat([
            Message(Question('Question 1'), answer=Answer('Answer 1')),
            Message(Question('Question 2'), answer=Answer('Answer 2')),
            # add message without an answer -> expect to be skipped
            Message(Question('Question 3'))
        ])

        # Make the request
        response = openai.request(question, chat, num_answers=2)

        # Assert the AIResponse
        self.assertIsInstance(response, AIResponse)
        self.assertEqual(len(response.messages), 2)
        self.assertEqual(response.messages[0].answer, 'Answer 1')
        self.assertEqual(response.messages[1].answer, 'Answer 2')
        self.assertIsNotNone(response.tokens)
        self.assertIsInstance(response.tokens, Tokens)
        assert response.tokens
        self.assertEqual(response.tokens.prompt, 53)
        self.assertEqual(response.tokens.completion, 6)
        self.assertEqual(response.tokens.total, 59)

        # Assert the mock call to openai.ChatCompletion.create
        mock_create.assert_called_once_with(
            model=f'{config.model}',
            messages=[
                {'role': 'system', 'content': f'{config.system}'},
                {'role': 'user', 'content': 'Question 1'},
                {'role': 'assistant', 'content': 'Answer 1'},
                {'role': 'user', 'content': 'Question 2'},
                {'role': 'assistant', 'content': 'Answer 2'},
                {'role': 'user', 'content': 'Question'}
            ],
            temperature=config.temperature,
            max_tokens=config.max_tokens,
            top_p=config.top_p,
            n=2,
            stream=True,
            frequency_penalty=config.frequency_penalty,
            presence_penalty=config.presence_penalty
        )