import json
import logging
import math
from typing import Any, TypedDict, override

from pydantic import BaseModel
from tcvdb_text.encoder import BM25Encoder  # type: ignore
from tcvectordb import RPCVectorDBClient, VectorDBException  # type: ignore
from tcvectordb.model import document, enum  # type: ignore
from tcvectordb.model import index as vdb_index  # type: ignore
from tcvectordb.model.document import AnnSearch, Filter, KeywordSearch, WeightedRerank  # type: ignore

from configs import dify_config
from core.rag.datasource.vdb.field import parse_metadata_json
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__)


class TencentParamsDict(TypedDict):
    url: str
    username: str | None
    key: str | None
    timeout: float


class TencentConfig(BaseModel):
    url: str
    api_key: str | None = None
    timeout: float = 30
    username: str | None = None
    database: str | None = None
    index_type: str = "HNSW"
    metric_type: str = "IP"
    shard: int = 1
    replicas: int = 2
    max_upsert_batch_size: int = 128
    enable_hybrid_search: bool = False  # Flag to enable hybrid search

    def to_tencent_params(self) -> TencentParamsDict:
        result: TencentParamsDict = {
            "url": self.url,
            "username": self.username,
            "key": self.api_key,
            "timeout": self.timeout,
        }
        return result


bm25 = BM25Encoder.default("zh")


