"""InterSystems IRIS vector database implementation for Dify.

This module provides vector storage and retrieval using IRIS native VECTOR type
with HNSW indexing for efficient similarity search.
"""

from __future__ import annotations

import json
import logging
import threading
import uuid
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any

from configs import dify_config
from configs.middleware.vdb.iris_config import IrisVectorConfig
from core.rag.datasource.vdb.field import parse_metadata_json
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 extensions.ext_redis import redis_client
from models.dataset import Dataset

if TYPE_CHECKING:
    import iris
else:
    try:
        import iris
    except ImportError:
        iris = None  # type: ignore[assignment]

logger = logging.getLogger(__name__)

# Singleton connection pool to minimize IRIS license usage
_pool_lock = threading.Lock()
_pool_instance: IrisConnectionPool | None = None


def get_iris_pool(config: IrisVectorConfig) -> IrisConnectionPool:
    """Get or create the global IRIS connection pool (singleton pattern)."""
    global _pool_instance  # pylint: disable=global-statement
    with _pool_lock:
        if _pool_instance is None:
            logger.info("Initializing IRIS connection pool")
            _pool_instance = IrisConnectionPool(config)
        return _pool_instance


class IrisConnectionPool:
    """Thread-safe connection pool for IRIS database."""

    def __init__(self, config: IrisVectorConfig) -> None:
        self.config = config
        self._pool: list[Any] = []
        self._lock = threading.Lock()
        self._min_size = config.IRIS_MIN_CONNECTION
        self._max_size = config.IRIS_MAX_CONNECTION
        self._in_use = 0
        self._schemas_initialized: set[str] = set()  # Cache for initialized schemas
        self._initialize_pool()

    def _initialize_pool(self) -> None:
        for _ in range(self._min_size):
            self._pool.append(self._create_connection())

    def _create_connection(self) -> Any:
        return iris.connect(
            hostname=self.config.IRIS_HOST,
            port=self.config.IRIS_SUPER_SERVER_PORT,
            namespace=self.config.IRIS_DATABASE,
            username=self.config.IRIS_USER,
            password=self.config.IRIS_PASSWORD,
        )

    def get_connection(self) -> Any:
        """Get a connection from pool or create new if available."""
        with self._lock:
            if self._pool:
                conn = self._pool.pop()
                self._in_use += 1
                return conn
            if self._in_use < self._max_size:
                conn = self._create_connection()
                self._in_use += 1
                return conn
            raise RuntimeError("Connection pool exhausted")

    def return_connection(self, conn: Any) -> None:
        """Return connection to pool after validating it."""
        if not conn:
            return

        # Validate connection health
        is_valid = False
        try:
            cursor = conn.cursor()
            cursor.execute("SELECT 1")
            cursor.close()
            is_valid = True
        except (OSError, RuntimeError) as e:
            logger.debug("Connection validation failed: %s", e)
            try:
                conn.close()
            except (OSError, RuntimeError):
                pass

        with self._lock:
            self._pool.append(conn if is_valid else self._create_connection())
            self._in_use -= 1

    def ensure_schema_exists(self, schema: str) -> None:
        """Ensure schema exists in IRIS database.

        This method is idempotent and thread-safe. It uses a memory cache to avoid
        redundant database queries for already-verified schemas.

        Args:
            schema: Schema name to ensure exists

        Raises:
            Exception: If schema creation fails
        """
        # Fast path: check cache first (no lock needed for read-only set lookup)
        if schema in self._schemas_initialized:
            return

        # Slow path: acquire lock and check again (double-checked locking)
        with self._lock:
            if schema in self._schemas_initialized:
                return

            # Get a connection to check/create schema
            conn = self._pool[0] if self._pool else self._create_connection()
            cursor = conn.cursor()
            try:
                # Check if schema exists using INFORMATION_SCHEMA
                check_sql = """
                    SELECT COUNT(*) FROM INFORMATION_SCHEMA.SCHEMATA
                    WHERE SCHEMA_NAME = ?
                """
                cursor.execute(check_sql, (schema,))  # Must be tuple or list
                exists = cursor.fetchone()[0] > 0

                if not exists:
                    # Schema doesn't exist, create it
                    cursor.execute(f"CREATE SCHEMA {schema}")
                    conn.commit()
                    logger.info("Created schema: %s", schema)
                else:
                    logger.debug("Schema already exists: %s", schema)

                # Add to cache to skip future checks
                self._schemas_initialized.add(schema)

            except Exception:
                conn.rollback()
                logger.exception("Failed to ensure schema %s exists", schema)
                raise
            finally:
                cursor.close()

    def close_all(self) -> None:
        """Close all connections (application shutdown only)."""
        with self._lock:
            for conn in self._pool:
                try:
                    conn.close()
                except (OSError, RuntimeError):
                    pass
            self._pool.clear()
            self._in_use = 0
            self._schemas_initialized.clear()


