import logging
from typing import List, Dict
from .knowledge_base_manager import KnowledgeBaseManager

try:
    from .vector_rag_manager import VectorRAGManager
    VECTOR_AVAILABLE = True
except ImportError:
    VECTOR_AVAILABLE = False

logger = logging.getLogger(__name__)

class RAGHelper:
    def __init__(self, knowledge_base_path: str = "knowledge_base", use_vector_search: bool = True):
        self.kb_manager = KnowledgeBaseManager(knowledge_base_path)
        self.kb_manager.scan_documents()
        self.use_vector_search = use_vector_search and VECTOR_AVAILABLE
        self.vector_rag = None
        if self.use_vector_search:
            try:
                self.vector_rag = VectorRAGManager(knowledge_base_path, vector_db_path="/app/data/vector_db")
                logger.info("Vector RAG initialized")
            except Exception as e:
                logger.warning(f"Vector RAG failed: {e}")
                self.use_vector_search = False
        self.max_context_docs = 3
        self.similarity_threshold = 0.30

    def _route_query_to_categories(self, query: str) -> List[str]:
        ql = query.lower()
        cat_keywords = {
            'policies': ['return','refund','warranty','policy','shipping'],
            'faqs': ['how','what','why','help','troubleshoot','problem'],
            'product_manuals': ['spec','manual','guide','connect','cable','hdmi','usb']
        }
        scores = {}
        for cat, kw in cat_keywords.items():
            score = sum(1 for k in kw if k in ql)
            if score:
                scores[cat] = score
        if scores:
            return sorted(scores, key=scores.get, reverse=True)
        return ['policies','faqs','product_manuals']

    def get_relevant_context(self, query: str) -> str:
        categories = self._route_query_to_categories(query)
        if self.use_vector_search and self.vector_rag:
            try:
                results = self.vector_rag.retrieve_and_rerank_filtered(
                    query, target_categories=categories, initial_k=20, final_k=3,
                    similarity_threshold=self.similarity_threshold
                )
                return self._build_context_from_results(results)
            except Exception as e:
                logger.warning(f"Vector search failed: {e}, falling back to keyword")
        return self._get_keyword_context(query)

    def _get_keyword_context(self, query: str) -> str:
        results = self.kb_manager.search_documents(query)
        context = ""
        for r in results[:self.max_context_docs]:
            content = self.kb_manager.load_document_content(r['path'])
            if content:
                context += f"\n--- {r['title']} (Category: {r['category']}) ---\n{content[:1500]}\n"
        return context

    def _build_context_from_results(self, results: List[Dict]) -> str:
        if not results:
            return ""
        chunks = []
        for res in results:
            doc = res['document']
            meta = res['metadata']
            chunks.append(f"--- {meta['document_title']} (Category: {meta['category']}) ---\n{doc[:1200]}")
        return "\n".join(chunks[:3])

    def get_knowledge_base_stats(self) -> Dict:
        stats = self.kb_manager.get_stats() if hasattr(self.kb_manager, 'get_stats') else {'total_documents': 0}
        stats['vector_search_available'] = self.use_vector_search
        return stats

    def ensure_vector_index(self, force_reindex: bool = False):
        if self.vector_rag:
            return self.vector_rag.index_documents(force_reindex=force_reindex)
        return {"error": "Vector not available"}