import unittest
import asyncio
from unittest.mock import patch, AsyncMock, MagicMock

# Import the class to be tested
from app.core.providers.llm.general import GeneralProvider

# A mock class to simulate the LiteLLM completion response object
# It mimics the nested structure: response.choices[0].message.content
class MockResponse:
    def __init__(self, content):
        self.choices = [self.MockChoice(content)]
        
    class MockChoice:
        def __init__(self, content):
            self.message = self.MockMessage(content)

        class MockMessage:
            def __init__(self, content):
                self.content = content

class TestGeneralProvider(unittest.TestCase):
    def setUp(self):
        """Set up a new GeneralProvider instance before each test."""
        self.model_name = "gemini/gemini-1.5-pro-latest"
        self.api_key = "test-api-key"
        self.provider = GeneralProvider(self.model_name, self.api_key)
        self.prompt = "What is the capital of France?"

    @patch('litellm.acompletion', new_callable=AsyncMock)
    def test_generate_response_success(self, mock_litellm_completion: AsyncMock):
        """
        Test that aforward returns the correct content on a successful call.
        """
        async def run_test():
            # Arrange: Configure the mock to return a successful response
            expected_content = "The capital of France is Paris."
            mock_litellm_completion.return_value = MockResponse(expected_content)

            # Act: Call the async method
            response_obj = await self.provider.aforward(prompt=self.prompt)
            response = response_obj.choices[0].message.content

            # Assert: Check the returned value and the mock's call arguments
            self.assertEqual(response, expected_content)
            
            # Assert that the mock was called with the correct parameters, including defaults from BaseLM
            mock_litellm_completion.assert_called_once_with(
                model=self.model_name,
                messages=[{"role": "user", "content": self.prompt}],
                api_key=self.api_key,
                temperature=0.0, # <-- Add this line
                max_tokens=10000000, # <-- Add this line
            )
        
        # Run the async test function
        asyncio.run(run_test())

    @patch('litellm.acompletion', new_callable=AsyncMock)
    def test_generate_response_failure(self, mock_litellm_completion: AsyncMock):
        """
        Test that aforward raises a RuntimeError on an API failure.
        """
        async def run_test():
            # Arrange: Configure the mock to raise an exception
            mock_litellm_completion.side_effect = Exception("API connection failed.")

            # Act & Assert: Use assertRaises to check for the RuntimeError
            with self.assertRaisesRegex(RuntimeError, f"Failed to get response from LiteLLM for model '{self.model_name}': API connection failed."):
                await self.provider.aforward(prompt=self.prompt)
        
        # Run the async test function
        asyncio.run(run_test())
