from __future__ import annotations

import hashlib
import os
import re
import unicodedata
from typing import Any, Dict, Optional

_CHROMADB_IMPORT_ERROR: Optional[str] = None
_chromadb_value: Any
try:
    import chromadb as _chromadb_imported
except Exception as exc:  # pragma: no cover - exercised in runtime-specific environments
    _CHROMADB_IMPORT_ERROR = str(exc)

    class _ChromaStub:
        HttpClient = None
        PersistentClient = None

    _chromadb_value = _ChromaStub()
else:
    _chromadb_value = _chromadb_imported

chromadb: Any = _chromadb_value


class VectorStore:
    """Thin stable wrapper over Chroma collections with namespace isolation."""

    MAX_NAMESPACE_LENGTH = 80

    def __init__(self) -> None:
        self.prefix = os.getenv("CHROMA_COLLECTION_PREFIX", "grantflow")
        self._collections: Dict[str, Any] = {}
        self._memory_store: Dict[str, Dict[str, Any]] = {}
        self._client_init_error: Optional[str] = None
        forced_memory_backend = str(os.getenv("GRANTFLOW_FORCE_INMEM_VECTOR_STORE", "")).strip().lower()
        self._force_inmem = forced_memory_backend in {"1", "true", "yes", "on"}

        self._chroma_host = os.getenv("CHROMA_HOST")
        # Keep a non-API default to avoid common localhost conflict with uvicorn (8000).
        self._chroma_port = int(os.getenv("CHROMA_PORT", "8001"))
        self._persist_dir = (
            os.getenv("CHROMA_PERSIST_DIRECTORY")
            or os.getenv("GRANTFLOW_CHROMA_DIR")
            or os.getenv("AIDGRAPH_CHROMA_DIR")
            or "./chroma_db"
        )
        self.client: Any = None

        if self._force_inmem:
            self._client_init_error = "vector store backend forced to in-memory via GRANTFLOW_FORCE_INMEM_VECTOR_STORE"
            self.client = None
            return

        if _CHROMADB_IMPORT_ERROR:
            self._client_init_error = f"chromadb import failed: {_CHROMADB_IMPORT_ERROR}"
            self.client = None
            return

        try:
            if self._chroma_host:
                self.client = chromadb.HttpClient(host=self._chroma_host, port=self._chroma_port)
            else:
                self.client = chromadb.PersistentClient(path=self._persist_dir)
        except Exception as exc:
            # Sandbox/runtime fallback: keep API alive and use in-memory vectors.
            self._client_init_error = str(exc)
            self.client = None

    @staticmethod
    def normalize_namespace(namespace: str) -> str:
        raw = str(namespace or "default").strip().lower()
        if not raw:
            return "default"

        # Keep normalization stable across unicode inputs by folding to ASCII where possible.
        # If folding drops all symbols, we still return a deterministic hashed namespace token.
        folded = unicodedata.normalize("NFKD", raw).encode("ascii", "ignore").decode("ascii")
        if folded.strip():
            raw = folded

        raw = raw.replace("/", "_").replace("\\", "_").replace(":", "_").replace(" ", "_")
        # Chroma collection names are safer with a constrained token set.
        normalized = re.sub(r"[^a-z0-9_.-]+", "_", raw)
        normalized = re.sub(r"_+", "_", normalized).strip("._-")
        if not normalized:
            digest = hashlib.sha1(raw.encode("utf-8", errors="ignore")).hexdigest()[:12]
            normalized = f"ns_{digest}"
        if len(normalized) > VectorStore.MAX_NAMESPACE_LENGTH:
            digest = hashlib.sha1(normalized.encode("utf-8", errors="ignore")).hexdigest()[:10]
            head_len = max(8, VectorStore.MAX_NAMESPACE_LENGTH - len(digest) - 1)
            head = normalized[:head_len].rstrip("._-") or "ns"
            normalized = f"{head}_{digest}"
        return normalized or "default"

    def _collection_name(self, namespace: str) -> str:
        ns = self.normalize_namespace(namespace)
        return f"{self.prefix}_{ns}"

    def namespace_trace(self, namespace: str) -> Dict[str, str]:
        requested = str(namespace or "").strip() or "default"
        normalized = self.normalize_namespace(requested)
        return {
            "namespace": requested,
            "namespace_normalized": normalized,
            "collection": self._collection_name(requested),
        }

    def _embed_texts(self, texts: list[str], dims: int = 16) -> list[list[float]]:
        """Deterministic lightweight embeddings for local/offline MVP smoke tests."""
        embeddings: list[list[float]] = []
        for text in texts:
            vec = [0.0] * dims
            data = (text or "").encode("utf-8", errors="ignore")
            digest = hashlib.sha256(data).digest()
            for i in range(dims):
                vec[i] = digest[i] / 255.0
            embeddings.append(vec)
        return embeddings

    def _cosine_like(self, a: list[float], b: list[float]) -> float:
        return sum(x * y for x, y in zip(a, b))

    def _ensure_memory_namespace(self, namespace: str) -> Dict[str, Any]:
        name = self._collection_name(namespace)
        if name not in self._memory_store:
            self._memory_store[name] = {"name": name, "rows": {}}
        return self._memory_store[name]

    def get_collection(self, namespace: str):
        if self.client is None:
            return None
        name = self._collection_name(namespace)
        if name not in self._collections:
            self._collections[name] = self.client.get_or_create_collection(name=name)
        return self._collections[name]

    def upsert(
        self,
        namespace: str,
        ids: list[str],
        documents: list[str],
        metadatas: Optional[list[dict]] = None,
    ) -> None:
        embeddings = self._embed_texts(documents)

        if self.client is None:
            ns = self._ensure_memory_namespace(namespace)
            rows = ns["rows"]
            for i, doc_id in enumerate(ids):
                rows[doc_id] = {
                    "id": doc_id,
                    "document": documents[i],
                    "metadata": (metadatas[i] if metadatas and i < len(metadatas) else None),
                    "embedding": embeddings[i],
                }
            return

        col = self.get_collection(namespace)
        kwargs: Dict[str, Any] = {
            "ids": ids,
            "documents": documents,
            "embeddings": embeddings,
        }
        if metadatas is not None:
            kwargs["metadatas"] = metadatas
        col.upsert(**kwargs)

    def query(
        self,
        namespace: str,
        query_texts: list[str] | str,
        n_results: int = 5,
        where: Optional[dict] = None,
        top_k: Optional[int] = None,
    ):
        if top_k is not None:
            n_results = top_k

        single_query = isinstance(query_texts, str)
        if single_query:
            query_list: list[str] = [str(query_texts)]
        else:
            query_list = [str(item or "") for item in query_texts]

        if self.client is None:
            ns = self._ensure_memory_namespace(namespace)
            rows = list(ns["rows"].values())
            docs_out = []
            metas_out = []
            ids_out = []
            for q in query_list:
                q_emb = self._embed_texts([q])[0]
                ranked = sorted(
                    rows,
                    key=lambda r: self._cosine_like(q_emb, r["embedding"]),
                    reverse=True,
                )[:n_results]
                docs_out.append([r["document"] for r in ranked])
                metas_out.append([r["metadata"] for r in ranked])
                ids_out.append([r["id"] for r in ranked])
            result = {"ids": ids_out, "documents": docs_out, "metadatas": metas_out}
            if single_query:
                return docs_out[0]
            return result

        col = self.get_collection(namespace)
        kwargs: Dict[str, Any] = {
            "query_embeddings": self._embed_texts(query_list),
            "n_results": n_results,
        }
        if where is not None:
            kwargs["where"] = where
        result = col.query(**kwargs)

        if single_query:
            return (result.get("documents") or [[]])[0]
        return result

    def get_stats(self, namespace: str) -> dict:
        trace = self.namespace_trace(namespace)
        if self.client is None:
            ns = self._ensure_memory_namespace(namespace)
            count = len(ns["rows"])
            return {
                "namespace": trace["namespace"],
                "namespace_normalized": trace["namespace_normalized"],
                "collection": ns["name"],
                "count": count,
                "document_count": count,
                "backend": "memory",
                "client_init_error": self._client_init_error,
            }

        col = self.get_collection(namespace)
        count = col.count()
        return {
            "namespace": trace["namespace"],
            "namespace_normalized": trace["namespace_normalized"],
            "collection": col.name,
            "count": count,
            "document_count": count,
            "backend": "chroma",
        }


vector_store = VectorStore()
