"""
RAG Query Module
Handles querying the RAG system and extracting answers with sources
"""
from typing import Tuple, Optional, List
from openai import OpenAI


class RAGQueryEngine:
    """Handles RAG queries with source attribution"""
    
    def __init__(self, client: OpenAI, vector_store_id: str, model: str = "gpt-4o-mini"):
        self.client = client
        self.vector_store_id = vector_store_id
        self.model = model
    
    def get_response_from_vectorstore(self, query: str):
        """
        Get response from vector store using OpenAI responses API
        
        Args:
            query: User query
            
        Returns:
            Response object or None if failed
        """
        try:
            response = self.client.responses.create(
                input=query,
                model=self.model,
                tools=[{
                    "type": "file_search",
                    "vector_store_ids": [self.vector_store_id],
                }]
            )
            
            # Check if response is valid
            if response and hasattr(response.output[1], 'content'):
                return response
            else:
                print("⚠️ Invalid response structure")
                return None
        except Exception as e:
            print(f"❌ Error during API call: {e}")
            return None
    
    def query(self, query: str) -> Tuple[str, str]:
        """
        Query the RAG model and return answer with sources
        
        Args:
            query: User query
            
        Returns:
            Tuple of (answer_text, footnotes)
        """
        response = self.get_response_from_vectorstore(query)
        
        if not response:
            return "That question is outside my area of expertise.", ""
        
        # Extract the answer text
        answer_text = response.output[1].content[0].text
        
        # Extract the source files
        footnotes = ""
        if hasattr(response.output[1].content[0], 'annotations'):
            annotations = response.output[1].content[0].annotations
            if annotations:
                # Get unique source files
                source_files = list(set([result.filename for result in annotations]))
                
                # Format the footnotes
                footnotes = "\n\n📚 **Sources:**\n"
                for i, filename in enumerate(source_files, 1):
                    # Remove the ".pdf" extension and format nicely
                    clean_name = filename.replace('.pdf', '')
                    footnotes += f"{i}. {clean_name}\n"
        
        return answer_text, footnotes
    
    def get_files_from_vector_store(self) -> List[str]:
        """
        Get list of files in the vector store
        
        Returns:
            List of filenames
        """
        try:
            query = "List all documents about Mercedes E-class ADAS features"
            response = self.get_response_from_vectorstore(query)
            
            file_list = []
            if response and hasattr(response.output[1].content[0], 'annotations'):
                annotations = response.output[1].content[0].annotations
                file_list = list(set([annotation.filename for annotation in annotations]))
                file_list.sort()
            
            # Fallback to default list if empty
            if not file_list:
                file_list = [
                    "Function of Active Distance Assist DISTRONIC.pdf",
                    "Function of Active Lane Change Assist.pdf",
                    "Function of Active Steering Assist.pdf",
                    "Function of Active Stop-and-Go Assist.pdf"
                ]
            
            return file_list
        except Exception as e:
            print(f"❌ Error getting files: {str(e)}")
            return []