# rag.py
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import google as genai
import os

class SimpleRAG:
    def __init__(self, api_key):
        # Initialize the embedding model and generative AI
        self.embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel("gemini-1.5-flash")
        self.index = None
        self.chunks = []
        self.is_initialized = False
        self.processing_status = None

    def chunk_text(self, text, chunk_size=700):
        """Split text into smaller chunks."""
        words = text.split()
        return [' '.join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]

    def process_search_data(self, search_data):
        """
        Process search result data and index it.
        'search_data' should be a list of job posting dictionaries.
        For each job posting, we combine key fields (e.g., job title and description) and then chunk the text.
        """
        try:
            self.processing_status = "Processing search data..."
            combined_text = ""
            for job in search_data:
                # Combine job title and job description (you can add more fields if needed)
                job_title = job.get('job_title', '')
                job_description = job.get('job_description', '')
                combined_text += f"Job Title: {job_title}. Description: {job_description}. "

            if not combined_text.strip():
                raise Exception("No text found in search results.")
            
            # Chunk the combined text
            self.chunks = self.chunk_text(combined_text)
            if not self.chunks:
                raise Exception("No content chunks were generated from search data.")
            
            # Generate embeddings and create the FAISS index
            embeddings = self.embedder.encode(self.chunks)
            vector_dimension = embeddings.shape[1]
            self.index = faiss.IndexFlatL2(vector_dimension)
            self.index.add(np.array(embeddings).astype('float32'))
            
            self.is_initialized = True
            self.processing_status = f"RAG system initialized with {len(self.chunks)} chunks."
            return {"status": "success", "message": self.processing_status}
        except Exception as e:
            self.processing_status = f"Error: {str(e)}"
            self.is_initialized = False
            return {"status": "error", "message": str(e)}

    def get_status(self):
        """Return the current processing status."""
        return {
            "is_initialized": self.is_initialized,
            "status": self.processing_status
        }

    def get_relevant_chunks(self, query, k=3):
        """Retrieve the top-k most relevant text chunks for a given query."""
        query_vector = self.embedder.encode([query])
        distances, chunk_indices = self.index.search(query_vector.astype('float32'), k)
        return [self.chunks[i] for i in chunk_indices[0]]

    def query(self, question):
        """Query the RAG system with a user question."""
        if not self.is_initialized:
            raise Exception("RAG system not initialized. Please process search data first.")
        try:
            context = self.get_relevant_chunks(question)
            prompt = f"""
            Based on the following context, provide a clear and concise answer.
            If the context doesn't contain enough relevant information, say "I don't have enough information to answer that question."

            Context:
            {' '.join(context)}

            Question: {question}
            """
            response = self.model.generate_content(prompt)
            return {
                "status": "success",
                "answer": response.text.strip(),
                "context": context
            }
        except Exception as e:
            return {
                "status": "error",
                "message": str(e)
            }