"""
RAG Builder - Build RAG knowledge base from processed documents
"""

import os
import json
from pathlib import Path
from typing import List, Dict, Optional, Union

try:
    from langchain_openai import ChatOpenAI
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_core.output_parsers import StrOutputParser
    from langchain_core.runnables import RunnablePassthrough
    from langchain_community.vectorstores import FAISS
    from langchain_text_splitters import RecursiveCharacterTextSplitter
    from langchain_core.documents import Document
    try:
        from langchain_huggingface import HuggingFaceEmbeddings
    except ImportError:
        from langchain_community.embeddings import HuggingFaceEmbeddings
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False


class RAGBuilder:
    """Build RAG knowledge base from documents"""
    
    def __init__(
        self,
        vector_store_path: Optional[Union[str, Path]] = None,
        embedding_model: str = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
    ):
        """
        Initialize RAG Builder
        
        Args:
            vector_store_path: Path to store vector database (default: ~/.mineru_rag/vector_db)
            embedding_model: Embedding model name
        """
        if not LANGCHAIN_AVAILABLE:
            raise ImportError(
                "LangChain is required for RAG functionality. "
                "Install it with: pip install mineru-rag[rag]"
            )
        
        # Set vector store path
        if vector_store_path is None:
            vector_store_path = Path.home() / ".mineru_rag" / "vector_db"
        self.vector_store_path = Path(vector_store_path)
        self.vector_store_path.mkdir(parents=True, exist_ok=True)
        
        # Initialize embeddings
        self.embedding_model = embedding_model
        self.embeddings = self._init_embeddings()
        
        # Initialize text splitter
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len,
            separators=["\n\n", "\n", "。", ".", " ", ""]
        )
        
        # Vector store
        self.vector_store: Optional[FAISS] = None
        self.doc_metadata: Dict[str, Dict] = {}
    
    def _init_embeddings(self):
        """Initialize embedding model"""
        try:
            # Set model cache directory
            model_cache_dir = self.vector_store_path / "models"
            model_cache_dir.mkdir(parents=True, exist_ok=True)
            os.environ['SENTENCE_TRANSFORMERS_HOME'] = str(model_cache_dir)
            os.environ['HF_HOME'] = str(model_cache_dir)
            
            # Detect device
            try:
                import torch
                device = 'cuda' if torch.cuda.is_available() else 'cpu'
            except ImportError:
                device = 'cpu'
            
            embeddings = HuggingFaceEmbeddings(
                model_name=self.embedding_model,
                model_kwargs={'device': device},
                encode_kwargs={'normalize_embeddings': True}
            )
            
            # Test embedding
            _ = embeddings.embed_query("test")
            return embeddings
        except Exception as e:
            raise RuntimeError(f"Failed to initialize embeddings: {str(e)}")
    
    def build_from_files(
        self,
        file_paths: List[Union[str, Path]],
        library_id: str = "default",
        metadata: Optional[Dict] = None
    ) -> bool:
        """
        Build vector database from markdown files
        
        Args:
            file_paths: List of markdown file paths
            library_id: Library identifier
            metadata: Optional metadata for documents
            
        Returns:
            True if successful
        """
        if not file_paths:
            raise ValueError("No files provided")
        
        all_documents = []
        local_metadata = {}
        
        for file_path in file_paths:
            file_path = Path(file_path)
            if not file_path.exists():
                print(f"⚠️  File not found: {file_path}")
                continue
            
            # Read markdown content
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
            except Exception as e:
                print(f"⚠️  Failed to read {file_path}: {str(e)}")
                continue
            
            # Split text
            chunks = self.text_splitter.split_text(content)
            
            # Create documents
            file_id = file_path.stem
            for i, chunk in enumerate(chunks):
                doc = Document(
                    page_content=chunk,
                    metadata={
                        'file_id': file_id,
                        'library_id': library_id,
                        'filename': file_path.name,
                        'chunk_index': i,
                        'total_chunks': len(chunks),
                        **(metadata or {})
                    }
                )
                all_documents.append(doc)
            
            local_metadata[file_id] = {
                'filename': file_path.name,
                'library_id': library_id,
                'chunk_count': len(chunks),
                **(metadata or {})
            }
        
        if not all_documents:
            raise ValueError("No valid documents to process")
        
        # Build vector store
        try:
            print(f"🔄 Building vector database from {len(all_documents)} chunks...")
            vector_store = FAISS.from_documents(all_documents, self.embeddings)
            
            # Save vector store
            store_path = self.vector_store_path / f"{library_id}_faiss"
            vector_store.save_local(str(store_path))
            print(f"✅ Vector database saved to: {store_path}")
            
            # Save metadata
            metadata_path = self.vector_store_path / f"{library_id}_metadata.json"
            with open(metadata_path, 'w', encoding='utf-8') as f:
                json.dump(local_metadata, f, ensure_ascii=False, indent=2)
            
            self.vector_store = vector_store
            self.doc_metadata = local_metadata
            
            return True
        except Exception as e:
            print(f"❌ Failed to build vector database: {str(e)}")
            raise
    
    def load_vector_store(self, library_id: str = "default") -> bool:
        """
        Load existing vector database
        
        Args:
            library_id: Library identifier
            
        Returns:
            True if successful
        """
        store_path = self.vector_store_path / f"{library_id}_faiss"
        
        if not store_path.exists():
            raise FileNotFoundError(f"Vector database not found: {store_path}")
        
        try:
            self.vector_store = FAISS.load_local(
                str(store_path),
                self.embeddings,
                allow_dangerous_deserialization=True
            )
            
            # Load metadata
            metadata_path = self.vector_store_path / f"{library_id}_metadata.json"
            if metadata_path.exists():
                with open(metadata_path, 'r', encoding='utf-8') as f:
                    self.doc_metadata = json.load(f)
            
            print(f"✅ Vector database loaded: {len(self.doc_metadata)} documents")
            return True
        except Exception as e:
            print(f"❌ Failed to load vector database: {str(e)}")
            raise
    
    def query(
        self,
        question: str,
        k: int = 4,
        file_id: Optional[str] = None
    ) -> Dict:
        """
        Query the RAG system
        
        Args:
            question: Question to ask
            k: Number of documents to retrieve
            file_id: Optional file ID to limit search
            
        Returns:
            Dictionary with answer and sources
        """
        if not self.vector_store:
            raise ValueError("Vector store not loaded. Call load_vector_store() first.")
        
        # Retrieve documents
        search_k = k * 10 if file_id else k
        retriever = self.vector_store.as_retriever(
            search_kwargs={"k": search_k}
        )
        
        docs = retriever.invoke(question)
        
        # Filter by file_id if specified
        if file_id:
            docs = [doc for doc in docs if doc.metadata.get('file_id') == file_id][:k]
        else:
            docs = docs[:k]
        
        # Format context
        context_parts = []
        for doc in docs:
            filename = doc.metadata.get('filename', 'Unknown')
            chunk_idx = doc.metadata.get('chunk_index', 0)
            context_parts.append(f"[{filename}, chunk {chunk_idx+1}]\n{doc.page_content}")
        
        context = "\n\n---\n\n".join(context_parts)
        
        # Prepare sources
        sources = []
        for doc in docs:
            sources.append({
                'filename': doc.metadata.get('filename', 'Unknown'),
                'file_id': doc.metadata.get('file_id', ''),
                'chunk_index': doc.metadata.get('chunk_index', 0),
                'content_preview': doc.page_content[:200] + "..."
            })
        
        return {
            'question': question,
            'context': context,
            'sources': sources,
            'num_sources': len(sources)
        }

