#!/usr/bin/env python3
"""
RAG Service Module for Themblem
基于 LangChain 的实时知识库检索服务，支持缓存优化
"""

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from .models import Article
from django.core.cache import cache
from typing import List, Dict, Any
import hashlib
import pickle
import time
from bs4 import BeautifulSoup

class CachedLangChainRAG:
    """基于缓存的实时LangChain RAG服务"""
    
    def __init__(self):
        # Use pre-downloaded model from container (relative to this file)
        import os
        from pathlib import Path
        
        # Model path relative to this file: api/products/rag_service.py -> ../../models/
        model_path = Path(__file__).parent.parent.parent / "models" / "text2vec-base-chinese"
        
        # Fall back to downloading if local model doesn't exist (for development)
        if not model_path.exists():
            model_path = "shibing624/text2vec-base-chinese"
        else:
            model_path = str(model_path)
        
        self.embeddings = HuggingFaceEmbeddings(
            model_name=model_path,
            model_kwargs={'device': 'cpu'},
            encode_kwargs={'normalize_embeddings': True}
        )
        
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            separators=["\n\n", "\n", "。", "！", "？", ".", "!", "?", "；", ";"]
        )
        
        # 缓存配置
        self.vector_cache_timeout = 86400 * 30 # 30天
        self.search_cache_timeout = 86400 # 1天
    
    def _get_article_content_hash(self, article: Article) -> str:
        """获取文章内容的hash值"""
        content = f"{article.title or ''}{article.body}{article.options or ''}"
        return hashlib.md5(content.encode('utf-8')).hexdigest()
    
    def _get_articles_signature(self, articles: List[Article]) -> str:
        """获取文章集合的签名"""
        signatures = []
        for article in articles:
            content_hash = self._get_article_content_hash(article)
            signatures.append(f"{article.id}_{content_hash}")
        
        # 排序确保一致性
        signatures.sort()
        combined = '|'.join(signatures)
        return hashlib.md5(combined.encode('utf-8')).hexdigest()
    
    def search(self, query: str, max_results: int = 3, tenant_id: int = None) -> List[Dict[str, Any]]:
        """带缓存的搜索"""
        start_time = time.time()
        print(f"\n🔍 RAG搜索开始: {query}")
        print(f"📊 参数: max_results={max_results}, tenant_id={tenant_id}")
        
        # 生成查询缓存key
        query_hash = hashlib.md5(query.encode('utf-8')).hexdigest()
        cache_key = f"rag_search_{query_hash}_{max_results}_{tenant_id}"
        
        # 尝试从缓存获取搜索结果
        cache_check_start = time.time()
        cached_result = cache.get(cache_key)
        cache_check_time = time.time() - cache_check_start
        
        if cached_result:
            total_time = time.time() - start_time
            print(f"✅ 缓存命中: {cache_key}")
            print(f"📚 返回缓存结果: {len(cached_result)} 条")
            print(f"⏱️ 缓存检查耗时: {cache_check_time:.3f}秒")
            print(f"⏱️ 总耗时: {total_time:.3f}秒")
            return cached_result
        
        print(f"❌ 缓存未命中: {cache_key}")
        print(f"⏱️ 缓存检查耗时: {cache_check_time:.3f}秒")
        print("🔄 执行实时搜索...")
        
        # 缓存未命中，执行搜索
        search_start_time = time.time()
        results = self._perform_search(query, max_results, tenant_id)
        search_time = time.time() - search_start_time
        
        # 缓存搜索结果
        cache_start_time = time.time()
        cache.set(cache_key, results, self.search_cache_timeout)
        cache_time = time.time() - cache_start_time
        
        total_time = time.time() - start_time
        print(f"💾 搜索结果已缓存，过期时间: {self.search_cache_timeout}秒")
        print(f"⏱️ 搜索耗时: {search_time:.3f}秒")
        print(f"⏱️ 缓存写入耗时: {cache_time:.3f}秒")
        print(f"⏱️ 总耗时: {total_time:.3f}秒")
        
        return results
    
    def _perform_search(self, query: str, max_results: int, tenant_id: int = None) -> List[Dict[str, Any]]:
        """执行实际搜索"""
        search_start_time = time.time()
        try:
            print(f"🔧 获取向量存储...")
            vector_start_time = time.time()
            # 获取或创建缓存的向量存储
            vectorstore = self._get_cached_vectorstore(tenant_id)
            vector_time = time.time() - vector_start_time
            print(f"⏱️ 向量存储获取耗时: {vector_time:.3f}秒")
            
            if not vectorstore:
                print("❌ 无法获取向量存储")
                return []
            
            print(f"🎯 执行相似度搜索，查询: '{query}'")
            similarity_start_time = time.time()
            # 执行相似度搜索
            docs = vectorstore.similarity_search_with_score(query, k=max_results)
            similarity_time = time.time() - similarity_start_time
            print(f"⏱️ 相似度搜索耗时: {similarity_time:.3f}秒")
            print(f"📊 找到 {len(docs)} 个候选文档")
            
            # 格式化结果
            format_start_time = time.time()
            results = []
            for i, (doc, score) in enumerate(docs, 1):
                result = {
                    'id': doc.metadata.get('article_id'),
                    'title': doc.metadata.get('title'),
                    'content': doc.page_content,
                    'score': float(score),
                    'tenant_id': doc.metadata.get('tenant_id'),
                    'url': doc.metadata.get('url')
                }
                results.append(result)
                
                print(f"  📄 结果 {i}:")
                print(f"     ID: {result['id']}")
                print(f"     标题: {result['title']}")
                print(f"     相关度: {result['score']:.4f}")
                print(f"     内容预览: {result['content'][:100]}...")
            
            format_time = time.time() - format_start_time
            total_search_time = time.time() - search_start_time
            print(f"⏱️ 结果格式化耗时: {format_time:.3f}秒")
            print(f"✅ 搜索完成，返回 {len(results)} 个结果")
            print(f"⏱️ 搜索总耗时: {total_search_time:.3f}秒")
            return results
            
        except Exception as e:
            total_search_time = time.time() - search_start_time
            print(f"❌ RAG搜索失败: {str(e)}")
            print(f"⏱️ 搜索耗时: {total_search_time:.3f}秒")
            import traceback
            traceback.print_exc()
            return []
    
    def _get_cached_vectorstore(self, tenant_id: int = None):
        """获取或创建缓存的向量存储"""
        vector_start_time = time.time()
        print(f"🔍 获取知识库文章...")
        
        articles_start_time = time.time()
        # 获取知识库文章列表
        articles = self._get_knowledge_base_articles(tenant_id)
        articles_time = time.time() - articles_start_time
        print(f"⏱️ 文章获取耗时: {articles_time:.3f}秒")
        
        if not articles:
            print("❌ 没有找到知识库文章")
            return None
        
        print(f"📚 找到 {len(articles)} 篇知识库文章")
        
        # 检查向量缓存
        cache_check_start = time.time()
        cache_key = self._get_vector_cache_key(articles, tenant_id)
        print(f"🔑 向量缓存key: {cache_key}")
        
        cached_vectors = cache.get(cache_key)
        cache_check_time = time.time() - cache_check_start
        print(f"⏱️ 向量缓存检查耗时: {cache_check_time:.3f}秒")
        
        if cached_vectors:
            restore_start_time = time.time()
            print("✅ 向量缓存命中，恢复向量存储")
            # 从缓存恢复向量存储
            result = self._restore_vectorstore(cached_vectors)
            restore_time = time.time() - restore_start_time
            print(f"⏱️ 向量存储恢复耗时: {restore_time:.3f}秒")
            
            total_time = time.time() - vector_start_time
            print(f"⏱️ 向量存储获取总耗时: {total_time:.3f}秒")
            return result
        
        print("❌ 向量缓存未命中，创建新的向量存储")
        create_start_time = time.time()
        # 缓存未命中，创建新的向量存储
        vectorstore = self._create_vectorstore(articles)
        create_time = time.time() - create_start_time
        print(f"⏱️ 向量存储创建耗时: {create_time:.3f}秒")
        
        # 缓存向量存储
        cache_start_time = time.time()
        self._cache_vectorstore(vectorstore, cache_key)
        cache_time = time.time() - cache_start_time
        print(f"⏱️ 向量存储缓存耗时: {cache_time:.3f}秒")
        
        total_time = time.time() - vector_start_time
        print(f"⏱️ 向量存储获取总耗时: {total_time:.3f}秒")
        return vectorstore
    
    def _get_knowledge_base_articles(self, tenant_id: int = None) -> List[Article]:
        """获取知识库文章（带缓存）"""
        articles_start_time = time.time()
        
        # 生成基于文章内容的缓存key
        cache_key = self._get_articles_cache_key(tenant_id)
        print(f"🔍 检查文章缓存: {cache_key}")
        
        cache_check_start = time.time()
        cached_articles = cache.get(cache_key)
        cache_check_time = time.time() - cache_check_start
        
        if cached_articles:
            total_time = time.time() - articles_start_time
            print(f"✅ 文章缓存命中，返回 {len(cached_articles)} 篇文章")
            print(f"⏱️ 缓存检查耗时: {cache_check_time:.3f}秒")
            print(f"⏱️ 文章获取总耗时: {total_time:.3f}秒")
            return cached_articles
        
        print("❌ 文章缓存未命中，查询数据库...")
        print(f"⏱️ 缓存检查耗时: {cache_check_time:.3f}秒")
        
        # 查询数据库
        db_start_time = time.time()
        filter_kwargs = {'is_platform_knowledge_base': True}
        if tenant_id:
            filter_kwargs['tenant_id'] = tenant_id
        
        articles = list(Article.objects.filter(**filter_kwargs))
        db_time = time.time() - db_start_time
        print(f"📊 数据库查询结果: {len(articles)} 篇文章")
        print(f"⏱️ 数据库查询耗时: {db_time:.3f}秒")
        
        # 缓存文章列表
        cache_start_time = time.time()
        cache.set(cache_key, articles, self.vector_cache_timeout)
        cache_time = time.time() - cache_start_time
        print(f"💾 文章列表已缓存，过期时间: {self.vector_cache_timeout}秒")
        print(f"⏱️ 文章缓存写入耗时: {cache_time:.3f}秒")
        
        total_time = time.time() - articles_start_time
        print(f"⏱️ 文章获取总耗时: {total_time:.3f}秒")
        return articles
    
    def _get_articles_cache_key(self, tenant_id: int = None) -> str:
        """生成基于文章内容的缓存key"""
        # 查询当前文章状态
        filter_kwargs = {'is_platform_knowledge_base': True}
        if tenant_id:
            filter_kwargs['tenant_id'] = tenant_id
        
        articles = Article.objects.filter(**filter_kwargs)
        
        # 生成文章内容哈希
        content_signature = self._get_articles_content_signature(articles)
        
        # 缓存key包含租户ID和内容签名
        cache_key = f"rag_kb_articles_{tenant_id}_{content_signature}"
        print(f"🔑 生成文章缓存key: {cache_key}")
        print(f"📊 内容签名: {content_signature}")
        
        return cache_key
    
    def _get_articles_content_signature(self, articles) -> str:
        """生成文章内容的签名"""
        signatures = []
        
        for article in articles:
            # 基于文章ID、标题、完整内容、选项等生成签名
            title = article.title or ''
            body = article.body or ''
            options = article.options or ''
            
            # 组合签名要素 - 直接哈希完整内容
            article_signature = f"{article.id}_{title}_{body}_{options}"
            signatures.append(hashlib.md5(article_signature.encode('utf-8')).hexdigest())
        
        # 排序确保一致性
        signatures.sort()
        combined = '|'.join(signatures)
        
        # 生成最终签名
        final_signature = hashlib.md5(combined.encode('utf-8')).hexdigest()
        return final_signature
    
    def _clean_html_content(self, html_content: str) -> str:
        """清理HTML内容，提取纯文本"""
        if not html_content:
            return ""
        
        # 使用BeautifulSoup解析HTML
        soup = BeautifulSoup(html_content, 'html.parser')
        
        # 移除script、style、head、title、meta、link等标签
        for tag in soup(["script", "style", "head", "title", "meta", "link", "noscript"]):
            tag.decompose()
        
        # 获取纯文本
        text = soup.get_text()
        
        # 清理空白字符：移除多余的空行和空格
        lines = (line.strip() for line in text.splitlines())
        chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
        text = ' '.join(chunk for chunk in chunks if chunk)
        
        return text
    
    def _get_vector_cache_key(self, articles: List[Article], tenant_id: int = None) -> str:
        """生成向量缓存key"""
        # 基于文章集合签名生成缓存key
        signature = self._get_articles_signature(articles)
        cache_key = f"rag_vectors_{tenant_id}_{signature}"
        print(f"🔑 生成向量缓存key: {cache_key}")
        return cache_key
    
    def _create_vectorstore(self, articles: List[Article]):
        """创建向量存储"""
        create_start_time = time.time()
        print(f"🔧 创建文档对象...")
        
        docs_start_time = time.time()
        documents = self._create_documents(articles)
        docs_time = time.time() - docs_start_time
        print(f"⏱️ 文档创建耗时: {docs_time:.3f}秒")
        
        if not documents:
            print("❌ 无法创建文档对象")
            return None
        
        print(f"📄 创建了 {len(documents)} 个文档片段")
        print(f"🎯 开始向量化...")
        
        embedding_start_time = time.time()
        vectorstore = FAISS.from_documents(
            documents=documents,
            embedding=self.embeddings
        )
        embedding_time = time.time() - embedding_start_time
        print(f"⏱️ 向量化耗时: {embedding_time:.3f}秒")
        
        total_time = time.time() - create_start_time
        print(f"✅ 向量存储创建完成，包含 {len(documents)} 个向量")
        print(f"⏱️ 向量存储创建总耗时: {total_time:.3f}秒")
        return vectorstore
    
    def _create_documents(self, articles: List[Article]) -> List[Document]:
        """创建文档对象"""
        docs_start_time = time.time()
        documents = []
        
        for article in articles:
            try:
                article_start_time = time.time()
                print(f"  📝 处理文章: {article.title or '无标题'} (ID: {article.id})")
                
                # 清理HTML内容
                clean_start_time = time.time()
                clean_body = self._clean_html_content(article.body)
                clean_time = time.time() - clean_start_time
                print(f"     🧹 HTML清理耗时: {clean_time:.3f}秒")
                
                split_start_time = time.time()
                chunks = self.text_splitter.split_text(clean_body)
                split_time = time.time() - split_start_time
                print(f"     ✂️ 分割为 {len(chunks)} 个片段，耗时: {split_time:.3f}秒")
                
                valid_chunks = 0
                for i, chunk in enumerate(chunks):
                    if len(chunk.strip()) < 50:
                        continue
                        
                    doc = Document(
                        page_content=chunk,
                        metadata={
                            'article_id': article.id,
                            'title': article.title or '无标题',
                            'tenant_id': article.tenant_id,
                            'chunk_index': i,
                            'url': article.url,
                            'source': 'platform_knowledge_base',
                            'html_cleaned': True
                        }
                    )
                    documents.append(doc)
                    valid_chunks += 1
                
                article_time = time.time() - article_start_time
                print(f"     ✅ 有效片段: {valid_chunks}/{len(chunks)}，文章处理耗时: {article_time:.3f}秒")
                    
            except Exception as e:
                print(f"❌ 处理文章 {article.id} 失败: {str(e)}")
                continue
        
        total_time = time.time() - docs_start_time
        print(f"📊 总共创建了 {len(documents)} 个有效文档片段")
        print(f"⏱️ 文档创建总耗时: {total_time:.3f}秒")
        return documents
    
    def _cache_vectorstore(self, vectorstore, cache_key: str):
        """缓存向量存储"""
        try:
            print(f"💾 序列化向量存储...")
            serialize_start_time = time.time()
            # 序列化向量存储
            serialized_data = pickle.dumps(vectorstore)
            serialize_time = time.time() - serialize_start_time
            print(f"📦 序列化完成，大小: {len(serialized_data)} 字节，耗时: {serialize_time:.3f}秒")
            
            cache_start_time = time.time()
            cache.set(cache_key, serialized_data, self.vector_cache_timeout)
            cache_time = time.time() - cache_start_time
            print(f"✅ 向量存储已缓存，过期时间: {self.vector_cache_timeout}秒，缓存写入耗时: {cache_time:.3f}秒")
        except Exception as e:
            print(f"❌ 缓存向量存储失败: {str(e)}")
            import traceback
            traceback.print_exc()
    
    def _restore_vectorstore(self, serialized_data: bytes):
        """从缓存恢复向量存储"""
        try:
            restore_start_time = time.time()
            result = pickle.loads(serialized_data)
            restore_time = time.time() - restore_start_time
            print(f"⏱️ 向量存储反序列化耗时: {restore_time:.3f}秒")
            return result
        except Exception as e:
            print(f"❌ 恢复向量存储失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
    
    def get_enhanced_context(self, query: str, max_results: int = 3, tenant_id: int = None) -> str:
        """获取增强的上下文信息"""
        context_start_time = time.time()
        print(f"\n📖 开始构建增强上下文...")
        
        results = self.search(query, max_results, tenant_id)
        
        if not results:
            total_time = time.time() - context_start_time
            print(f"⏱️ 上下文构建耗时: {total_time:.3f}秒")
            return "未找到相关知识库内容。"
        
        # 构建上下文
        build_start_time = time.time()
        context = "基于知识库检索到的相关信息：\n\n"
        for i, result in enumerate(results, 1):
            context += f"【{i}】{result['title']}\n"
            context += f"相关度：{result['score']:.3f}\n"
            context += f"内容：{result['content']}\n\n"
        
        build_time = time.time() - build_start_time
        total_time = time.time() - context_start_time
        print(f"⏱️ 上下文构建耗时: {build_time:.3f}秒")
        print(f"⏱️ 上下文构建总耗时: {total_time:.3f}秒")
        
        return context
    
    def clear_old_cache(self, tenant_id: int = None):
        """清理旧的缓存，避免缓存碎片"""
        try:
            # 清理旧的搜索缓存
            old_search_pattern = f"rag_search_*_{tenant_id}_*"
            self._clear_cache_by_pattern(old_search_pattern)
            
            # 清理旧的向量缓存
            old_vector_pattern = f"rag_vectors_{tenant_id}_*"
            self._clear_cache_by_pattern(old_vector_pattern)
            
            print(f"🧹 已清理租户 {tenant_id} 的旧缓存")
        except Exception as e:
            print(f"⚠️ 清理旧缓存失败: {e}")
    
    def _clear_cache_by_pattern(self, pattern: str):
        """根据模式清理缓存（Django缓存限制，这里只是标记）"""
        # 注意：Django的默认缓存后端不支持模式删除
        # 这里标记需要清理的缓存，实际清理在下次访问时处理
        cache.set(f"rag_cache_invalidated_{pattern}", time.time(), 300)  # 5分钟标记
    
    def _is_cache_invalidated(self, cache_key: str) -> bool:
        """检查缓存是否被标记为失效"""
        # 检查是否有相关的失效标记
        for pattern in ["rag_search_", "rag_vectors_", "rag_kb_articles_"]:
            if pattern in cache_key:
                invalidated_key = f"rag_cache_invalidated_{pattern}*"
                if cache.get(invalidated_key):
                    return True
        return False
