import asyncio
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from typing import List, Dict, Any, Optional
from datetime import datetime
import uuid
import logging
from config import Config

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class RAGDocumentStore:
    """RAG-enabled document store with vector similarity search"""
    
    def __init__(self):
        self.config = Config()
        self._initialize_vector_db()
        self.embeddings = None
        self.sentence_model = None
        self.embed_function = None
        self._initialize_text_splitter()
        self.documents_metadata = {}
        logger.info("RAG Document Store initialized (embeddings deferred)")
    
    def _ensure_embeddings_initialized(self):
        """Lazy initialization of embeddings"""
        if self.embed_function is not None:
            return
        
        try:
            if self.config.USE_OPENAI_EMBEDDINGS and self.config.OPENAI_API_KEY:
                self.embeddings = OpenAIEmbeddings(api_key=self.config.OPENAI_API_KEY,base_url=self.config.BASE_URL,model=self.config.EMBEDDING_MODEL)
                self.embed_function = self._openai_embed
                logger.info("Using OpenAI embeddings")
            else:
                self.sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
                self.embed_function = self._sentence_transformer_embed
                logger.info(f"Using SentenceTransformer: {self.config.EMBEDDING_MODEL}")
        except Exception as e:
            logger.error(f"Failed to initialize embeddings: {e}")
            raise
    
    def _initialize_vector_db(self):
        """Initialize ChromaDB"""
        try:
            self.chroma_client = chromadb.PersistentClient(
                path=self.config.VECTOR_DB_PATH,
                settings=Settings(anonymized_telemetry=False)
            )
            
            self.collection = self.chroma_client.get_or_create_collection(
                name="documents",
                metadata={"hnsw:space": "cosine"}
            )
            logger.info(f"Vector database initialized at {self.config.VECTOR_DB_PATH}")
        except Exception as e:
            logger.error(f"Failed to initialize vector database: {e}")
            raise
    
    def _initialize_text_splitter(self):
        """Initialize text splitter"""
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.config.CHUNK_SIZE,
            chunk_overlap=self.config.CHUNK_OVERLAP,
            length_function=len
        )
    
    def _sentence_transformer_embed(self, texts: List[str]) -> List[List[float]]:
        """Create embeddings using SentenceTransformer"""
        try:
            embeddings = self.sentence_model.encode(texts, show_progress_bar=True)
            return embeddings.tolist()
        except Exception as e:
            logger.error(f"SentenceTransformer embedding failed: {e}")
            raise
    
    async def _openai_embed(self, texts: List[str]) -> List[List[float]]:
        """Create embeddings using OpenAI with T-Systems batch size limits"""
        try:
            # T-Systems has a batch size limit of 128
            max_batch_size = 100  # Use 100 to be safe
            # T-Systems has a batch size limit of 128
            max_batch_size = 100  # Use 100 to be safe
            all_embeddings = []
            
            logger.info(f"Creating embeddings for {len(texts)} texts in batches of {max_batch_size}")
            
            # Process in batches
            for i in range(0, len(texts), max_batch_size):
                batch = texts[i:i + max_batch_size]
                batch_num = i//max_batch_size + 1
                total_batches = (len(texts) + max_batch_size - 1)//max_batch_size
                
                logger.info(f"Processing embedding batch {batch_num}/{total_batches} ({len(batch)} texts)")
                
                try:
                    # Create embeddings for this batch
                    batch_embeddings = await self.embeddings.aembed_documents(batch)
                    all_embeddings.extend(batch_embeddings)
                    
                    # Small delay between batches to avoid rate limiting
                    if i + max_batch_size < len(texts):
                        await asyncio.sleep(0.2)  # 200ms delay
                        
                except Exception as batch_error:
                    logger.error(f"Batch {batch_num} failed: {batch_error}")
                    # Try with even smaller batch size
                    if len(batch) > 50:
                        logger.info("Retrying with smaller batch size...")
                        for j in range(0, len(batch), 50):
                            mini_batch = batch[j:j + 50]
                            mini_embeddings = await self.embeddings.aembed_documents(mini_batch)
                            all_embeddings.extend(mini_embeddings)
                            await asyncio.sleep(0.1)
                    else:
                        raise batch_error
            
            logger.info(f"Successfully created {len(all_embeddings)} embeddings")
            return all_embeddings
            
        except Exception as e:
            logger.error(f"OpenAI embedding failed: {e}")
            raise
    
    # async def add_document(self, doc_id: str, filename: str, file_type: str, 
    #                       content: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
    #     """Add document with vector embeddings"""
    #     try:
    #         logger.info(f"Adding document: {filename} ({len(content)} characters)")
            
    #         # Split content into chunks
    #         chunks = self.text_splitter.split_text(content)
    #         logger.info(f"Created {len(chunks)} chunks for {filename}")
            
    #         if not chunks:
    #             return {"error": "No content chunks created"}
            
    #         # Create embeddings
    #         embeddings = await self.embed_function(chunks) if asyncio.iscoroutinefunction(self.embed_function) else self.embed_function(chunks)
            
    #         # Prepare data for ChromaDB
    #         chunk_ids = [f"{doc_id}_chunk_{i}" for i in range(len(chunks))]
            
    #         chunk_metadata = []
    #         for i, chunk in enumerate(chunks):
    #             chunk_metadata.append({
    #                 "doc_id": doc_id,
    #                 "filename": filename,
    #                 "file_type": file_type,
    #                 "chunk_index": i,
    #                 "chunk_size": len(chunk),
    #                 "upload_time": datetime.now().isoformat(),
    #                 **metadata
    #             })
            
    #         # Add to vector database
    #         self.collection.add(
    #             embeddings=embeddings,
    #             documents=chunks,
    #             metadatas=chunk_metadata,
    #             ids=chunk_ids
    #         )
            
    #         # Store document metadata
    #         self.documents_metadata[doc_id] = {
    #             "filename": filename,
    #             "file_type": file_type,
    #             "chunk_count": len(chunks),
    #             "total_chars": len(content),
    #             "metadata": metadata,
    #             "upload_time": datetime.now().isoformat()
    #         }
            
    #         logger.info(f"Successfully added {filename} with {len(chunks)} chunks")
    #         return {
    #             "doc_id": doc_id,
    #             "chunks_created": len(chunks),
    #             "status": "success"
    #         }
            
    #     except Exception as e:
    #         logger.error(f"Failed to add document {filename}: {e}")
    #         return {"error": str(e)}


    async def add_document(self, doc_id: str, filename: str, file_type: str, 
                        content: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
        """Add document with vector embeddings and batched ChromaDB insertion"""
        try:
            logger.info(f"Adding document: {filename} ({len(content)} characters)")
            
            # Split content into chunks
            chunks = self.text_splitter.split_text(content)
            logger.info(f"Created {len(chunks)} chunks for {filename}")
            
            if not chunks:
                return {"error": "No content chunks created"}
            
            # Create embeddings (already batched)
            self._ensure_embeddings_initialized()
            embeddings = await self.embed_function(chunks) if asyncio.iscoroutinefunction(self.embed_function) else self.embed_function(chunks)
            
            # Prepare data for ChromaDB
            chunk_ids = [f"{doc_id}_chunk_{i}" for i in range(len(chunks))]
            
            chunk_metadata = []
            for i, chunk in enumerate(chunks):
                chunk_metadata.append({
                    "doc_id": doc_id,
                    "filename": filename,
                    "file_type": file_type,
                    "chunk_index": i,
                    "chunk_size": len(chunk),
                    "upload_time": datetime.now().isoformat(),
                    **metadata
                })
            
            # Add to ChromaDB in batches
            await self._add_to_chromadb_batched(embeddings, chunks, chunk_metadata, chunk_ids)
            
            # Store document metadata
            self.documents_metadata[doc_id] = {
                "filename": filename,
                "file_type": file_type,
                "chunk_count": len(chunks),
                "total_chars": len(content),
                "metadata": metadata,
                "upload_time": datetime.now().isoformat()
            }
            
            logger.info(f"Successfully added {filename} with {len(chunks)} chunks")
            return {
                "doc_id": doc_id,
                "chunks_created": len(chunks),
                "status": "success"
            }
            
        except Exception as e:
            logger.error(f"Failed to add document {filename}: {e}")
            return {"error": str(e)}

    async def _add_to_chromadb_batched(self, embeddings: List[List[float]], chunks: List[str], 
                                    chunk_metadata: List[Dict], chunk_ids: List[str]):
        """Add data to ChromaDB in batches to avoid batch size limits"""
        try:
            # ChromaDB batch size limit (use conservative value)
            chromadb_batch_size = 5000
            total_chunks = len(chunks)
            
            logger.info(f"Adding {total_chunks} chunks to ChromaDB in batches of {chromadb_batch_size}")
            
            for i in range(0, total_chunks, chromadb_batch_size):
                end_idx = min(i + chromadb_batch_size, total_chunks)
                
                # Get batch data
                batch_embeddings = embeddings[i:end_idx]
                batch_chunks = chunks[i:end_idx]
                batch_metadata = chunk_metadata[i:end_idx]
                batch_ids = chunk_ids[i:end_idx]
                
                # Add batch to ChromaDB
                self.collection.add(
                    embeddings=batch_embeddings,
                    documents=batch_chunks,
                    metadatas=batch_metadata,
                    ids=batch_ids
                )
                
                batch_num = i // chromadb_batch_size + 1
                total_batches = (total_chunks + chromadb_batch_size - 1) // chromadb_batch_size
                
                logger.info(f"Added ChromaDB batch {batch_num}/{total_batches} ({len(batch_chunks)} chunks)")
                
                # Small delay between batches to be gentle on ChromaDB
                if end_idx < total_chunks:
                    await asyncio.sleep(0.1)
            
            logger.info(f"Successfully added all {total_chunks} chunks to ChromaDB")
            
        except Exception as e:
            logger.error(f"Failed to add to ChromaDB: {e}")
            raise

    async def semantic_search(self, query: str, n_results: int = 10, 
                             doc_ids: Optional[List[str]] = None) -> List[Dict[str, Any]]:
        """Perform semantic similarity search"""
        try:
            logger.info(f"Performing semantic search: '{query}' (n_results={n_results})")
            
            self._ensure_embeddings_initialized()
            query_embedding = await self.embed_function([query]) if asyncio.iscoroutinefunction(self.embed_function) else self.embed_function([query])
            query_embedding = await self.embed_function([query]) if asyncio.iscoroutinefunction(self.embed_function) else self.embed_function([query])
            
            # Prepare filters
            where_filter = {}
            if doc_ids:
                where_filter["doc_id"] = {"$in": doc_ids}
            
            # Search in ChromaDB
            results = self.collection.query(
                query_embeddings=query_embedding,
                n_results=min(n_results, self.collection.count()),
                where=where_filter if where_filter else None,
                include=["documents", "metadatas", "distances"]
            )
            
            # Format results
            formatted_results = []
            if results["documents"] and results["documents"][0]:
                for i in range(len(results["documents"][0])):
                    similarity_score = 1 - results["distances"][0][i] if results["distances"] else 0
                    formatted_results.append({
                        "content": results["documents"][0][i],
                        "metadata": results["metadatas"][0][i],
                        "similarity_score": similarity_score,
                        "chunk_id": results["ids"][0][i] if "ids" in results else f"chunk_{i}"
                    })
            
            logger.info(f"Found {len(formatted_results)} semantic matches")
            return formatted_results
            
        except Exception as e:
            logger.error(f"Semantic search failed: {e}")
            return []
    
    async def reconstruct_document_metadata(self):
        """Reconstruct document metadata from ChromaDB"""
        try:
            # Get all documents from ChromaDB
            results = self.collection.get(include=["metadatas"])
            
            if not results["metadatas"]:
                logger.info("No chunks found in ChromaDB")
                return
            
            # Reconstruct metadata by grouping chunks by document
            reconstructed_docs = {}
            
            for metadata in results["metadatas"]:
                doc_id = metadata.get("doc_id")
                if doc_id and doc_id not in reconstructed_docs:
                    reconstructed_docs[doc_id] = {
                        "filename": metadata.get("filename", "unknown"),
                        "file_type": metadata.get("file_type", "unknown"),
                        "chunk_count": 0,
                        "total_chars": 0,
                        "metadata": {
                            "file_size_bytes": metadata.get("file_size_bytes", 0),
                            "file_size_mb": metadata.get("file_size_mb", 0),
                            "file_extension": metadata.get("file_extension", ""),
                            "processing_strategy": metadata.get("processing_strategy", "unknown")
                        },
                        "upload_time": metadata.get("upload_time", "unknown")
                    }
                
                if doc_id:
                    reconstructed_docs[doc_id]["chunk_count"] += 1
                    reconstructed_docs[doc_id]["total_chars"] += metadata.get("chunk_size", 0)
            
            # Update documents_metadata
            self.documents_metadata = reconstructed_docs
            
            logger.info(f"Reconstructed metadata for {len(reconstructed_docs)} documents")
            return reconstructed_docs
            
        except Exception as e:
            logger.error(f"Failed to reconstruct document metadata: {e}")
            return {}

    async def get_document_summary(self) -> Dict[str, Any]:
        """Get summary of all documents with auto-reconstruction"""
        try:
            total_chunks = self.collection.count()
            
            # If we have chunks but no metadata, try to reconstruct
            if total_chunks > 0 and len(self.documents_metadata) == 0:
                logger.info("Found chunks but no metadata, attempting reconstruction...")
                await self.reconstruct_document_metadata()
            
            return {
                "total_documents": len(self.documents_metadata),
                "total_chunks": total_chunks,
                "documents": self.documents_metadata,
                "vector_db_status": "active",
                "embedding_model": self.config.EMBEDDING_MODEL if not self.config.USE_OPENAI_EMBEDDINGS else "OpenAI"
            }
        except Exception as e:
            logger.error(f"Failed to get document summary: {e}")
            return {"error": str(e)}
    
    def clear_all_documents(self):
        """Clear all documents from the store"""
        try:
            # Delete the collection and recreate it
            self.chroma_client.delete_collection("documents")
            self.collection = self.chroma_client.get_or_create_collection(
                name="documents",
                metadata={"hnsw:space": "cosine"}
            )
            self.documents_metadata.clear()
            logger.info("All documents cleared from RAG store")
        except Exception as e:
            logger.error(f"Failed to clear documents: {e}")
            raise
