import chromadb
from chromadb.config import Settings as ChromaSettings
from app.config import settings
from app.services.embedding_service import get_embedding_service
from typing import List, Dict, Any
import logging

logger = logging.getLogger(__name__)


class VectorStore:
    def __init__(self):
        self._client = chromadb.PersistentClient(
            path=settings.chroma_persist_directory,
            settings=ChromaSettings(anonymized_telemetry=False),
        )
        self._collection = self._client.get_or_create_collection(
            name=settings.chroma_collection_name,
            metadata={"hnsw:space": "cosine"},
        )
        logger.info(f"ChromaDB collection '{settings.chroma_collection_name}' ready")

    def upsert(
        self,
        ids: List[str],
        embeddings: List[List[float]],
        documents: List[str],
        metadatas: List[Dict[str, Any]],
    ) -> None:
        self._collection.upsert(
            ids=ids,
            embeddings=embeddings,
            documents=documents,
            metadatas=metadatas,
        )

    def delete_by_sources(self, sources: List[str]) -> None:
        """Remove all chunks belonging to the given source files. Called before
        re-ingesting a file so edited documents don't leave stale vectors
        behind (chunk IDs are content-hashed, so edits change the IDs)."""
        if not sources:
            return
        self._collection.delete(where={"source": {"$in": sources}})

    def query(self, embedding: List[float], top_k: int) -> Dict[str, Any]:
        return self._collection.query(
            query_embeddings=[embedding],
            n_results=top_k,
            include=["documents", "metadatas", "distances"],
        )

    def count(self) -> int:
        return self._collection.count()

    def reset(self) -> None:
        self._client.delete_collection(settings.chroma_collection_name)
        self._collection = self._client.get_or_create_collection(
            name=settings.chroma_collection_name,
            metadata={"hnsw:space": "cosine"},
        )
        logger.info("Collection reset")

    def list_sources(self) -> List[Dict[str, Any]]:
        """Return per-source stats aggregated from all chunk metadata.
        Category is read from stored metadata when available, otherwise
        derived from the top-level folder in the source path so that
        documents ingested before the category field was added still
        display correctly."""
        from pathlib import Path as _Path

        result = self._collection.get(include=["metadatas"])
        metadatas = result.get("metadatas") or []

        sources: Dict[str, Dict[str, Any]] = {}
        for meta in metadatas:
            source = meta.get("source", "")
            stored_cat = meta.get("category", "")
            if not stored_cat:
                parts = _Path(source).parts
                stored_cat = parts[0] if len(parts) > 1 else ""

            if source not in sources:
                sources[source] = {
                    "source": source,
                    "category": stored_cat,
                    "title": meta.get("title", ""),
                    "chunk_count": 0,
                }
            sources[source]["chunk_count"] += 1

        return sorted(sources.values(), key=lambda x: x["source"])

    def is_healthy(self) -> bool:
        try:
            self._collection.count()
            return True
        except Exception:
            return False


_instance: VectorStore | None = None


def get_vector_store() -> VectorStore:
    global _instance
    if _instance is None:
        _instance = VectorStore()
    return _instance