import json
import logging
from typing import Any, override
from uuid import UUID, uuid4

from numpy import ndarray
from pgvecto_rs.sqlalchemy import VECTOR  # type: ignore
from pydantic import BaseModel, model_validator
from sqlalchemy import Float, create_engine, insert, select, text
from sqlalchemy import text as sql_text
from sqlalchemy.dialects import postgresql
from sqlalchemy.orm import Mapped, Session, mapped_column, sessionmaker

from configs import dify_config
from core.rag.datasource.vdb.vector_base import BaseVector
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 dify_vdb_pgvecto_rs.collection import CollectionORM
from extensions.ext_redis import redis_client
from models.dataset import Dataset

logger = logging.getLogger(__name__)


class PgvectoRSConfig(BaseModel):
    host: str
    port: int
    user: str
    password: str
    database: str

    @model_validator(mode="before")
    @classmethod
    def validate_config(cls, values: dict[str, Any]):
        if not values["host"]:
            raise ValueError("config PGVECTO_RS_HOST is required")
        if not values["port"]:
            raise ValueError("config PGVECTO_RS_PORT is required")
        if not values["user"]:
            raise ValueError("config PGVECTO_RS_USER is required")
        if not values["password"]:
            raise ValueError("config PGVECTO_RS_PASSWORD is required")
        if not values["database"]:
            raise ValueError("config PGVECTO_RS_DATABASE is required")
        return values


class PGVectoRS(BaseVector):
    def __init__(self, collection_name: str, config: PgvectoRSConfig, dim: int):
        super().__init__(collection_name)
        self._client_config = config
        self._url = (
            f"postgresql+psycopg2://{config.user}:{config.password}@{config.host}:{config.port}/{config.database}"
        )
        self._client = create_engine(self._url)
        with sessionmaker(bind=self._client).begin() as session:
            session.execute(text("CREATE EXTENSION IF NOT EXISTS vectors"))
        self._fields: list[str] = []

        class _Table(CollectionORM):
            __tablename__ = collection_name
            __table_args__ = {"extend_existing": True}
            id: Mapped[UUID] = mapped_column(  # pyrefly: ignore[missing-override-decorator]
                postgresql.UUID(as_uuid=True),
                primary_key=True,
            )
            text: Mapped[str]  # pyrefly: ignore[missing-override-decorator]
            meta: Mapped[dict[str, Any]] = mapped_column(postgresql.JSONB)  # pyrefly: ignore[missing-override-decorator]
            vector: Mapped[ndarray] = mapped_column(VECTOR(dim))  # pyrefly: ignore[missing-override-decorator]

        self._table = _Table
        self._distance_op = "<=>"

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

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

    def create_collection(self, dimension: int):
        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
            index_name = f"{self._collection_name}_embedding_index"
            with sessionmaker(bind=self._client).begin() as session:
                create_statement = sql_text(f"""
                    CREATE TABLE IF NOT EXISTS {self._collection_name} (
                        id UUID PRIMARY KEY,
                        text TEXT NOT NULL,
                        meta JSONB NOT NULL,
                        vector vector({dimension}) NOT NULL
                    ) using heap;
                """)
                session.execute(create_statement)
                index_statement = sql_text(f"""
                        CREATE INDEX IF NOT EXISTS {index_name}
                        ON {self._collection_name} USING vectors(vector vector_l2_ops)
                        WITH (options = $$
                                optimizing.optimizing_threads = 30
                                segment.max_growing_segment_size = 2000
                                segment.max_sealed_segment_size = 30000000
                                [indexing.hnsw]
                                m=30
                                ef_construction=500
                                $$);
                    """)
                session.execute(index_statement)
            redis_client.set(collection_exist_cache_key, 1, ex=3600)

    @override
    def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs):
        pks = []
        with sessionmaker(bind=self._client).begin() as session:
            for document, embedding in zip(documents, embeddings):
                pk = uuid4()
                session.execute(
                    insert(self._table).values(
                        id=pk,
                        text=document.page_content,
                        meta=document.metadata,
                        vector=embedding,
                    ),
                )
                pks.append(pk)

        return pks

    @override
    def get_ids_by_metadata_field(self, key: str, value: str):
        result = None
        with Session(self._client) as session:
            select_statement = sql_text(f"SELECT id FROM {self._collection_name} WHERE meta->>:key = :value")
            result = session.execute(select_statement, {"key": key, "value": value}).fetchall()
        if result:
            return [item[0] for item in result]
        else:
            return None

    @override
    def delete_by_metadata_field(self, key: str, value: str):
        ids = self.get_ids_by_metadata_field(key, value)
        if ids:
            with sessionmaker(bind=self._client).begin() as session:
                select_statement = sql_text(f"DELETE FROM {self._collection_name} WHERE id = ANY(:ids)")
                session.execute(select_statement, {"ids": ids})

    @override
    def delete_by_ids(self, ids: list[str]):
        with Session(self._client) as session:
            select_statement = sql_text(
                f"SELECT id FROM {self._collection_name} WHERE meta->>'doc_id' = ANY (:doc_ids); "
            )
            result = session.execute(select_statement, {"doc_ids": ids}).fetchall()
        if result:
            ids = [item[0] for item in result]
            if ids:
                with sessionmaker(bind=self._client).begin() as session:
                    select_statement = sql_text(f"DELETE FROM {self._collection_name} WHERE id = ANY(:ids)")
                    session.execute(select_statement, {"ids": ids})

    @override
    def delete(self):
        with sessionmaker(bind=self._client).begin() as session:
            session.execute(sql_text(f"DROP TABLE IF EXISTS {self._collection_name}"))

    @override
    def text_exists(self, id: str) -> bool:
        with Session(self._client) as session:
            select_statement = sql_text(
                f"SELECT id FROM {self._collection_name} WHERE meta->>'doc_id' = :doc_id limit 1"
            )
            result = session.execute(select_statement, {"doc_id": id}).fetchall()
        return len(result) > 0

    @override
    def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
        with Session(self._client) as session:
            stmt = (
                select(
                    self._table,
                    self._table.vector.op(self._distance_op, return_type=Float)(
                        query_vector,
                    ).label("distance"),
                )
                .limit(kwargs.get("top_k", 4))
                .order_by("distance")
            )
            document_ids_filter = kwargs.get("document_ids_filter")
            if document_ids_filter:
                stmt = stmt.where(self._table.meta["document_id"].in_(document_ids_filter))
            res = session.execute(stmt)
            results = [(row[0], row[1]) for row in res]

        # Organize results.
        docs = []
        for record, dis in results:
            metadata = record.meta
            score = 1 - dis
            metadata["score"] = score
            score_threshold = float(kwargs.get("score_threshold") or 0.0)
            if score >= score_threshold:
                doc = Document(page_content=record.text, metadata=metadata)
                docs.append(doc)
        return docs

    @override
    def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
        return []


class PGVectoRSFactory(AbstractVectorFactory):
    @override
    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> PGVectoRS:
        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.PGVECTO_RS, collection_name))
        dim = len(embeddings.embed_query("pgvecto_rs"))

        return PGVectoRS(
            collection_name=collection_name,
            config=PgvectoRSConfig(
                host=dify_config.PGVECTO_RS_HOST or "localhost",
                port=dify_config.PGVECTO_RS_PORT or 5432,
                user=dify_config.PGVECTO_RS_USER or "postgres",
                password=dify_config.PGVECTO_RS_PASSWORD or "",
                database=dify_config.PGVECTO_RS_DATABASE or "postgres",
            ),
            dim=dim,
        )
