"""Postgres/pgvector storage backend for LlamaIndex RAG.

Requires:
  uv sync --extra rag --extra pgvector

Environment:
  - `DATABASE_URL` (or pass `database_url=`) like:
      postgresql://user:pass@localhost:5432/vector_db
"""

from __future__ import annotations

# pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false

import os
from pathlib import Path
from typing import Any, Dict
from ketju.rag.llama_index.contracts import LlamaIndexStorageBackend


class PgVectorStorageBackend(LlamaIndexStorageBackend):
    """Storage strategy using Postgres + pgvector + postgres doc/index stores."""

    def __init__(
        self,
        *,
        database_url: str | None = None,
        table_name: str = "ketju_vectors",
        schema_name: str = "public",
        embed_dim: int | None = None,
        hybrid_search: bool = False,
        text_search_config: str | None = None,
        hnsw_kwargs: dict[str, Any] | None = None,
    ) -> None:
        self._storage_contexts: Dict[str, Any] = {}
        self.database_url = database_url or os.getenv("DATABASE_URL")
        if not self.database_url:
            raise ValueError("DATABASE_URL is required (or pass database_url=...).")

        self.table_name = table_name
        self.schema_name = schema_name
        self.embed_dim = embed_dim
        self.hybrid_search = hybrid_search
        self.text_search_config = text_search_config
        self.hnsw_kwargs = hnsw_kwargs or {
            "hnsw_m": 16,
            "hnsw_ef_construction": 64,
            "hnsw_ef_search": 40,
            "hnsw_dist_method": "vector_cosine_ops",
        }

    def _db_url(self):
        from sqlalchemy import make_url  # type: ignore

        return make_url(self.database_url or "")

    def _get_vector_store(self) -> Any:
        try:
            from llama_index.vector_stores.postgres import PGVectorStore  # type: ignore
            from llama_index.core import Settings  # type: ignore
        except ModuleNotFoundError as e:
            raise ModuleNotFoundError(
                "pgvector backend not installed. Install with:\n"
                "  uv sync --extra rag --extra pgvector\n"
                "Or (pip):\n"
                "  pip install 'ketju[rag,pgvector]'"
            ) from e

        url = self._db_url()

        dim = self.embed_dim
        if dim is None:
            # Best-effort: some embedding models expose dimension as `embed_dim`.
            dim = getattr(Settings.embed_model, "embed_dim", None)
        if dim is None:
            raise ValueError(
                "Embedding dimension is unknown. Pass embed_dim=... when constructing PgVectorStorageBackend."
            )

        params: dict[str, Any] = {
            "database": url.database or "postgres",
            "host": url.host or "localhost",
            "password": url.password or "",
            "port": url.port or 5432,
            "user": url.username or "postgres",
            "table_name": self.table_name,
            "schema_name": self.schema_name,
            "embed_dim": int(dim),
            "hybrid_search": self.hybrid_search,
            "hnsw_kwargs": self.hnsw_kwargs,
        }
        # PGVectorStore validates this as a string if provided; omit when None.
        if self.text_search_config is not None:
            params["text_search_config"] = self.text_search_config

        return PGVectorStore.from_params(**params)  # type: ignore[no-any-return]

    def _get_docstore(self) -> Any:
        try:
            from llama_index.storage.docstore.postgres import PostgresDocumentStore  # type: ignore
        except ModuleNotFoundError as e:
            raise ModuleNotFoundError(
                "Postgres docstore backend not installed. Install with:\n"
                "  uv sync --extra rag --extra pgvector"
            ) from e

        db_url = self.database_url
        if db_url is None:  # pragma: no cover
            raise ValueError("DATABASE_URL is required (or pass database_url=...).")

        return PostgresDocumentStore.from_uri(  # type: ignore[no-any-return]
            uri=db_url,
            table_name=f"{self.table_name}_docstore",
            schema_name=self.schema_name,
        )

    def _get_index_store(self) -> Any:
        try:
            from llama_index.storage.index_store.postgres import PostgresIndexStore  # type: ignore
        except ModuleNotFoundError as e:
            raise ModuleNotFoundError(
                "Postgres index-store backend not installed. Install with:\n"
                "  uv sync --extra rag --extra pgvector"
            ) from e

        db_url = self.database_url
        if db_url is None:  # pragma: no cover
            raise ValueError("DATABASE_URL is required (or pass database_url=...).")

        return PostgresIndexStore.from_uri(  # type: ignore[no-any-return]
            uri=db_url,
            table_name=f"{self.table_name}_indexstore",
            schema_name=self.schema_name,
        )

    def _get_graph_store(self) -> Any:
        # Keep graph stores in-memory by default (no file persistence).
        from llama_index.core.graph_stores.simple import SimpleGraphStore  # type: ignore

        return SimpleGraphStore()

    def _get_property_graph_store(self) -> Any:
        from llama_index.core.graph_stores.simple_labelled import (  # type: ignore
            SimplePropertyGraphStore,
        )

        return SimplePropertyGraphStore()

    def create_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> Any:
        # Build storage context without `persist_dir` to avoid disk-backed stores.
        from llama_index.core import StorageContext  # type: ignore

        return StorageContext.from_defaults(
            vector_store=self._get_vector_store(),
            index_store=self._get_index_store(),
            docstore=self._get_docstore(),
            graph_store=self._get_graph_store(),
            property_graph_store=self._get_property_graph_store(),
        )

    def get_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> Any:
        if name in self._storage_contexts:
            return self._storage_contexts[name]
        storage_context = self.create_storage_context(index_dir=index_dir, data_dir=data_dir, name=name)
        self._storage_contexts[name] = storage_context
        return storage_context

    def save(self, *, vector_index: Any, index_dir: Path, data_dir: Path) -> None:
        # Vector/doc/index stores are already persisted in Postgres.
        return

