"""
Weaviate vector database implementation for Dify's RAG system.

This module provides integration with Weaviate vector database for storing and retrieving
document embeddings used in retrieval-augmented generation workflows.
"""

import atexit
import datetime
import json
import logging
import threading
import uuid as _uuid
from typing import Any, override
from urllib.parse import urlparse

import weaviate
import weaviate.classes.config as wc
from pydantic import BaseModel, model_validator
from weaviate.classes.data import DataObject
from weaviate.classes.init import Auth
from weaviate.classes.query import Filter, MetadataQuery
from weaviate.exceptions import UnexpectedStatusCodeError, WeaviateQueryError

from configs import dify_config
from core.rag.datasource.vdb.field import Field
from core.rag.datasource.vdb.vector_base import BaseVector, VectorIndexStructDict
from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
from core.rag.datasource.vdb.vector_type import VectorType
from core.rag.embedding.embedding_base import Embeddings
from core.rag.models.document import Document
from extensions.ext_redis import redis_client
from models.dataset import Dataset

logger = logging.getLogger(__name__)

_weaviate_client: weaviate.WeaviateClient | None = None
_weaviate_client_lock = threading.Lock()


def _shutdown_weaviate_client() -> None:
    """
    Best-effort shutdown hook to close the module-level Weaviate client.

    This is registered with atexit so that HTTP/gRPC resources are released
    when the Python interpreter exits.
    """
    global _weaviate_client

    # Ensure thread-safety when accessing the shared client instance
    with _weaviate_client_lock:
        client = _weaviate_client
        _weaviate_client = None

    if client is not None:
        try:
            client.close()
        except Exception:
            # Best-effort cleanup; log at debug level and ignore errors.
            logger.debug("Failed to close Weaviate client during shutdown", exc_info=True)


# Register the shutdown hook once per process.
atexit.register(_shutdown_weaviate_client)


class WeaviateConfig(BaseModel):
    """
    Configuration model for Weaviate connection settings.

    Attributes:
        endpoint: Weaviate server endpoint URL
        grpc_endpoint: Optional Weaviate gRPC server endpoint URL
        api_key: Optional API key for authentication
        batch_size: Number of objects to batch per insert operation
    """

    endpoint: str
    grpc_endpoint: str | None = None
    api_key: str | None = None
    batch_size: int = 100

    @model_validator(mode="before")
    @classmethod
    def validate_config(cls, values: dict[str, Any]) -> dict[str, Any]:
        """Validates that required configuration values are present."""
        if not values["endpoint"]:
            raise ValueError("config WEAVIATE_ENDPOINT is required")
        return values