class IrisVector(BaseVector):
    """IRIS vector database implementation using native VECTOR type and HNSW indexing."""

    # Fallback score for full-text search when Rank function unavailable or TEXT_INDEX disabled
    _FULL_TEXT_FALLBACK_SCORE = 0.5

    def __init__(self, collection_name: str, config: IrisVectorConfig) -> None:
        super().__init__(collection_name)
        self.config = config
        self.table_name = f"embedding_{collection_name}".upper()
        self.schema = config.IRIS_SCHEMA or "dify"
        self.pool = get_iris_pool(config)

    def get_type(self) -> str:
        return VectorType.IRIS

    @contextmanager
    def _get_cursor(self):
        """Context manager for database cursor with connection pooling."""
        conn = self.pool.get_connection()
        cursor = conn.cursor()
        try:
            yield cursor
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            cursor.close()
            self.pool.return_connection(conn)

    def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs) -> list[str]:
        dimension = len(embeddings[0])
        self._create_collection(dimension)
        return self.add_texts(texts, embeddings)

    def add_texts(self, documents: list[Document], embeddings: list[list[float]], **_kwargs) -> list[str]:
        """Add documents with embeddings to the collection."""
        added_ids = []
        with self._get_cursor() as cursor:
            for i, doc in enumerate(documents):
                doc_id = doc.metadata.get("doc_id", str(uuid.uuid4())) if doc.metadata else str(uuid.uuid4())
                metadata = json.dumps(doc.metadata) if doc.metadata else "{}"
                embedding_str = json.dumps(embeddings[i])

                sql = f"INSERT INTO {self.schema}.{self.table_name} (id, text, meta, embedding) VALUES (?, ?, ?, ?)"
                cursor.execute(sql, (doc_id, doc.page_content, metadata, embedding_str))
                added_ids.append(doc_id)

        return added_ids

    def text_exists(self, id: str) -> bool:  # pylint: disable=redefined-builtin
        try:
            with self._get_cursor() as cursor:
                sql = f"SELECT 1 FROM {self.schema}.{self.table_name} WHERE id = ?"
                cursor.execute(sql, (id,))
                return cursor.fetchone() is not None
        except (OSError, RuntimeError, ValueError):
            return False

    def delete_by_ids(self, ids: list[str]) -> None:
        if not ids:
            return

        with self._get_cursor() as cursor:
            placeholders = ",".join(["?" for _ in ids])
            sql = f"DELETE FROM {self.schema}.{self.table_name} WHERE id IN ({placeholders})"
            cursor.execute(sql, ids)

    def delete_by_metadata_field(self, key: str, value: str) -> None:
        """Delete documents by metadata field (JSON LIKE pattern matching)."""
        with self._get_cursor() as cursor:
            pattern = f'%"{key}": "{value}"%'
            sql = f"DELETE FROM {self.schema}.{self.table_name} WHERE meta LIKE ?"
            cursor.execute(sql, (pattern,))

    def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
        """Search similar documents using VECTOR_COSINE with HNSW index."""
        top_k = kwargs.get("top_k", 4)
        score_threshold = float(kwargs.get("score_threshold") or 0.0)
        embedding_str = json.dumps(query_vector)

        with self._get_cursor() as cursor:
            sql = f"""
                SELECT TOP {top_k} id, text, meta, VECTOR_COSINE(embedding, ?) as score
                FROM {self.schema}.{self.table_name}
                ORDER BY score DESC
            """
            cursor.execute(sql, (embedding_str,))

            docs = []
            for row in cursor.fetchall():
                if len(row) >= 4:
                    text, meta_str, score = row[1], row[2], float(row[3])
                    if score >= score_threshold:
                        metadata = parse_metadata_json(meta_str)
                        metadata["score"] = score
                        docs.append(Document(page_content=text, metadata=metadata))
            return docs

    def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
        """Search documents by full-text using iFind index with BM25 relevance scoring.

        When IRIS_TEXT_INDEX is enabled, this method uses the auto-generated Rank
        function from %iFind.Index.Basic to calculate BM25 relevance scores. The Rank
        function is automatically created with naming: {schema}.{table_name}_{index}Rank

        Args:
            query: Search query string
            **kwargs: Optional parameters including top_k, document_ids_filter

        Returns:
            List of Document objects with relevance scores in metadata["score"]
        """
        top_k = kwargs.get("top_k", 5)
        document_ids_filter = kwargs.get("document_ids_filter")

        with self._get_cursor() as cursor:
            if self.config.IRIS_TEXT_INDEX:
                # Use iFind full-text search with auto-generated Rank function
                text_index_name = f"idx_{self.table_name}_text"
                # IRIS removes underscores from function names
                table_no_underscore = self.table_name.replace("_", "")
                index_no_underscore = text_index_name.replace("_", "")
                rank_function = f"{self.schema}.{table_no_underscore}_{index_no_underscore}Rank"

                # Build WHERE clause with document ID filter if provided
                where_clause = f"WHERE %ID %FIND search_index({text_index_name}, ?)"
                # First param for Rank function, second for FIND
                params = [query, query]

                if document_ids_filter:
                    # Add document ID filter
                    placeholders = ",".join("?" * len(document_ids_filter))
                    where_clause += f" AND JSON_VALUE(meta, '$.document_id') IN ({placeholders})"
                    params.extend(document_ids_filter)

                sql = f"""
                    SELECT TOP {top_k}
                        id,
                        text,
                        meta,
                        {rank_function}(%ID, ?) AS score
                    FROM {self.schema}.{self.table_name}
                    {where_clause}
                    ORDER BY score DESC
                """

                logger.debug(
                    "iFind search: query='%s', index='%s', rank='%s'",
                    query,
                    text_index_name,
                    rank_function,
                )

                try:
                    cursor.execute(sql, params)
                except Exception:  # pylint: disable=broad-exception-caught
                    # Fallback to query without Rank function if it fails
                    logger.warning(
                        "Rank function '%s' failed, using fixed score",
                        rank_function,
                        exc_info=True,
                    )
                    sql_fallback = f"""
                        SELECT TOP {top_k} id, text, meta, {self._FULL_TEXT_FALLBACK_SCORE} AS score
                        FROM {self.schema}.{self.table_name}
                        {where_clause}
                    """
                    # Skip first param (for Rank function)
                    cursor.execute(sql_fallback, params[1:])
            else:
                # Fallback to LIKE search (IRIS_TEXT_INDEX disabled)
                from libs.helper import (  # pylint: disable=import-outside-toplevel
                    escape_like_pattern,
                )

                escaped_query = escape_like_pattern(query)
                query_pattern = f"%{escaped_query}%"

                # Build WHERE clause with document ID filter if provided
                where_clause = "WHERE text LIKE ? ESCAPE '\\\\'"
                params = [query_pattern]

                if document_ids_filter:
                    placeholders = ",".join("?" * len(document_ids_filter))
                    where_clause += f" AND JSON_VALUE(meta, '$.document_id') IN ({placeholders})"
                    params.extend(document_ids_filter)

                sql = f"""
                    SELECT TOP {top_k} id, text, meta, {self._FULL_TEXT_FALLBACK_SCORE} AS score
                    FROM {self.schema}.{self.table_name}
                    {where_clause}
                    ORDER BY LENGTH(text) ASC
                """

                logger.debug(
                    "LIKE fallback (TEXT_INDEX disabled): query='%s'",
                    query_pattern,
                )
                cursor.execute(sql, params)

            docs = []
            for row in cursor.fetchall():
                # Expecting 4 columns: id, text, meta, score
                if len(row) >= 4:
                    text_content = row[1]
                    meta_str = row[2]
                    score_value = row[3]

                    metadata = parse_metadata_json(meta_str)
                    # Add score to metadata for hybrid search compatibility
                    score = float(score_value) if score_value is not None else 0.0
                    metadata["score"] = score

                    docs.append(Document(page_content=text_content, metadata=metadata))

            logger.info(
                "Full-text search completed: query='%s', results=%d/%d",
                query,
                len(docs),
                top_k,
            )

            if not docs:
                logger.warning("Full-text search for '%s' returned no results", query)

            return docs

    def delete(self) -> None:
        """Delete the entire collection (drop table - permanent)."""
        with self._get_cursor() as cursor:
            sql = f"DROP TABLE {self.schema}.{self.table_name}"
            cursor.execute(sql)

    def _create_collection(self, dimension: int) -> None:
        """Create table with VECTOR column and HNSW index.

        Uses Redis lock to prevent concurrent creation attempts across multiple
        API server instances (api, worker, worker_beat).
        """
        cache_key = f"vector_indexing_{self._collection_name}"
        lock_name = f"{cache_key}_lock"

        with redis_client.lock(lock_name, timeout=20):  # pylint: disable=not-context-manager
            if redis_client.get(cache_key):
                return

            # Ensure schema exists (idempotent, cached after first call)
            self.pool.ensure_schema_exists(self.schema)

            with self._get_cursor() as cursor:
                # Create table with VECTOR column
                sql = f"""
                    CREATE TABLE {self.schema}.{self.table_name} (
                        id VARCHAR(255) PRIMARY KEY,
                        text CLOB,
                        meta CLOB,
                        embedding VECTOR(DOUBLE, {dimension})
                    )
                """
                logger.info("Creating table: %s.%s", self.schema, self.table_name)
                cursor.execute(sql)

                # Create HNSW index for vector similarity search
                index_name = f"idx_{self.table_name}_embedding"
                sql_index = (
                    f"CREATE INDEX {index_name} ON {self.schema}.{self.table_name} "
                    "(embedding) AS HNSW(Distance='Cosine')"
                )
                logger.info("Creating HNSW index: %s", index_name)
                cursor.execute(sql_index)
                logger.info("HNSW index created successfully: %s", index_name)

                # Create full-text search index if enabled
                logger.info(
                    "IRIS_TEXT_INDEX config value: %s (type: %s)",
                    self.config.IRIS_TEXT_INDEX,
                    type(self.config.IRIS_TEXT_INDEX),
                )
                if self.config.IRIS_TEXT_INDEX:
                    text_index_name = f"idx_{self.table_name}_text"
                    language = self.config.IRIS_TEXT_INDEX_LANGUAGE
                    # Fixed: Removed extra parentheses and corrected syntax
                    sql_text_index = f"""
                        CREATE INDEX {text_index_name} ON {self.schema}.{self.table_name} (text)
                        AS %iFind.Index.Basic
                        (LANGUAGE = '{language}', LOWER = 1, INDEXOPTION = 0)
                    """
                    logger.info(
                        "Creating text index: %s with language: %s",
                        text_index_name,
                        language,
                    )
                    logger.info("SQL for text index: %s", sql_text_index)
                    cursor.execute(sql_text_index)
                    logger.info("Text index created successfully: %s", text_index_name)
                else:
                    logger.warning("Text index creation skipped - IRIS_TEXT_INDEX is disabled")

            redis_client.set(cache_key, 1, ex=3600)


