import pytest
from typing import List
import asyncio
from unittest.mock import MagicMock, AsyncMock
from sqlalchemy.orm import Session

# Import the pipeline
from app.core.pipelines.rag_pipeline import RagPipeline
from app.db import models # Import your SQLAlchemy models for mocking history
from app.core.retrievers.base_retriever import Retriever


# --- Mock Classes ---

class MockRetriever(Retriever):
    """A mock retriever that returns a predefined list of strings."""
    def __init__(self, name: str, mock_data: List[str]):
        self.name = name
        self.mock_data = mock_data

    def retrieve_context(self, question: str, db: Session) -> List[str]:
        return self.mock_data

# --- Fixtures ---

@pytest.fixture
def mock_db():
    """A mock SQLAlchemy session."""
    return MagicMock()

@pytest.fixture
def mock_llm_provider():
    """
    Mocks the llm_provider that the pipeline now uses.
    """
    mock_provider = MagicMock()
    mock_completion = MagicMock()
    mock_completion.choices = [MagicMock()]
    mock_completion.choices[0].message.content = "Mocked LLM answer."
    mock_provider.acompletion = AsyncMock(return_value=mock_completion)
    return mock_provider


# --- Test Cases ---

@pytest.mark.asyncio
async def test_forward_pass_with_defaults(mock_db, mock_llm_provider):
    """Test a successful forward pass using default processors."""
    pipeline = RagPipeline()
    
    question = "What is the capital of France?"
    history = [models.Message(sender="user", content="Hello there."), models.Message(sender="assistant", content="Hi.")]
    context_chunks = ["Context 1.", "Context 2."]
    
    response = await pipeline.forward(
        question=question, 
        history=history, 
        context_chunks=context_chunks,
        llm_provider=mock_llm_provider
    )
    
    # Assert that the completion method was called
    mock_llm_provider.acompletion.assert_called_once()
    assert response == "Mocked LLM answer."

@pytest.mark.asyncio
async def test_forward_with_custom_processors(mock_db, mock_llm_provider):
    """Test that custom processors are used correctly."""
    
    def custom_context_processor(contexts: List[str]) -> str:
        return "CUSTOM_CONTEXT: " + " | ".join(contexts)

    def custom_history_formatter(history: List[models.Message]) -> str:
        return "<HISTORY> " + " ".join([m.content for m in history]) + " </HISTORY>"

    def custom_response_processor(response: str) -> str:
        return f"FINAL: {response.upper()}"

    pipeline = RagPipeline(
        context_postprocessor=custom_context_processor,
        history_formatter=custom_history_formatter,
        response_postprocessor=custom_response_processor
    )

    question = "Custom question?"
    history = [models.Message(sender="user", content="User message.")]
    context_chunks = ["Context A", "Context B"]
    
    response = await pipeline.forward(
        question=question, 
        history=history, 
        context_chunks=context_chunks,
        llm_provider=mock_llm_provider
    )

    mock_llm_provider.acompletion.assert_called_once()
    assert response == "FINAL: MOCKED LLM ANSWER."

@pytest.mark.asyncio
async def test_empty_context_and_history_handling(mock_db, mock_llm_provider):
    """Test behavior with empty context and chat history."""
    pipeline = RagPipeline()
    
    question = "No context question."
    history = []
    context_chunks = []
    
    response = await pipeline.forward(
        question=question, 
        history=history, 
        context_chunks=context_chunks,
        llm_provider=mock_llm_provider
    )
    
    mock_llm_provider.acompletion.assert_called_once()
    assert response == "Mocked LLM answer."