class TencentVector(BaseVector):
    field_id: str = "id"
    field_vector: str = "vector"
    field_text: str = "text"
    field_metadata: str = "metadata"

    def __init__(self, collection_name: str, config: TencentConfig):
        super().__init__(collection_name)
        self._client_config = config
        self._client = RPCVectorDBClient(**self._client_config.to_tencent_params())
        self._enable_hybrid_search = False
        self._dimension = 1024
        self._init_database()
        self._load_collection()

    def _load_collection(self):
        """
        Check if the collection supports hybrid search.
        """
        if self._client_config.enable_hybrid_search:
            self._enable_hybrid_search = True
            if self._has_collection():
                coll = self._client.describe_collection(
                    database_name=self._client_config.database, collection_name=self.collection_name
                )
                has_hybrid_search = False
                for idx in coll.indexes:
                    if idx.name == "sparse_vector":
                        has_hybrid_search = True
                    elif idx.name == "vector":
                        self._dimension = idx.dimension
                if not has_hybrid_search:
                    self._enable_hybrid_search = False

    def _init_database(self):
        return self._client.create_database_if_not_exists(database_name=self._client_config.database)

    @override
    def get_type(self) -> str:
        return VectorType.TENCENT

    def to_index_struct(self) -> VectorIndexStructDict:
        result: VectorIndexStructDict = {
            "type": self.get_type(),
            "vector_store": {"class_prefix": self._collection_name},
        }
        return result

    def _has_collection(self) -> bool:
        return bool(
            self._client.exists_collection(
                database_name=self._client_config.database, collection_name=self.collection_name
            )
        )

    def _create_collection(self, dimension: int):
        self._dimension = dimension
        lock_name = f"vector_indexing_lock_{self._collection_name}"
        with redis_client.lock(lock_name, timeout=20):
            collection_exist_cache_key = f"vector_indexing_{self._collection_name}"
            if redis_client.get(collection_exist_cache_key):
                return

            if self._has_collection():
                return

            index_type = None
            for k, v in enum.IndexType.__members__.items():
                if k == self._client_config.index_type:
                    index_type = v
            if index_type is None:
                raise ValueError("unsupported index_type")
            metric_type = None
            for k, v in enum.MetricType.__members__.items():
                if k == self._client_config.metric_type:
                    metric_type = v
            if metric_type is None:
                raise ValueError("unsupported metric_type")
            params = vdb_index.HNSWParams(m=16, efconstruction=200)
            index_id = vdb_index.FilterIndex(self.field_id, enum.FieldType.String, enum.IndexType.PRIMARY_KEY)
            index_vector = vdb_index.VectorIndex(
                self.field_vector,
                dimension,
                index_type,
                metric_type,
                params,
            )
            index_metadate = vdb_index.FilterIndex(self.field_metadata, enum.FieldType.Json, enum.IndexType.FILTER)
            index_sparse_vector = vdb_index.SparseIndex(
                name="sparse_vector",
                field_type=enum.FieldType.SparseVector,
                index_type=enum.IndexType.SPARSE_INVERTED,
                metric_type=enum.MetricType.IP,
            )
            indexes = [index_id, index_vector, index_metadate]
            if self._enable_hybrid_search:
                indexes.append(index_sparse_vector)
            try:
                self._client.create_collection(
                    database_name=self._client_config.database,
                    collection_name=self._collection_name,
                    shard=self._client_config.shard,
                    replicas=self._client_config.replicas,
                    description="Collection for Dify",
                    indexes=indexes,
                )
            except VectorDBException as e:
                if "fieldType:json" not in e.message:
                    raise e
                # vdb version not support json, use string
                index_metadate = vdb_index.FilterIndex(
                    self.field_metadata, enum.FieldType.String, enum.IndexType.FILTER
                )
                indexes = [index_id, index_vector, index_metadate]
                if self._enable_hybrid_search:
                    indexes.append(index_sparse_vector)
                self._client.create_collection(
                    database_name=self._client_config.database,
                    collection_name=self._collection_name,
                    shard=self._client_config.shard,
                    replicas=self._client_config.replicas,
                    description="Collection for Dify",
                    indexes=indexes,
                )
            redis_client.set(collection_exist_cache_key, 1, ex=3600)

    @override
    def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs):
        self._create_collection(len(embeddings[0]))
        self.add_texts(texts, embeddings)

    @override
    def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs):
        texts = [doc.page_content for doc in documents]
        metadatas = [doc.metadata for doc in documents]
        total_count = len(embeddings)
        batch_size = self._client_config.max_upsert_batch_size
        batch = math.ceil(total_count / batch_size)
        for j in range(batch):
            docs = []
            start_idx = j * batch_size
            end_idx = min(total_count, (j + 1) * batch_size)
            for i in range(start_idx, end_idx):
                if metadatas is None:
                    continue
                metadata = metadatas[i] or {}
                doc = document.Document(
                    id=metadata.get("doc_id"),
                    vector=embeddings[i],
                    text=texts[i],
                    metadata=metadata,
                )
                if self._enable_hybrid_search:
                    doc.__dict__["sparse_vector"] = bm25.encode_texts(texts[i])
                docs.append(doc)
            self._client.upsert(
                database_name=self._client_config.database,
                collection_name=self.collection_name,
                documents=docs,
                timeout=self._client_config.timeout,
            )

    @override
    def text_exists(self, id: str) -> bool:
        docs = self._client.query(
            database_name=self._client_config.database, collection_name=self.collection_name, document_ids=[id]
        )
        if docs and len(docs) > 0:
            return True
        return False

    @override
    def delete_by_ids(self, ids: list[str]):
        if not ids:
            return

        total_count = len(ids)
        batch_size = self._client_config.max_upsert_batch_size
        batch = math.ceil(total_count / batch_size)

        for j in range(batch):
            start_idx = j * batch_size
            end_idx = min(total_count, (j + 1) * batch_size)
            batch_ids = ids[start_idx:end_idx]

            self._client.delete(
                database_name=self._client_config.database, collection_name=self.collection_name, document_ids=batch_ids
            )

    @override
    def delete_by_metadata_field(self, key: str, value: str):
        self._client.delete(
            database_name=self._client_config.database,
            collection_name=self.collection_name,
            filter=Filter(Filter.In(f"metadata.{key}", [value])),
        )

    @override
    def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
        document_ids_filter = kwargs.get("document_ids_filter")
        filter = None
        if document_ids_filter:
            filter = Filter(Filter.In("metadata.document_id", document_ids_filter))
        res = self._client.search(
            database_name=self._client_config.database,
            collection_name=self.collection_name,
            vectors=[query_vector],
            filter=filter,
            params=document.HNSWSearchParams(ef=kwargs.get("ef", 10)),
            retrieve_vector=False,
            limit=kwargs.get("top_k", 4),
            timeout=self._client_config.timeout,
        )
        score_threshold = float(kwargs.get("score_threshold") or 0.0)
        return self._get_search_res(res, score_threshold)

    @override
    def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
        document_ids_filter = kwargs.get("document_ids_filter")
        filter = None
        if document_ids_filter:
            filter = Filter(Filter.In("metadata.document_id", document_ids_filter))
        if not self._enable_hybrid_search:
            return []
        res = self._client.hybrid_search(
            database_name=self._client_config.database,
            collection_name=self.collection_name,
            ann=[
                AnnSearch(
                    field_name="vector",
                    data=[0.0] * self._dimension,
                )
            ],
            match=[
                KeywordSearch(
                    field_name="sparse_vector",
                    data=bm25.encode_queries(query),
                ),
            ],
            rerank=WeightedRerank(
                field_list=["vector", "sparse_vector"],
                weight=[0, 1],
            ),
            retrieve_vector=False,
            limit=kwargs.get("top_k", 4),
            filter=filter,
        )
        score_threshold = float(kwargs.get("score_threshold") or 0.0)
        return self._get_search_res(res, score_threshold)

    def _get_search_res(self, res: list | None, score_threshold: float) -> list[Document]:
        docs: list[Document] = []
        if res is None or len(res) == 0:
            return docs

        for result in res[0]:
            raw_meta = result.get(self.field_metadata)
            # Compatible with version 1.1.3 and below: str means old driver.
            score = (1 - result.get("score", 0.0)) if isinstance(raw_meta, str) else result.get("score", 0.0)
            meta = parse_metadata_json(raw_meta)
            if score >= score_threshold:
                meta["score"] = score
                doc = Document(page_content=result.get(self.field_text), metadata=meta)
                docs.append(doc)
        return docs

    @override
    def delete(self):
        if self._has_collection():
            self._client.drop_collection(
                database_name=self._client_config.database, collection_name=self.collection_name
            )


class TencentVectorFactory(AbstractVectorFactory):
    @override
    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TencentVector:
        if dataset.index_struct_dict:
            class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
            collection_name = class_prefix.lower()
        else:
            dataset_id = dataset.id
            collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower()
            dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.TENCENT, collection_name))

        return TencentVector(
            collection_name=collection_name,
            config=TencentConfig(
                url=dify_config.TENCENT_VECTOR_DB_URL or "",
                api_key=dify_config.TENCENT_VECTOR_DB_API_KEY,
                timeout=dify_config.TENCENT_VECTOR_DB_TIMEOUT,
                username=dify_config.TENCENT_VECTOR_DB_USERNAME,
                database=dify_config.TENCENT_VECTOR_DB_DATABASE,
                shard=dify_config.TENCENT_VECTOR_DB_SHARD,
                replicas=dify_config.TENCENT_VECTOR_DB_REPLICAS,
                enable_hybrid_search=dify_config.TENCENT_VECTOR_DB_ENABLE_HYBRID_SEARCH or False,
            ),
        )
