from __future__ import annotations

import os
from pathlib import Path
from typing import List, Optional

from langchain_core.documents import Document

from src.config import settings
from src.logger import logger


def get_vectorstore(
    documents: List[Document],
    embedding_model,
    collection_name: str,
    persist_directory: str,
    store_type: Optional[str] = None,
):
    """
    Create and populate a vector store.
    store_type: "faiss" | "chroma" — if omitted, falls back to settings.vector_db_type.
    """
    store_type = (store_type or settings.vector_db_type or "faiss").lower().strip()

    logger.info(
        "Building vector store",
        context={
            "type": store_type,
            "docs": len(documents),
            "collection": collection_name,
            "persist_dir": persist_directory,
        },
    )

    Path(persist_directory).mkdir(parents=True, exist_ok=True)

    if store_type == "chroma":
        from langchain_chroma import Chroma

        store = Chroma(
            collection_name=collection_name,
            persist_directory=persist_directory,
            embedding_function=embedding_model,
        )
        batch_size = int(os.getenv("VECTORSTORE_BATCH_SIZE", "50"))
        total = len(documents)
        for i in range(0, total, batch_size):
            end = min(i + batch_size, total)
            logger.info(
                f"Ingesting batch {i // batch_size + 1}/{(total - 1) // batch_size + 1}"
            )
            store.add_documents(documents[i:end])
        logger.info("Chroma ingestion complete", context={"path": persist_directory})
        return store

    if store_type == "faiss":
        from langchain_community.vectorstores import FAISS

        # Prefer Embeddings object; ST wrappers should implement embed_documents/embed_query
        store = FAISS.from_documents(documents, embedding_model)
        store.save_local(persist_directory)
        logger.info("FAISS store saved", context={"path": persist_directory})
        return store

    raise ValueError(f"Unsupported vector_db_type: {store_type}")