import pytest
from unittest.mock import patch, MagicMock
import oci
from src.providers.llm_provider import chat_llm_with_usage, classification_llm

@patch("oci.generative_ai_inference.GenerativeAiInferenceClient.chat")
def test_chat_llm_retry_on_timeout(mock_chat):
    """
    Testa se o chat_llm_with_usage executa retries quando recebe um ServiceError (504).
    Simula 2 falhas seguidas e sucesso na 3ª tentativa.
    """
    # ── Montagem do Cenário ──────────────────────────────────────────────────
    # Criamos o erro que o retry deve capturar
    timeout_error = oci.exceptions.ServiceError(
        status=504,
        code="GatewayTimeout",
        message="Request timed out",
        headers={},
        operation_name="Chat"
    )
    
    # Simula resposta de sucesso na 3ª chamada
    mock_response = MagicMock()
    chat_response = mock_response.data.chat_response
    
    # Configuramos os valores de retorno para serem tipos primitivos (não Mocks)
    chat_response.choices = [
        MagicMock(
            message=MagicMock(content=[MagicMock(text="Sucesso após retry!")]),
            finish_reason="stop"
        )
    ]
    chat_response.model = "cohere.command-r-plus"
    chat_response.usage.prompt_tokens = 10
    chat_response.usage.completion_tokens = 10
    chat_response.usage.total_tokens = 20
    
    # Configuramos os side_effects: Falha, Falha, Sucesso
    mock_chat.side_effect = [timeout_error, timeout_error, mock_response]

    # ── Execução ─────────────────────────────────────────────────────────────
    prompt = "Teste de retry"
    result = chat_llm_with_usage(classification_llm, prompt)

    # ── Verificação ──────────────────────────────────────────────────────────
    # O resultado final deve ser o da 3ª chamada
    assert result.content == "Sucesso após retry!"
    
    # Verificamos se o client.chat foi chamado exatamente 3 vezes
    assert mock_chat.call_count == 3