"""Vector store using ChromaDB for semantic memory search."""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

from loguru import logger

from miobot.utils.helpers import ensure_dir

if TYPE_CHECKING:
    from miobot.agent.embedding import EmbeddingProvider


@dataclass
class MemoryChunk:
    """A memory chunk with embedding."""
    id: str
    content: str
    metadata: dict
    timestamp: str
    score: float = 0.0  # Similarity score from search


class VectorStore:
    """ChromaDB-based vector store for semantic memory search."""

    def __init__(
        self,
        workspace: Path,
        embedding_provider: EmbeddingProvider,
        collection_name: str = "memory",
    ):
        self._embedding_provider = embedding_provider
        
        # Initialize ChromaDB
        try:
            import chromadb
        except ImportError:
            raise ImportError(
                "chromadb not installed. "
                "Install with: pip install chromadb"
            )
        
        # Persist storage path
        self._store_dir = ensure_dir(workspace / "vector_store")
        self._client = chromadb.PersistentClient(path=str(self._store_dir))
        
        # Get or create collection
        self._collection = self._client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"}
        )
        
        logger.info(
            "VectorStore initialized: collection={}, dimension={}",
            collection_name,
            embedding_provider.dimension
        )

    def add(
        self,
        id: str,
        content: str,
        metadata: dict | None = None,
    ) -> None:
        """Add a memory chunk to the store."""
        embedding = self._embedding_provider.encode(content)
        
        # ChromaDB requires metadata to be a non-empty dict
        if metadata:
            self._collection.add(
                ids=[id],
                embeddings=[embedding],
                documents=[content],
                metadatas=[metadata]
            )
        else:
            self._collection.add(
                ids=[id],
                embeddings=[embedding],
                documents=[content]
            )
        
        logger.debug("Added memory chunk: {}", id)

    def add_batch(
        self,
        ids: list[str],
        contents: list[str],
        metadata: list[dict] | None = None,
    ) -> None:
        """Add multiple memory chunks at once."""
        embeddings = self._embedding_provider.encode_batch(contents)
        
        if metadata:
            self._collection.add(
                ids=ids,
                embeddings=embeddings,
                documents=contents,
                metadatas=metadata
            )
        else:
            self._collection.add(
                ids=ids,
                embeddings=embeddings,
                documents=contents
            )
        
        logger.debug("Added {} memory chunks", len(ids))

    def search(
        self,
        query: str,
        n_results: int = 5,
        where: dict | None = None,
    ) -> list[MemoryChunk]:
        """Search for similar memory chunks."""
        query_embedding = self._embedding_provider.encode(query)
        
        results = self._collection.query(
            query_embeddings=[query_embedding],
            n_results=n_results,
            where=where,
        )
        
        chunks = []
        if results['ids'] and results['ids'][0]:
            for i, id in enumerate(results['ids'][0]):
                meta = results['metadatas'][0][i] if results['metadatas'] and results['metadatas'][0] else None
                metadata = meta or {}
                distance = results['distances'][0][i] if results.get('distances') else 0
                chunk = MemoryChunk(
                    id=id,
                    content=results['documents'][0][i],
                    metadata=metadata,
                    timestamp=metadata.get('timestamp', ''),
                    score=1 - distance,  # Convert distance to similarity
                )
                chunks.append(chunk)
        
        logger.debug("Search '{}' returned {} results", query, len(chunks))
        return chunks

    def get(self, id: str) -> MemoryChunk | None:
        """Get a memory chunk by ID."""
        result = self._collection.get(ids=[id])
        
        if not result['ids']:
            return None
        
        meta = result['metadatas'][0] if result['metadatas'] else None
        metadata = meta or {}
        return MemoryChunk(
            id=result['ids'][0],
            content=result['documents'][0],
            metadata=metadata,
            timestamp=metadata.get('timestamp', ''),
        )

    def delete(self, id: str) -> None:
        """Delete a memory chunk by ID."""
        self._collection.delete(ids=[id])
        logger.debug("Deleted memory chunk: {}", id)

    def update(
        self,
        id: str,
        content: str | None = None,
        metadata: dict | None = None,
    ) -> None:
        """Update a memory chunk."""
        update_dict = {}
        
        if content is not None:
            embedding = self._embedding_provider.encode(content)
            update_dict['embeddings'] = [embedding]
            update_dict['documents'] = [content]
        
        if metadata is not None:
            update_dict['metadatas'] = [metadata]
        
        if update_dict:
            self._collection.update(ids=[id], **update_dict)
            logger.debug("Updated memory chunk: {}", id)

    def count(self) -> int:
        """Return the number of memory chunks."""
        return self._collection.count()

    def clear(self) -> None:
        """Clear all memory chunks."""
        self._client.delete_collection(self._collection.name)
        self._collection = self._client.get_or_create_collection(
            name=self._collection.name,
            metadata={"hnsw:space": "cosine"}
        )
        logger.info("VectorStore cleared")

    def list_all(self, limit: int = 100) -> list[MemoryChunk]:
        """List all memory chunks (for debugging)."""
        result = self._collection.get(limit=limit)
        
        chunks = []
        if result['ids']:
            for i, id in enumerate(result['ids']):
                meta = result['metadatas'][i] if result['metadatas'] else None
                metadata = meta or {}
                chunks.append(MemoryChunk(
                    id=id,
                    content=result['documents'][i],
                    metadata=metadata,
                    timestamp=metadata.get('timestamp', ''),
                ))
        
        return chunks