class WeaviateVector(BaseVector):
    """
    Weaviate vector database implementation for document storage and retrieval.

    Handles creation, insertion, deletion, and querying of document embeddings
    in a Weaviate collection.
    """

    _DOCUMENT_ID_PROPERTY = "document_id"

    def __init__(self, collection_name: str, config: WeaviateConfig, attributes: list):
        """
        Initializes the Weaviate vector store.

        Args:
            collection_name: Name of the Weaviate collection
            config: Weaviate configuration settings
            attributes: List of metadata attributes to store
        """
        super().__init__(collection_name)
        self._client = self._init_client(config)
        self._attributes = attributes

    def _init_client(self, config: WeaviateConfig) -> weaviate.WeaviateClient:
        """
        Initializes and returns a connected Weaviate client.

        Configures both HTTP and gRPC connections with proper authentication.
        """
        global _weaviate_client
        if _weaviate_client and _weaviate_client.is_ready():
            return _weaviate_client

        with _weaviate_client_lock:
            if _weaviate_client and _weaviate_client.is_ready():
                return _weaviate_client

            p = urlparse(config.endpoint)
            host = p.hostname or config.endpoint.replace("https://", "").replace("http://", "")
            http_secure = p.scheme == "https"
            http_port = p.port or (443 if http_secure else 80)

            # Parse gRPC configuration
            if config.grpc_endpoint:
                # Urls without scheme won't be parsed correctly in some python versions,
                # see https://bugs.python.org/issue27657
                grpc_endpoint_with_scheme = (
                    config.grpc_endpoint if "://" in config.grpc_endpoint else f"grpc://{config.grpc_endpoint}"
                )
                grpc_p = urlparse(grpc_endpoint_with_scheme)
                grpc_host = grpc_p.hostname or "localhost"
                grpc_port = grpc_p.port or (443 if grpc_p.scheme == "grpcs" else 50051)
                grpc_secure = grpc_p.scheme == "grpcs"
            else:
                # Infer from HTTP endpoint as fallback
                grpc_host = host
                grpc_secure = http_secure
                grpc_port = 443 if grpc_secure else 50051

            client = weaviate.connect_to_custom(
                http_host=host,
                http_port=http_port,
                http_secure=http_secure,
                grpc_host=grpc_host,
                grpc_port=grpc_port,
                grpc_secure=grpc_secure,
                auth_credentials=Auth.api_key(config.api_key) if config.api_key else None,
                skip_init_checks=True,  # Skip PyPI version check to avoid unnecessary HTTP requests
            )

            if not client.is_ready():
                raise ConnectionError("Vector database is not ready")

            _weaviate_client = client
            return client

    @override
    def get_type(self) -> str:
        """Returns the vector database type identifier."""
        return VectorType.WEAVIATE

    def get_collection_name(self, dataset: Dataset) -> str:
        """
        Retrieves or generates the collection name for a dataset.

        Uses existing index structure if available, otherwise generates from dataset ID.
        """
        if dataset.index_struct_dict:
            class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
            if not class_prefix.endswith("_Node"):
                class_prefix += "_Node"
            return class_prefix

        dataset_id = dataset.id
        return Dataset.gen_collection_name_by_id(dataset_id)

    def to_index_struct(self) -> VectorIndexStructDict:
        """Returns the index structure dictionary for persistence."""
        result: VectorIndexStructDict = {
            "type": self.get_type(),
            "vector_store": {"class_prefix": self._collection_name},
        }
        return result

    @override
    def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs):
        """
        Creates a new collection and adds initial documents with embeddings.
        """
        self._create_collection()
        self.add_texts(texts, embeddings)

    def _create_collection(self):
        """
        Creates the Weaviate collection with required schema if it doesn't exist.

        Uses Redis locking to prevent concurrent creation attempts.
        """
        lock_name = f"vector_indexing_lock_{self._collection_name}"
        with redis_client.lock(lock_name, timeout=20):
            cache_key = f"vector_indexing_{self._collection_name}"
            if redis_client.get(cache_key):
                return

            try:
                if not self._client.collections.exists(self._collection_name):
                    tokenization = (
                        wc.Tokenization(dify_config.WEAVIATE_TOKENIZATION)
                        if dify_config.WEAVIATE_TOKENIZATION
                        else wc.Tokenization.WORD
                    )
                    self._client.collections.create(
                        name=self._collection_name,
                        properties=[
                            wc.Property(
                                name=Field.TEXT_KEY.value,
                                data_type=wc.DataType.TEXT,
                                tokenization=tokenization,
                            ),
                            wc.Property(name="document_id", data_type=wc.DataType.TEXT),
                            wc.Property(name="doc_id", data_type=wc.DataType.TEXT),
                            wc.Property(name="doc_type", data_type=wc.DataType.TEXT),
                            wc.Property(name="chunk_index", data_type=wc.DataType.INT),
                            wc.Property(name="is_summary", data_type=wc.DataType.BOOL),
                            wc.Property(name="original_chunk_id", data_type=wc.DataType.TEXT),
                        ],
                        vector_config=wc.Configure.Vectors.self_provided(),
                    )

                self._ensure_properties()
                redis_client.set(cache_key, 1, ex=3600)
            except Exception as e:
                logger.exception("Error creating collection %s", self._collection_name)
                raise

    def _ensure_properties(self) -> None:
        """
        Ensures all required properties exist in the collection schema.

        Adds missing properties if the collection exists but lacks them.
        """
        if not self._client.collections.exists(self._collection_name):
            return

        col = self._client.collections.use(self._collection_name)
        cfg = col.config.get()
        existing = {p.name for p in (cfg.properties or [])}

        to_add = []
        if "document_id" not in existing:
            to_add.append(wc.Property(name="document_id", data_type=wc.DataType.TEXT))
        if "doc_id" not in existing:
            to_add.append(wc.Property(name="doc_id", data_type=wc.DataType.TEXT))
        if "doc_type" not in existing:
            to_add.append(wc.Property(name="doc_type", data_type=wc.DataType.TEXT))
        if "chunk_index" not in existing:
            to_add.append(wc.Property(name="chunk_index", data_type=wc.DataType.INT))
        if "is_summary" not in existing:
            to_add.append(wc.Property(name="is_summary", data_type=wc.DataType.BOOL))
        if "original_chunk_id" not in existing:
            to_add.append(wc.Property(name="original_chunk_id", data_type=wc.DataType.TEXT))

        for prop in to_add:
            try:
                col.config.add_property(prop)
            except Exception as e:
                logger.warning("Could not add property %s: %s", prop.name, e)

    @override
    def _get_uuids(self, documents: list[Document]) -> list[str]:
        """
        Generates deterministic UUIDs for documents based on their content.

        Uses UUID5 with URL namespace to ensure consistent IDs for identical content.
        """
        URL_NAMESPACE = _uuid.UUID("6ba7b811-9dad-11d1-80b4-00c04fd430c8")

        uuids = []
        for doc in documents:
            uuid_val = _uuid.uuid5(URL_NAMESPACE, doc.page_content)
            uuids.append(str(uuid_val))

        return uuids

    @override
    def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs):
        """
        Adds documents with their embeddings to the collection.

        Batches insertions for efficiency and returns the list of inserted object IDs.
        """
        uuids = self._get_uuids(documents)
        texts = [d.page_content for d in documents]
        metadatas = [d.metadata for d in documents]

        col = self._client.collections.use(self._collection_name)
        objs: list[DataObject] = []
        ids_out: list[str] = []

        for i, text in enumerate(texts):
            props: dict[str, Any] = {Field.TEXT_KEY.value: text}
            meta = metadatas[i] or {}
            for k, v in meta.items():
                props[k] = self._json_serializable(v)

            candidate = uuids[i] if uuids else None
            uid = candidate if (candidate and self._is_uuid(candidate)) else str(_uuid.uuid4())
            ids_out.append(uid)

            vec_payload = None
            if embeddings and i < len(embeddings) and embeddings[i]:
                vec_payload = {"default": embeddings[i]}

            objs.append(
                DataObject(
                    uuid=uid,
                    properties=props,  # type: ignore[arg-type]  # mypy incorrectly infers DataObject signature
                    vector=vec_payload,
                )
            )

        with col.batch.dynamic() as batch:
            for obj in objs:
                batch.add_object(properties=obj.properties, uuid=obj.uuid, vector=obj.vector)

        return ids_out

    def _is_uuid(self, val: str) -> bool:
        """Validates whether a string is a valid UUID format."""
        try:
            _uuid.UUID(str(val))
            return True
        except Exception:
            return False

    @override
    def delete_by_metadata_field(self, key: str, value: str) -> None:
        """Deletes all objects matching a specific metadata field value."""
        if not self._client.collections.exists(self._collection_name):
            return

        col = self._client.collections.use(self._collection_name)
        col.data.delete_many(where=Filter.by_property(key).equal(value))

    @override
    def delete(self):
        """Deletes the entire collection from Weaviate."""
        if self._client.collections.exists(self._collection_name):
            self._client.collections.delete(self._collection_name)

    @override
    def text_exists(self, id: str) -> bool:
        """Checks if a document with the given doc_id exists in the collection."""
        if not self._client.collections.exists(self._collection_name):
            return False

        col = self._client.collections.use(self._collection_name)
        res = col.query.fetch_objects(
            filters=Filter.by_property("doc_id").equal(id),
            limit=1,
            return_properties=["doc_id"],
        )

        return len(res.objects) > 0

    @override
    def delete_by_ids(self, ids: list[str]) -> None:
        """
        Deletes objects by their UUID identifiers.

        Silently ignores 404 errors for non-existent IDs.
        """
        if not self._client.collections.exists(self._collection_name):
            return

        col = self._client.collections.use(self._collection_name)

        for uid in ids:
            try:
                col.data.delete_by_id(uid)
            except UnexpectedStatusCodeError as e:
                if getattr(e, "status_code", None) != 404:
                    raise

    @override
    def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
        """
        Performs vector similarity search using the provided query vector.

        Filters by document IDs if provided and applies score threshold.
        Returns documents sorted by relevance score.
        """
        if not self._client.collections.exists(self._collection_name):
            return []

        col = self._client.collections.use(self._collection_name)
        props = list({*self._attributes, self._DOCUMENT_ID_PROPERTY, Field.TEXT_KEY.value})

        where = None
        doc_ids = kwargs.get("document_ids_filter") or []
        if doc_ids:
            where = Filter.by_property(self._DOCUMENT_ID_PROPERTY).contains_any(doc_ids)

        top_k = int(kwargs.get("top_k", 4))
        score_threshold = float(kwargs.get("score_threshold") or 0.0)

        try:
            res = col.query.near_vector(
                near_vector=query_vector,
                limit=top_k,
                return_properties=props,
                return_metadata=MetadataQuery(distance=True),
                include_vector=False,
                filters=where,
                target_vector="default",
            )
        except WeaviateQueryError:
            self._ensure_properties()
            res = col.query.near_vector(
                near_vector=query_vector,
                limit=top_k,
                return_properties=props,
                return_metadata=MetadataQuery(distance=True),
                include_vector=False,
                filters=where,
                target_vector="default",
            )

        docs: list[Document] = []
        for obj in res.objects:
            properties = dict(obj.properties or {})
            text = properties.pop(Field.TEXT_KEY.value, "")
            if obj.metadata and obj.metadata.distance is not None:
                distance = obj.metadata.distance
            else:
                distance = 1.0
            score = 1.0 - distance

            if score > score_threshold:
                properties["score"] = score
                docs.append(Document(page_content=text, metadata=properties))

        docs.sort(key=lambda d: d.metadata.get("score", 0.0), reverse=True)
        return docs

    @override
    def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
        """
        Performs BM25 full-text search on document content.

        Filters by document IDs if provided and returns matching documents with vectors.
        """
        if not self._client.collections.exists(self._collection_name):
            return []

        col = self._client.collections.use(self._collection_name)
        props = list({*self._attributes, Field.TEXT_KEY.value})

        where = None
        doc_ids = kwargs.get("document_ids_filter") or []
        if doc_ids:
            where = Filter.by_property(self._DOCUMENT_ID_PROPERTY).contains_any(doc_ids)

        top_k = int(kwargs.get("top_k", 4))

        try:
            res = col.query.bm25(
                query=query,
                query_properties=[Field.TEXT_KEY.value],
                limit=top_k,
                return_properties=props,
                include_vector=True,
                filters=where,
            )
        except WeaviateQueryError:
            self._ensure_properties()
            res = col.query.bm25(
                query=query,
                query_properties=[Field.TEXT_KEY.value],
                limit=top_k,
                return_properties=props,
                include_vector=True,
                filters=where,
            )

        docs: list[Document] = []
        for obj in res.objects:
            properties = dict(obj.properties or {})
            text = properties.pop(Field.TEXT_KEY.value, "")

            vec = obj.vector
            if isinstance(vec, dict):
                vec = vec.get("default") or next(iter(vec.values()), None)

            docs.append(Document(page_content=text, vector=vec, metadata=properties))
        return docs

    def _json_serializable(self, value: Any) -> Any:
        """Converts values to JSON-serializable format, handling datetime objects."""
        if isinstance(value, datetime.datetime):
            return value.isoformat()
        return value


class WeaviateVectorFactory(AbstractVectorFactory):
    """Factory class for creating WeaviateVector instances."""

    @override
    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> WeaviateVector:
        """
        Initializes a WeaviateVector instance for the given dataset.

        Uses existing collection name from dataset index structure or generates a new one.
        Updates dataset index structure if not already set.
        """
        if dataset.index_struct_dict:
            class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
            collection_name = class_prefix
        else:
            dataset_id = dataset.id
            collection_name = Dataset.gen_collection_name_by_id(dataset_id)
            dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.WEAVIATE, collection_name))
        return WeaviateVector(
            collection_name=collection_name,
            config=WeaviateConfig(
                endpoint=dify_config.WEAVIATE_ENDPOINT or "",
                grpc_endpoint=dify_config.WEAVIATE_GRPC_ENDPOINT or "",
                api_key=dify_config.WEAVIATE_API_KEY,
                batch_size=dify_config.WEAVIATE_BATCH_SIZE,
            ),
            attributes=attributes,
        )
