# -*- coding: utf-8 -*-
from typing import List
import numpy as np
from config import APP_CONFIG

try:
    from sklearn.feature_extraction.text import TfidfVectorizer
    from sklearn.metrics.pairwise import cosine_similarity
    SKLEARN_AVAILABLE = True
except ImportError:
    SKLEARN_AVAILABLE = False

class ResearchRAGEngine:
    def __init__(self, raw_texts: List[str]):
        self.chunks = []
        self.vectorizer = None
        self.tfidf_matrix = None
        self.chunk_size = APP_CONFIG["rag_chunk_size"]
        self._initialize_chunks(raw_texts)

    def _initialize_chunks(self, raw_texts: List[str]):
        for text in raw_texts:
            for i in range(0, len(text), self.chunk_size):
                chunk = text[i:i + self.chunk_size].strip()
                if len(chunk) > 50:
                    self.chunks.append(chunk)
                    
        if SKLEARN_AVAILABLE and self.chunks:
            self.vectorizer = TfidfVectorizer(stop_words='english', max_df=0.85)
            try:
                self.tfidf_matrix = self.vectorizer.fit_transform(self.chunks)
            except ValueError:
                self.vectorizer = None

    def get_relevant_context(self, query: str, top_k: int = None) -> str:
        if not self.chunks: 
            return ""
        if top_k is None: 
            top_k = APP_CONFIG["rag_top_k"]
        
        relevant_chunks = []
        if self.vectorizer and self.tfidf_matrix is not None:
            query_vec = self.vectorizer.transform([query])
            similarities = cosine_similarity(query_vec, self.tfidf_matrix).flatten()
            top_indices = similarities.argsort()[-top_k:][::-1]
            relevant_chunks = [self.chunks[i] for i in top_indices if similarities[i] > 0.05]
        else:
            scores = []
            query_words = set(query.lower().split())
            for chunk in self.chunks:
                chunk_lower = chunk.lower()
                score = sum(1 for word in query_words if word in chunk_lower)
                scores.append(score)
            top_indices = np.argsort(scores)[-top_k:]
            relevant_chunks = [self.chunks[i] for i in top_indices if scores[i] > 0]
            
        if not relevant_chunks: 
            return ""
        return "\n--- Target Context Snippet ---\n" + "\n\n".join(relevant_chunks)

def initialize_rag(texts: List[str]):
    return ResearchRAGEngine(texts)