from __future__ import annotations

from pathlib import Path
from typing import Dict, List, Protocol, Sequence

import anyio
import chromadb
from chromadb.api import ClientAPI
from chromadb.api import Collection as CollectionType

from ..interfaces import VectorStoreWriter
from ..models import Chunk


class ChromaClientFactory(Protocol):
    def client(self) -> ClientAPI: ...


class LocalChromaClientFactory(ChromaClientFactory):
    """Lazy factory for PersistentClient instances."""

    def __init__(self, *, persist_dir: Path) -> None:
        self._persist_dir = persist_dir
        self._client: ClientAPI | None = None

    def client(self) -> ClientAPI:
        if self._client is None:
            self._persist_dir.mkdir(parents=True, exist_ok=True)
            self._client = chromadb.PersistentClient(path=str(self._persist_dir))
        return self._client


class InMemoryChromaClientFactory(ChromaClientFactory):
    """Factory that returns an in-memory Chroma client for tests."""

    def __init__(self) -> None:
        self._client: ClientAPI | None = None

    def client(self) -> ClientAPI:
        if self._client is None:
            self._client = chromadb.Client()
        return self._client


class ChromaVectorStoreWriter(VectorStoreWriter):
    """Persists chunks + embeddings into a Chroma collection."""

    def __init__(
        self,
        *,
        client_factory: ChromaClientFactory,
        batch_size: int = 64,
    ) -> None:
        self._client_factory = client_factory
        self._batch_size = batch_size
        self._collections: Dict[str, CollectionType] = {}

    def _collection(self, name: str) -> CollectionType:
        if name not in self._collections:
            client = self._client_factory.client()
            self._collections[name] = client.get_or_create_collection(name=name)
        return self._collections[name]

    async def write(
        self,
        collection_name: str,
        chunks: Sequence[Chunk],
        embeddings: Sequence[Sequence[float]],
    ) -> Sequence[str]:
        if len(chunks) != len(embeddings):
            raise ValueError("Chunk count must match embedding count.")

        def _run() -> List[str]:
            collection = self._collection(collection_name)
            ids: List[str] = []
            for start in range(0, len(chunks), self._batch_size):
                end = start + self._batch_size
                batch_chunks = chunks[start:end]
                batch_embeddings = embeddings[start:end]
                batch_embeddings_list: list[Sequence[float]] = list(batch_embeddings)
                collection.add(
                    ids=[chunk.chunk_id for chunk in batch_chunks],
                    documents=[chunk.text for chunk in batch_chunks],
                    metadatas=[dict(chunk.metadata) for chunk in batch_chunks],
                    embeddings=batch_embeddings_list,
                )
                ids.extend(chunk.chunk_id for chunk in batch_chunks)
            return ids

        return await anyio.to_thread.run_sync(_run)
