"""
ChromaDB Wrapper for MCP Compatibility

This module provides a compatibility wrapper for ChromaDB 1.0.4
to work with MIRA's consciousness evolution platform.
"""

# Apply SQLite fix BEFORE importing ChromaDB
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
import chromadb_sqlite_fix

import chromadb
from chromadb.config import Settings
from pathlib import Path
import json
import time
from typing import Dict, List, Optional, Any
import logging

logger = logging.getLogger(__name__)


class ChromaDBWrapper:
    """
    A wrapper around ChromaDB 0.3.29 that provides a stable interface
    for MIRA's consciousness evolution platform.
    """
    
    def __init__(self, persist_directory: str = ".mira/chromadb"):
        """Initialize ChromaDB wrapper with compatibility fixes."""
        self.persist_directory = Path(persist_directory)
        self.persist_directory.mkdir(parents=True, exist_ok=True)
        
        # Initialize ChromaDB client
        self._init_client()
        
        # Track collections internally
        self._collections = {}
        
    def _init_client(self):
        """Initialize ChromaDB client with proper settings."""
        # Use PersistentClient for ChromaDB 1.0.4
        self.client = chromadb.PersistentClient(
            path=str(self.persist_directory),
            settings=Settings(
                anonymized_telemetry=False,
                allow_reset=True
            )
        )
    
    def create_collection(self, name: str, metadata: Optional[Dict] = None) -> 'CollectionWrapper':
        """Create a new collection with wrapper."""
        try:
            # Delete if exists to ensure clean creation
            try:
                self.client.delete_collection(name)
            except:
                pass
            
            # Ensure metadata is not empty (ChromaDB 1.0.4 requirement)
            if not metadata:
                metadata = {'created_by': 'MIRA', 'version': '1.0.4'}
            
            # Create collection
            collection = self.client.create_collection(
                name=name,
                metadata=metadata
            )
            
            # Wrap in our wrapper
            wrapped = CollectionWrapper(collection, self)
            self._collections[name] = wrapped
            
            logger.info(f"Created collection: {name}")
            return wrapped
            
        except Exception as e:
            logger.error(f"Failed to create collection {name}: {e}")
            raise
    
    def get_collection(self, name: str) -> Optional['CollectionWrapper']:
        """Get an existing collection with wrapper."""
        try:
            # Check cache first
            if name in self._collections:
                return self._collections[name]
            
            # Get from ChromaDB
            collection = self.client.get_collection(name)
            wrapped = CollectionWrapper(collection, self)
            self._collections[name] = wrapped
            
            return wrapped
            
        except Exception as e:
            logger.debug(f"Collection {name} not found: {e}")
            return None
    
    def get_or_create_collection(self, name: str, metadata: Optional[Dict] = None) -> 'CollectionWrapper':
        """Get or create a collection with wrapper."""
        collection = self.get_collection(name)
        if collection is None:
            # Ensure metadata is not empty (ChromaDB 1.0.4 requirement)
            if not metadata:
                metadata = {'created_by': 'MIRA', 'version': '1.0.4'}
            collection = self.create_collection(name, metadata)
        return collection
    
    def list_collections(self) -> List[str]:
        """List all collection names."""
        try:
            collections = self.client.list_collections()
            if collections is None:
                return []
            return [col.name for col in collections]
        except Exception as e:
            logger.error(f"Error listing collections: {e}")
            return []
    
    def delete_collection(self, name: str):
        """Delete a collection."""
        try:
            self.client.delete_collection(name)
            if name in self._collections:
                del self._collections[name]
            logger.info(f"Deleted collection: {name}")
        except Exception as e:
            logger.error(f"Failed to delete collection {name}: {e}")


class CollectionWrapper:
    """
    A wrapper around ChromaDB Collection that ensures compatibility.
    """
    
    def __init__(self, collection, client: ChromaDBWrapper):
        """Initialize collection wrapper."""
        self._collection = collection
        self._client = client
        self.name = collection.name if collection else "unknown"
        self.metadata = getattr(collection, 'metadata', {})
    
    def add(self, documents: List[str], metadatas: List[Dict], ids: List[str], 
            embeddings: Optional[List[List[float]]] = None):
        """Add documents to collection."""
        if self._collection is None:
            raise ValueError("Collection is None")
        
        # ChromaDB 1.0.4 can handle various metadata types
        cleaned_metadatas = metadatas
        
        # Add to collection
        self._collection.add(
            documents=documents,
            metadatas=cleaned_metadatas,
            ids=ids,
            embeddings=embeddings
        )
    
    def get(self, ids: Optional[List[str]] = None, where: Optional[Dict] = None, 
            limit: Optional[int] = None) -> Dict[str, Any]:
        """Get documents from collection."""
        if self._collection is None:
            return {'documents': [], 'metadatas': [], 'ids': []}
        
        return self._collection.get(ids=ids, where=where, limit=limit)
    
    def query(self, query_embeddings: Optional[List[List[float]]] = None,
              query_texts: Optional[List[str]] = None, n_results: int = 10,
              where: Optional[Dict] = None) -> Dict[str, Any]:
        """Query collection."""
        if self._collection is None:
            return {'documents': [[]], 'metadatas': [[]], 'ids': [[]], 'distances': [[]]}
        
        return self._collection.query(
            query_embeddings=query_embeddings,
            query_texts=query_texts,
            n_results=n_results,
            where=where
        )
    
    def count(self) -> int:
        """Count documents in collection."""
        if self._collection is None:
            return 0
        
        try:
            return self._collection.count()
        except:
            # Fallback for older versions
            result = self._collection.get()
            return len(result.get('ids', []))
    
    def delete(self, ids: List[str]):
        """Delete documents by ID."""
        if self._collection is None:
            return
        
        # ChromaDB 0.3.29 might not have delete, so we work around it
        if hasattr(self._collection, 'delete'):
            self._collection.delete(ids=ids)
        else:
            logger.warning("Delete method not available")


# Global wrapper instance
_wrapper_instance = None

def get_chromadb_wrapper(persist_directory: str = ".mira/chromadb") -> ChromaDBWrapper:
    """Get or create ChromaDB wrapper instance."""
    global _wrapper_instance
    if _wrapper_instance is None:
        _wrapper_instance = ChromaDBWrapper(persist_directory)
    return _wrapper_instance