class IrisVectorFactory(AbstractVectorFactory):
    """Factory for creating IrisVector instances."""

    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> IrisVector:
        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)
            index_struct_dict = self.gen_index_struct_dict(VectorType.IRIS, collection_name)
            dataset.index_struct = json.dumps(index_struct_dict)

        return IrisVector(
            collection_name=collection_name,
            config=IrisVectorConfig(
                IRIS_HOST=dify_config.IRIS_HOST,
                IRIS_SUPER_SERVER_PORT=dify_config.IRIS_SUPER_SERVER_PORT,
                IRIS_USER=dify_config.IRIS_USER,
                IRIS_PASSWORD=dify_config.IRIS_PASSWORD,
                IRIS_DATABASE=dify_config.IRIS_DATABASE,
                IRIS_SCHEMA=dify_config.IRIS_SCHEMA,
                IRIS_CONNECTION_URL=dify_config.IRIS_CONNECTION_URL,
                IRIS_MIN_CONNECTION=dify_config.IRIS_MIN_CONNECTION,
                IRIS_MAX_CONNECTION=dify_config.IRIS_MAX_CONNECTION,
                IRIS_TEXT_INDEX=dify_config.IRIS_TEXT_INDEX,
                IRIS_TEXT_INDEX_LANGUAGE=dify_config.IRIS_TEXT_INDEX_LANGUAGE,
            ),
        )
