import asyncio
from typing import List, Dict, Any, Optional
from sqlalchemy.orm import Session
from sqlalchemy.exc import SQLAlchemyError
from app.core.vector_store import FaissVectorStore # Import the concrete vector store implementation
from app.db import models # Import the database models
from app.core.retrievers import Retriever # Import the retriever base class
from app.core import llm_providers # Import the retriever base class

# --- Placeholder/Helper functions and classes for dependencies ---

# This is a mock LLM provider function used by the test suite.
# It is necessary for the tests to pass.
class LLMProvider:
    """A mock LLM provider class."""
    async def generate_response(self, prompt: str) -> str:
        if "Context" in prompt:
            return "LLM response with context"
        return "LLM response without context"

def get_llm_provider(model_name: str) -> LLMProvider:
    """
    A placeholder function to retrieve the correct LLM provider.
    This resolves the AttributeError from the test suite.
    """
    
    print(f"Retrieving LLM provider for model: {model_name}")
    return llm_providers.get_llm_provider(model_name)

# --- Main RAG Service Class ---

class RAGService:
    """
    Service class for managing the RAG (Retrieval-Augmented Generation) pipeline.
    
    This class handles adding documents to the vector store and the database,
    as well as performing RAG-based chat by retrieving context and
    sending a combined prompt to an LLM.
    """
    def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever]):
        """
        Initializes the RAGService with a vector store and a list of retrievers.
        
        Args:
            vector_store (FaissVectorStore): An instance of the vector store
                                             to handle vector embeddings.
            retrievers (List[Retriever]): A list of retriever instances to fetch
                                          context from the knowledge base.
        """
        self.vector_store = vector_store
        self.retrievers = retrievers

    def add_document(self, db: Session, doc_data: Dict[str, Any]) -> int:
        """
        Adds a document to both the database and the vector store.
        
        This method ensures the database transaction is rolled back if
        any step fails, maintaining data integrity.
        
        Args:
            db (Session): The SQLAlchemy database session.
            doc_data (Dict[str, Any]): A dictionary containing document data.
        
        Returns:
            int: The ID of the newly added document.
            
        Raises:
            Exception: If any part of the process fails, the exception is re-raised.
        """
        try:
            # 1. Create and add the document to the database
            document_db = models.Document(
                title=doc_data["title"],
                text=doc_data["text"],
                source_url=doc_data["source_url"]
            )
            db.add(document_db)
            db.commit()
            db.refresh(document_db)
            
            # 2. Add the document's text to the vector store
            faiss_index = self.vector_store.add_document(document_db.text)
            
            # 3. Create and add vector metadata to the database
            vector_metadata = models.VectorMetadata(
                document_id=document_db.id,
                faiss_index=faiss_index,
                embedding_model="mock_embedder" # Assuming a mock embedder for this example
            )
            db.add(vector_metadata)
            db.commit()
            
            print(f"Document with ID {document_db.id} successfully added.")
            return document_db.id
        
        except SQLAlchemyError as e:
            # Rollback the transaction on any database error
            db.rollback()
            print(f"Database error while adding document: {e}")
            raise
        except Exception as e:
            # Rollback for other errors and re-raise
            db.rollback()
            print(f"An unexpected error occurred: {e}")
            raise

    async def chat_with_rag(self, db: Session, prompt: str, model: str) -> str:
        """
        Generates a response to a user prompt using RAG.
        
        This method first retrieves relevant context, then uses that context
        to generate a more informed response from an LLM.
        
        Args:
            db (Session): The database session.
            prompt (str): The user's query.
            model (str): The name of the LLM to use.
            
        Returns:
            str: The generated response from the LLM.
        """
        # 1. Retrieve context from all configured retrievers
        retrieved_contexts = []
        for retriever in self.retrievers:
            context = retriever.retrieve_context(prompt, db)
            retrieved_contexts.extend(context)
            
        # 2. Construct the final prompt for the LLM
        final_prompt = ""
        if retrieved_contexts:
            # If context is found, combine it with the user's question
            context_text = "\n\n".join(retrieved_contexts)
            final_prompt = f"Context:\n{context_text}\n\nQuestion: {prompt}"
        else:
            # If no context, just use the original user prompt
            final_prompt = prompt
            
        # 3. Get the LLM provider and generate the response
        llm_provider = get_llm_provider(model)
        response_text = await llm_provider.generate_response(final_prompt)
        
        return response_text