import re
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

class ResumeRAGStore:
    def __init__(self, chunk_size: int = 250, overlap: int = 50):
        self.chunk_size = chunk_size
        self.overlap = overlap
        self.chunks = []
        self.vectorizer = None
        self.tfidf_matrix = None

    def index_resume_text(self, text: str):
        """
        Chunks raw resume text and indexes it using TF-IDF for RAG retrieval.
        """
        if not text or not text.strip():
            self.chunks = ["No resume content indexed."]
            return

        # Split into paragraphs and sentences
        lines = [line.strip() for line in text.splitlines() if line.strip()]
        self.chunks = []

        current_chunk = []
        current_len = 0

        for line in lines:
            current_chunk.append(line)
            current_len += len(line)
            if current_len >= self.chunk_size:
                self.chunks.append(" ".join(current_chunk))
                current_chunk = current_chunk[-1:]  # Keep last line for overlap
                current_len = len(current_chunk[0]) if current_chunk else 0

        if current_chunk:
            self.chunks.append(" ".join(current_chunk))

        if not self.chunks:
            self.chunks = [text]

        # Build TF-IDF index
        self.vectorizer = TfidfVectorizer(stop_words="english")
        try:
            self.tfidf_matrix = self.vectorizer.fit_transform(self.chunks)
        except Exception as e:
            print(f"[RAGStore] TFIDF indexing warning: {e}")
            self.tfidf_matrix = None

    def retrieve_context(self, query: str, top_k: int = 3) -> str:
        """
        Retrieves top_k most relevant resume passages matching the query.
        """
        if not self.chunks or self.vectorizer is None or self.tfidf_matrix is None:
            return "\n".join(self.chunks[:top_k])

        try:
            query_vec = self.vectorizer.transform([query])
            scores = cosine_similarity(query_vec, self.tfidf_matrix).flatten()
            top_indices = np.argsort(scores)[::-1][:top_k]
            
            relevant_chunks = [self.chunks[i] for i in top_indices if scores[i] > 0.05]
            if not relevant_chunks:
                relevant_chunks = self.chunks[:top_k]

            return "\n\n---\n\n".join(relevant_chunks)
        except Exception as e:
            print(f"[RAGStore] Context retrieval error: {e}")
            return "\n".join(self.chunks[:top_k])