"""Vector store — ChromaDB integration for semantic search.

The ChromaDB index is treated as a derived/cached store — it lives in data/
(git-ignored) and is rebuilt on startup from git-tracked content:
  - docs/**/*.md (agent-owned documents)
  - state/knowledge.json (knowledge bank entries)
"""

from __future__ import annotations

import json
import logging
import sys
import uuid
from pathlib import Path

logger = logging.getLogger(__name__)


def _chromadb_settings():
    """Build ChromaDB Settings. Use Segment API (pure Python) when running in PyInstaller bundle."""
    from chromadb.config import Settings

    kwargs = {"anonymized_telemetry": False}
    if getattr(sys, "frozen", False):
        kwargs["chroma_api_impl"] = "chromadb.api.segment.SegmentAPI"
    return Settings(**kwargs)


class VectorStore:
    """ChromaDB-backed vector store for semantic search over knowledge entries."""

    def __init__(self, persist_dir: Path) -> None:
        self.persist_dir = persist_dir
        self._client = None
        self._collection = None

    def initialize(self) -> None:
        """Initialize the ChromaDB client and collection."""
        import chromadb

        self.persist_dir.mkdir(parents=True, exist_ok=True)
        self._client = chromadb.PersistentClient(
            path=str(self.persist_dir),
            settings=_chromadb_settings(),
        )
        self._collection = self._client.get_or_create_collection(
            name="vastian_knowledge_bank",
            metadata={"description": "Vastian Network knowledge entries"},
        )
        logger.info(
            "VectorStore initialized at %s (entries: %d)",
            self.persist_dir,
            self._collection.count(),
        )

    def add(self, content: str, metadata: dict | None = None) -> str:
        """Add a document to the vector store. Returns the generated ID."""
        assert self._collection is not None, "VectorStore not initialized"

        doc_id = uuid.uuid4().hex[:12]
        # ChromaDB metadata values must be str, int, float, or bool
        clean_metadata = {}
        if metadata:
            for k, v in metadata.items():
                if isinstance(v, (str, int, float, bool)):
                    clean_metadata[k] = v
                else:
                    clean_metadata[k] = str(v)

        self._collection.add(
            ids=[doc_id],
            documents=[content],
            metadatas=[clean_metadata] if clean_metadata else None,
        )
        return doc_id

    def search(
        self,
        query: str,
        top_k: int = 5,
        where: dict | None = None,
    ) -> list[dict]:
        """Search for similar documents. Returns list of {id, content, metadata, distance}."""
        assert self._collection is not None, "VectorStore not initialized"

        kwargs: dict = {
            "query_texts": [query],
            "n_results": top_k,
        }
        if where:
            kwargs["where"] = where

        try:
            results = self._collection.query(**kwargs)
        except Exception:
            logger.exception("Vector search failed")
            return []

        entries = []
        if results and results["ids"]:
            for i, doc_id in enumerate(results["ids"][0]):
                entries.append(
                    {
                        "id": doc_id,
                        "content": results["documents"][0][i] if results["documents"] else "",
                        "metadata": results["metadatas"][0][i] if results["metadatas"] else {},
                        "distance": results["distances"][0][i] if results["distances"] else None,
                    }
                )
        return entries

    def delete(self, doc_id: str) -> None:
        """Delete a document by ID."""
        assert self._collection is not None, "VectorStore not initialized"
        self._collection.delete(ids=[doc_id])

    @property
    def count(self) -> int:
        if self._collection is None:
            return 0
        return self._collection.count()

    def rebuild_index(
        self,
        docs_dir: Path | None = None,
        state_dir: Path | None = None,
        force: bool = False,
    ) -> int:
        """Rebuild the vector index from git-tracked content.

        Clears the existing collection and re-indexes:
          1. All Markdown files in docs/**/*.md
          2. All knowledge entries from state/knowledge.json

        By default, only rebuilds if the index has fewer than 10 entries
        (indicating a fresh machine). Pass force=True to always rebuild.

        Returns the total number of documents indexed, or -1 if skipped.
        """
        if self._collection is None:
            logger.debug("VectorStore not initialized — skipping rebuild")
            return -1

        current_count = self._collection.count()
        if not force and current_count >= 10:
            logger.info(
                "VectorStore already has %d entries — skipping rebuild (use force=True to override)",
                current_count,
            )
            return -1

        # Clear existing data
        try:
            existing_ids = self._collection.get()["ids"]
            if existing_ids:
                # Delete in batches to avoid oversized requests
                batch_size = 500
                for i in range(0, len(existing_ids), batch_size):
                    self._collection.delete(ids=existing_ids[i : i + batch_size])
        except Exception:
            logger.debug("Could not clear existing collection, continuing")

        total = 0

        # 1. Index docs/**/*.md
        if docs_dir and docs_dir.exists():
            for md_file in sorted(docs_dir.rglob("*.md")):
                try:
                    content = md_file.read_text(encoding="utf-8")
                    if not content.strip():
                        continue
                    # Use relative path as stable ID so re-indexing is idempotent
                    rel_path = str(md_file.relative_to(docs_dir))
                    doc_id = f"doc-{rel_path.replace(chr(92), '-').replace('/', '-')}"

                    # Chunk large documents (ChromaDB has token limits)
                    chunks = _chunk_text(content, max_chars=2000)
                    for ci, chunk in enumerate(chunks):
                        chunk_id = f"{doc_id}-{ci}" if len(chunks) > 1 else doc_id
                        self._collection.add(
                            ids=[chunk_id],
                            documents=[chunk],
                            metadatas=[
                                {
                                    "type": "document",
                                    "path": rel_path,
                                    "author": md_file.parent.name,
                                }
                            ],
                        )
                        total += 1
                except Exception:
                    logger.debug("Failed to index %s", md_file)

        # 2. Index state/knowledge.json
        knowledge_path = (state_dir or Path("./state")) / "knowledge.json"
        if knowledge_path.exists():
            try:
                entries = json.loads(knowledge_path.read_text(encoding="utf-8"))
                if isinstance(entries, list):
                    for entry in entries:
                        content = entry.get("content", "")
                        if not content.strip():
                            continue
                        entry_id = entry.get("id", uuid.uuid4().hex[:12])
                        metadata = {
                            "type": "knowledge_entry",
                            "author": entry.get("author", ""),
                            "tags": entry.get("tags", ""),
                        }
                        self._collection.add(
                            ids=[entry_id],
                            documents=[content[:2000]],
                            metadatas=[metadata],
                        )
                        total += 1
            except Exception:
                logger.debug("Failed to index knowledge.json")

        logger.info("VectorStore rebuilt: %d documents indexed", total)
        return total


def _chunk_text(text: str, max_chars: int = 2000) -> list[str]:
    """Split text into chunks of approximately max_chars, breaking at newlines."""
    if len(text) <= max_chars:
        return [text]

    chunks: list[str] = []
    current = ""
    for line in text.split("\n"):
        if len(current) + len(line) + 1 > max_chars and current:
            chunks.append(current)
            current = line
        else:
            current = current + "\n" + line if current else line
    if current:
        chunks.append(current)
    return chunks
