"""
OceanBase storage implementation

This module provides OceanBase-based storage for memory data.
"""
import heapq
import json
import logging
import uuid
from typing import Any, Dict, List, Optional
from concurrent.futures import ThreadPoolExecutor
from powermem.storage.base import VectorStoreBase, OutputData
from powermem.utils.utils import serialize_datetime, generate_snowflake_id

try:
    from pyobvector import (
        VECTOR,
        ObVecClient,
        cosine_distance,
        inner_product,
        l2_distance,
        VecIndexType,
    )
    from pyobvector.schema import ReplaceStmt
    from sqlalchemy import JSON, Column, String, Table, func, ColumnElement, BigInteger
    from sqlalchemy import text, and_, or_, not_, select, bindparam, literal_column
    from sqlalchemy.dialects.mysql import LONGTEXT
except ImportError as e:
    raise ImportError(
        f"Required dependencies not found: {e}. Please install pyobvector and sqlalchemy."
    )

from powermem.storage.oceanbase import constants

logger = logging.getLogger(__name__)

class OceanBaseVectorStore(VectorStoreBase):
    """OceanBase vector store implementation"""

    def __init__(
            self,
            collection_name: str,
            connection_args: Optional[Dict[str, Any]] = None,
            vidx_metric_type: str = constants.DEFAULT_OCEANBASE_VECTOR_METRIC_TYPE,
            vidx_algo_params: Optional[Dict] = None,
            index_type: str = constants.DEFAULT_INDEX_TYPE,
            embedding_model_dims: Optional[int] = None,
            primary_field: str = constants.DEFAULT_PRIMARY_FIELD,
            vector_field: str = constants.DEFAULT_VECTOR_FIELD,
            text_field: str = constants.DEFAULT_TEXT_FIELD,
            metadata_field: str = constants.DEFAULT_METADATA_FIELD,
            vidx_name: str = constants.DEFAULT_VIDX_NAME,
            normalize: bool = False,
            include_sparse: bool = False,
            auto_configure_vector_index: bool = True,
            # Connection parameters (for compatibility with config)
            host: Optional[str] = None,
            port: Optional[str] = None,
            user: Optional[str] = None,
            password: Optional[str] = None,
            db_name: Optional[str] = None,
            hybrid_search: bool = True,
            fulltext_parser: str = constants.DEFAULT_FULLTEXT_PARSER,
            vector_weight: float = 0.5,
            fts_weight: float = 0.5,
            reranker: Optional[Any] = None,
            **kwargs,
    ):
        """
        Initialize the OceanBase vector store.

        Args:
            collection_name (str): Name of the collection/table.
            connection_args (Optional[Dict[str, Any]]): Connection parameters for OceanBase.
            vidx_metric_type (str): Metric method of distance between vectors.
            vidx_algo_params (Optional[Dict]): Index parameters.
            index_type (str): Type of vector index to use.
            embedding_model_dims (Optional[int]): Dimension of vectors.
            primary_field (str): Name of the primary key column.
            vector_field (str): Name of the vector column.
            text_field (str): Name of the text column.
            metadata_field (str): Name of the metadata column.
            vidx_name (str): Name of the vector index.
            normalize (bool): Whether to perform L2 normalization on vectors.
            include_sparse (bool): Whether to include sparse vector support.
            auto_configure_vector_index (bool): Whether to automatically configure vector index settings.
            host (Optional[str]): OceanBase server host.
            port (Optional[str]): OceanBase server port.
            user (Optional[str]): OceanBase username.
            password (Optional[str]): OceanBase password.
            db_name (Optional[str]): OceanBase database name.
            hybrid_search (bool): Whether to use hybrid search.
            vector_weight (float): Weight for vector search in hybrid search (default: 1.0).
            fts_weight (float): Weight for full-text search in hybrid search (default: 1.0).
        """
        self.normalize = normalize
        self.include_sparse = include_sparse
        self.auto_configure_vector_index = auto_configure_vector_index
        self.hybrid_search = hybrid_search
        self.fulltext_parser = fulltext_parser
        self.vector_weight = vector_weight
        self.fts_weight = fts_weight
        self.reranker = reranker

        # Validate fulltext parser
        if self.fulltext_parser not in constants.OCEANBASE_SUPPORTED_FULLTEXT_PARSERS:
            supported = ', '.join(constants.OCEANBASE_SUPPORTED_FULLTEXT_PARSERS)
            raise ValueError(
                f"Invalid fulltext parser: {self.fulltext_parser}. "
                f"Supported parsers are: {supported}"
            )

        # Handle connection arguments - prioritize individual parameters over connection_args
        if connection_args is None:
            connection_args = {}

        # Merge individual connection parameters with connection_args
        final_connection_args = {
            "host": host or connection_args.get("host", constants.DEFAULT_OCEANBASE_CONNECTION["host"]),
            "port": port or connection_args.get("port", constants.DEFAULT_OCEANBASE_CONNECTION["port"]),
            "user": user or connection_args.get("user", constants.DEFAULT_OCEANBASE_CONNECTION["user"]),
            "password": password or connection_args.get("password", constants.DEFAULT_OCEANBASE_CONNECTION["password"]),
            "db_name": db_name or connection_args.get("db_name", constants.DEFAULT_OCEANBASE_CONNECTION["db_name"]),
        }

        self.connection_args = final_connection_args

        self.index_type = index_type.upper()
        if self.index_type not in constants.OCEANBASE_SUPPORTED_VECTOR_INDEX_TYPES:
            raise ValueError(
                f"`index_type` should be one of "
                f"{list(constants.OCEANBASE_SUPPORTED_VECTOR_INDEX_TYPES.keys())}. "
                f"Got {self.index_type}"
            )

        # Set default parameters based on index type
        if vidx_algo_params is None:
            index_param_map = constants.OCEANBASE_BUILD_PARAMS_MAPPING
            self.vidx_algo_params = index_param_map[self.index_type].copy()

            if self.index_type == "IVF_PQ" and "m" not in self.vidx_algo_params:
                self.vidx_algo_params["m"] = 3
        else:
            self.vidx_algo_params = vidx_algo_params.copy()

        # Set field names
        self.collection_name = collection_name
        self.embedding_model_dims = int(embedding_model_dims) if embedding_model_dims is not None else None
        self.primary_field = primary_field
        self.vector_field = vector_field
        self.text_field = text_field
        self.metadata_field = metadata_field
        self.vidx_name = vidx_name
        self.sparse_vector_field = "sparse_embedding"
        self.fulltext_field = "fulltext_content"

        # Set up vector index parameters
        self.vidx_metric_type = vidx_metric_type.lower()

        # Initialize client
        self._create_client(**kwargs)
        assert self.obvector is not None

        # Autoconfigure vector index settings if enabled
        if self.auto_configure_vector_index:
            self._configure_vector_index_settings()

        self._create_col()

    def _create_client(self, **kwargs):
        """Create and initialize the OceanBase vector client."""
        host = self.connection_args.get("host")
        port = self.connection_args.get("port")
        user = self.connection_args.get("user")
        password = self.connection_args.get("password")
        db_name = self.connection_args.get("db_name")

        self.obvector = ObVecClient(
            uri=f"{host}:{port}",
            user=user,
            password=password,
            db_name=db_name,
            **kwargs,
        )

    def _configure_vector_index_settings(self):
        """Configure OceanBase vector index settings automatically."""
        try:
            logger.info("Configuring OceanBase vector index settings...")

            # Set vector memory limit percentage
            with self.obvector.engine.connect() as conn:
                conn.execute(text("ALTER SYSTEM SET ob_vector_memory_limit_percentage = 30"))
                conn.commit()
                logger.info("Set ob_vector_memory_limit_percentage = 30")

            logger.info("OceanBase vector index configuration completed")

        except Exception as e:
            logger.warning(f"Failed to configure vector index settings: {e}")
            logger.warning("   Vector index functionality may not work properly")

    def _create_table_with_index_by_embedding_model_dims(self) -> None:
        """Create table with vector index based on embedding dimension."""
        cols = [
            # Primary key - Snowflake ID (BIGINT without AUTO_INCREMENT)
            Column(self.primary_field, BigInteger, primary_key=True, autoincrement=False),
            # Vector field
            Column(self.vector_field, VECTOR(self.embedding_model_dims)),
            # Text content field
            Column(self.text_field, LONGTEXT),
            # Metadata field (JSON)
            Column(self.metadata_field, JSON),
            Column("user_id", String(128)),  # User identifier
            Column("agent_id", String(128)),  # Agent identifier
            Column("run_id", String(128)),  # Run identifier
            Column("actor_id", String(128)),  # Actor identifier
            Column("hash", String(32)),  # MD5 hash (32 chars)
            Column("created_at", String(128)),
            Column("updated_at", String(128)),
            Column("category", String(64)),  # Category name
            Column(self.fulltext_field, LONGTEXT)
        ]

        # Add hybrid search columns if enabled
        if self.include_sparse:
            cols.append(Column(self.sparse_vector_field, JSON))

        # Create vector index
        vidx_params = self.obvector.prepare_index_params()
        vidx_params.add_index(
            field_name=self.vector_field,
            index_type=constants.OCEANBASE_SUPPORTED_VECTOR_INDEX_TYPES[self.index_type],
            index_name=self.vidx_name,
            metric_type=self.vidx_metric_type,
            params=self.vidx_algo_params,
        )

        # Add sparse vector index if enabled
        if self.include_sparse:
            logger.warning("Sparse vector indexing not fully implemented yet")

        # Create table with vector index first
        self.obvector.create_table_with_index_params(
            table_name=self.collection_name,
            columns=cols,
            indexes=None,
            vidxs=vidx_params,
            partitions=None,
        )

        logger.debug("DEBUG: Table '%s' created successfully", self.collection_name)

    def _normalize(self, vector: List[float]) -> List[float]:
        """Normalize vector using L2 normalization."""
        import numpy as np
        arr = np.array(vector)
        norm = np.linalg.norm(arr)
        if norm == 0:
            return vector
        arr = arr / norm
        return arr.tolist()

    def _get_distance_function(self, metric_type: str):
        """Get the appropriate distance function for the given metric type."""
        if metric_type == "inner_product":
            return inner_product
        elif metric_type == "l2":
            return l2_distance
        elif metric_type == "cosine":
            return cosine_distance
        else:
            raise ValueError(f"Unsupported metric type: {metric_type}")

    def _get_default_search_params(self) -> dict:
        """Get default search parameters based on index type."""
        search_param_map = constants.OCEANBASE_SEARCH_PARAMS_MAPPING
        return search_param_map.get(
            self.index_type, constants.DEFAULT_OCEANBASE_HNSW_SEARCH_PARAM
        )

    def create_col(self, name: str, vector_size: Optional[int] = None, distance: str = "l2"):
        """Create a new collection."""
        try:
            if vector_size is None:
                raise ValueError("vector_size must be specified to create a collection.")
            distance = distance.lower()
            if distance not in ("l2", "inner_product", "cosine"):
                raise ValueError("distance must be one of 'l2', 'inner_product', or 'cosine'.")
            self.embedding_model_dims = int(vector_size)
            self.vidx_metric_type = distance
            self.collection_name = name

            self._create_col()
            logger.info(f"Successfully created collection '{name}' with vector size {vector_size} and distance '{distance}'")
            
        except ValueError as e:
            logger.error(f"Invalid parameters for creating collection: {e}")
            raise
        except Exception as e:
            logger.error(f"Failed to create collection '{name}': {e}", exc_info=True)
            raise

    def _create_col(self):
        """Create a new collection."""

        if self.embedding_model_dims is None:
            raise ValueError(
                "embedding_model_dims is required for OceanBase vector operations. "
                "Please configure embedding_model_dims in your OceanBaseConfig."
            )

        # Set up vector index parameters
        if self.vidx_metric_type not in ("l2", "inner_product", "cosine"):
            raise ValueError(
                "`vidx_metric_type` should be set in `l2`/`inner_product`/`cosine`."
            )

        # Only create table if it doesn't exist (preserve existing data)
        if not self.obvector.check_table_exists(self.collection_name):
            self._create_table_with_index_by_embedding_model_dims()
            logger.info(f"Created new table {self.collection_name}")
        else:
            logger.info(f"Table {self.collection_name} already exists, preserving existing data")
            # Check if the existing table's vector dimension matches the requested dimension
            existing_dim = self._get_existing_vector_dimension()
            if existing_dim is not None and existing_dim != self.embedding_model_dims:
                raise ValueError(
                    f"Vector dimension mismatch: existing table '{self.collection_name}' has "
                    f"vector dimension {existing_dim}, but requested dimension is {self.embedding_model_dims}. "
                    f"Please use a different collection name or delete the existing table."
                )

        if self.hybrid_search:
            self._check_and_create_fulltext_index()
        self.table = Table(self.collection_name, self.obvector.metadata_obj, autoload_with=self.obvector.engine)

    def insert(self,
               vectors: List[List[float]],
               payloads: Optional[List[Dict]] = None,
               ids: Optional[List[str]] = None) -> List[int]:
        """
        Insert vectors into the collection.
        
        Args:
            vectors: List of vectors to insert
            payloads: Optional list of payload dictionaries
            ids: Deprecated parameter (ignored), IDs are now generated using Snowflake algorithm
            
        Returns:
            List[int]: List of generated Snowflake IDs
        """
        try:
            if not vectors:
                return []

            if payloads is None:
                payloads = [{} for _ in vectors]

            # Generate Snowflake IDs for each vector
            generated_ids = [generate_snowflake_id() for _ in range(len(vectors))]

            # Prepare data for insertion with explicit IDs
            data: List[Dict[str, Any]] = []
            for vector, payload, vector_id in zip(vectors, payloads, generated_ids):
                record = self._build_record_for_insert(vector, payload)
                # Explicitly set the primary key field with Snowflake ID
                record[self.primary_field] = vector_id
                data.append(record)

            # Use transaction to ensure atomicity of insert
            table = Table(self.collection_name, self.obvector.metadata_obj, 
                         autoload_with=self.obvector.engine)
            
            with self.obvector.engine.connect() as conn:
                with conn.begin():
                    # Execute REPLACE INTO (upsert) statement
                    upsert_stmt = ReplaceStmt(table).values(data)
                    conn.execute(upsert_stmt)
            
            logger.debug(f"Successfully inserted {len(vectors)} vectors, generated Snowflake IDs: {generated_ids}")
            return generated_ids
            
        except Exception as e:
            logger.error(f"Failed to insert vectors into collection '{self.collection_name}': {e}", exc_info=True)
            raise

    def _parse_metadata(self, metadata_json):
        """
        Parse metadata from OceanBase.
        
        SQLAlchemy's JSON type automatically deserializes to dict, but this method
        handles backward compatibility with legacy string-serialized data.
        """
        if isinstance(metadata_json, dict):
            # SQLAlchemy JSON type returns dict directly (preferred path)
            return metadata_json
        elif isinstance(metadata_json, str):
            # Legacy compatibility: handle manually serialized strings
            try:
                # First attempt to parse
                metadata = json.loads(metadata_json)
                # Check if it's still a string (double encoded - legacy bug)
                if isinstance(metadata, str):
                    try:
                        # Second attempt to parse
                        metadata = json.loads(metadata)
                    except json.JSONDecodeError:
                        metadata = {}
                return metadata
            except json.JSONDecodeError:
                return {}
        else:
            return {}

    def _generate_where_clause(self, filters: Optional[Dict] = None) -> Optional[List]:
        """
        Generate a properly formatted where clause for OceanBase.

        Args:
            filters (Optional[Dict]): The filter conditions.
                Supports both simple and complex formats:

                Simple format (Open Source):
                - Simple values: {"field": "value"} -> field = 'value'
                - Comparison ops: {"field": {"gte": 10, "lte": 20}}
                - List values: {"field": ["a", "b", "c"]} -> field IN ('a', 'b', 'c')

                Complex format (Platform):
                - AND logic: {"AND": [{"user_id": "alice"}, {"category": "food"}]}
                - OR logic: {"OR": [{"rating": {"gte": 4.0}}, {"priority": "high"}]}
                - Nested: {"AND": [{"user_id": "alice"}, {"OR": [{"rating": {"gte": 4.0}}, {"priority": "high"}]}]}

        Returns:
            Optional[List]: List of SQLAlchemy ColumnElement objects for where clause.
        """

        def get_column(key) -> ColumnElement:
            """Get the appropriate column element for a field."""
            if key in self.table.c:
                return self.table.c[key]
            else:
                # Use ->> operator for unquoted JSON extract (MySQL/PostgreSQL)
                return self.table.c[self.metadata_field].op("->>")(f"$.{key}")

        def build_condition(key, value):
            """Build a single condition."""
            column = get_column(key)

            if isinstance(value, dict):
                # Handle comparison operators
                conditions = []
                for op, op_value in value.items():
                    op = op.lstrip("$")
                    match op:
                        case "eq":
                            conditions.append(column == op_value)
                        case "ne":
                            conditions.append(column != op_value)
                        case "gt":
                            conditions.append(column > op_value)
                        case "gte":
                            conditions.append(column >= op_value)
                        case "lt":
                            conditions.append(column < op_value)
                        case "lte":
                            conditions.append(column <= op_value)
                        case "in":
                            if not isinstance(op_value, list):
                                raise TypeError(f"Value for $in must be a list, got {type(op_value)}")
                            conditions.append(column.in_(op_value))
                        case "nin":
                            if not isinstance(op_value, list):
                                raise TypeError(f"Value for $nin must be a list, got {type(op_value)}")
                            conditions.append(~column.in_(op_value))
                        case "like":
                            conditions.append(column.like(str(op_value)))
                        case "ilike":
                            conditions.append(column.ilike(str(op_value)))
                        case _:
                            raise ValueError(f"Unsupported operator: {op}")
                return and_(*conditions) if conditions else None
            elif value is None:
                return column.is_(None)
            else:
                return column == value

        def process_condition(cond):
            """Process a single condition, handling nested AND/OR logic."""
            if isinstance(cond, dict):
                # Handle complex filters with AND/OR
                if "AND" in cond:
                    and_conditions = [process_condition(item) for item in cond["AND"]]
                    and_conditions = [c for c in and_conditions if c is not None]
                    return and_(*and_conditions) if and_conditions else None
                elif "OR" in cond:
                    or_conditions = [process_condition(item) for item in cond["OR"]]
                    or_conditions = [c for c in or_conditions if c is not None]
                    return or_(*or_conditions) if or_conditions else None
                else:
                    # Simple key-value filters
                    conditions = []
                    for k, v in cond.items():
                        expr = build_condition(k, v)
                        if expr is not None:
                            conditions.append(expr)
                    return and_(*conditions) if conditions else None
            elif isinstance(cond, list):
                subconditions = [process_condition(c) for c in cond]
                subconditions = [c for c in subconditions if c is not None]
                return and_(*subconditions) if subconditions else None
            else:
                return None

        # Handle complex filters with AND/OR
        result = process_condition(filters)
        return [result] if result is not None else None

    def _parse_row(self, row) -> tuple:
        """Parse a database result row. Returns up to 12 fields, padding with None if needed."""
        padded_row = list(row) + [None] * (12 - len(row))
        return tuple(padded_row[:12])

    def _build_standard_metadata(self, user_id: str, agent_id: str, run_id: str,
                                 actor_id: str, hash_val: str, created_at: str,
                                 updated_at: str, category: str, metadata_json: str) -> Dict:
        """Build standard metadata dictionary from row fields."""
        # Parse the JSON metadata first - this contains user-defined metadata
        user_metadata = self._parse_metadata(metadata_json)
        
        # Build complete payload with standard fields at top level and user metadata nested
        metadata = {
            "user_id": user_id,
            "agent_id": agent_id,
            "run_id": run_id,
            "actor_id": actor_id,
            "hash": hash_val,
            "created_at": created_at,
            "updated_at": updated_at,
            "category": category,
            # Store user metadata as nested structure to preserve it
            "metadata": user_metadata
        }

        return metadata

    def _create_output_data(self, vector_id: int, text_content: str, score: float,
                            metadata: Dict) -> OutputData:
        """Create an OutputData object with standard structure."""
        return OutputData(
            id=vector_id,
            score=score,
            payload={
                "data": text_content,
                **metadata
            }
        )

    def _build_record_for_insert(self, vector: List[float], payload: Dict) -> Dict[str, Any]:
        """
        Build a record dictionary for insertion with all standard fields.
        Note: Primary key (id) should be set explicitly before insertion.
        """
        # Serialize metadata to handle datetime objects
        metadata = payload.get("metadata", {})
        serialized_metadata = serialize_datetime(metadata) if metadata else {}
        
        record = {
            # Primary key (id) will be set explicitly in insert() method with Snowflake ID
            self.vector_field: (
                vector if not self.normalize else self._normalize(vector)
            ),
            self.text_field: payload.get("data", ""),
            self.metadata_field: serialized_metadata,
            "user_id": payload.get("user_id", ""),
            "agent_id": payload.get("agent_id", ""),
            "run_id": payload.get("run_id", ""),
            "actor_id": payload.get("actor_id", ""),
            "hash": payload.get("hash", ""),
            "created_at": serialize_datetime(payload.get("created_at", "")),
            "updated_at": serialize_datetime(payload.get("updated_at", "")),
            "category": payload.get("category", ""),
        }

        # Add hybrid search fields if enabled
        if self.include_sparse and "sparse_embedding" in payload:
            record[self.sparse_vector_field] = payload["sparse_embedding"]  # SQLAlchemy JSON type handles serialization automatically

        # Always add full-text content (enabled by default)
        fulltext_content = payload.get("fulltext_content") or payload.get("data", "")
        record[self.fulltext_field] = fulltext_content

        return record

    def search(self,
               query: str,
               vectors: List[List[float]],
               limit: int = 5,
               filters: Optional[Dict] = None) -> list[OutputData]:
        # Check if hybrid search is enabled, and we have query text
        # Full-text search is always enabled by default
        if self.hybrid_search and query:
            return self._hybrid_search(query, vectors, limit, filters)
        else:
            return self._vector_search(query, vectors, limit, filters)

    def _vector_search(self,
                       query: str,
                       vectors: List[List[float]],
                       limit: int = 5,
                       filters: Optional[Dict] = None) -> list[OutputData]:
        """Perform pure vector search."""
        try:
            # Handle both cases: single vector or list of vectors
            # If vectors is a single vector (list of floats), use it directly
            if isinstance(vectors, list) and len(vectors) > 0 and isinstance(vectors[0], (int, float)):
                query_vector = vectors
            # If vectors is a list of vectors, use the first one
            elif isinstance(vectors, list) and len(vectors) > 0 and isinstance(vectors[0], list):
                query_vector = vectors[0]
            else:
                logger.warning("Invalid vector format provided for search")
                return []

            # Build where clause from filters
            where_clause = self._generate_where_clause(filters)

            # Perform vector search - pyobvector expects a single vector, not a list of vectors
            results = self.obvector.ann_search(
                table_name=self.collection_name,
                vec_data=query_vector if not self.normalize else self._normalize(query_vector),
                vec_column_name=self.vector_field,
                distance_func=self._get_distance_function(self.vidx_metric_type),
                with_dist=True,
                topk=limit,
                output_column_names=[
                    self.text_field,
                    self.metadata_field,
                    self.primary_field,
                    "user_id",
                    "agent_id",
                    "run_id",
                    "actor_id",
                    "hash",
                    "created_at",
                    "updated_at",
                    "category",
                ],
                where_clause=where_clause,
            )

            # Convert results to OutputData objects
            search_results = []
            for row in results.fetchall():
                (text_content, metadata_json, vector_id, user_id, agent_id, run_id,
                 actor_id, hash_val, created_at, updated_at, category, distance) = self._parse_row(row)

                # Build standard metadata
                metadata = self._build_standard_metadata(
                    user_id, agent_id, run_id, actor_id, hash_val,
                    created_at, updated_at, category, metadata_json
                )

                # Convert distance to score based on metric type
                # Handle None distance (shouldn't happen but be defensive)
                if distance is None:
                    logger.warning(f"Distance is None for vector_id {vector_id}, using default score 0.0")
                    score = 0.0
                elif self.vidx_metric_type == "l2":
                    # For L2 distance, lower is better, so we can use 1/(1+distance) or just use distance
                    score = float(distance)
                elif self.vidx_metric_type == "cosine":
                    # For cosine distance, lower is better
                    score = float(distance)
                elif self.vidx_metric_type == "inner_product":
                    # For inner product, higher is better, so we negate the distance
                    score = -float(distance)
                else:
                    score = float(distance)

                search_results.append(self._create_output_data(vector_id, text_content, score, metadata))
            logger.debug(f"_vector_search results, len : {len(search_results)}")
            return search_results
            
        except Exception as e:
            logger.error(f"Vector search failed in collection '{self.collection_name}': {e}", exc_info=True)
            raise

    def _fulltext_search(self, query: str, limit: int = 5, filters: Optional[Dict] = None) -> list[OutputData]:
        """Perform full-text search using OceanBase FTS with parameterized queries including score."""
        # Skip search if query is empty
        if not query or not query.strip():
            logger.debug("Full-text search query is empty, returning empty results.")
            return []
        
        # Generate where clause from filters using the existing method
        filter_where_clause = self._generate_where_clause(filters)

        # Build the full-text search condition using SQLAlchemy text with parameters
        # Use the same parameter format that SQLAlchemy will use for other parameters
        fts_condition = text(f"MATCH({self.fulltext_field}) AGAINST(:query IN NATURAL LANGUAGE MODE)").bindparams(
            bindparam("query", query)
        )

        # Combine FTS condition with filter conditions
        where_conditions = [fts_condition]
        if filter_where_clause:
            where_conditions.extend(filter_where_clause)

        # Build custom query to include score field
        try:
            # Build select statement with specific columns AND score
            columns = [
                self.table.c[self.text_field],
                self.table.c[self.metadata_field],
                self.table.c[self.primary_field],
                self.table.c["user_id"],
                self.table.c["agent_id"],
                self.table.c["run_id"],
                self.table.c["actor_id"],
                self.table.c["hash"],
                self.table.c["created_at"],
                self.table.c["updated_at"],
                self.table.c["category"],
                # Add the score calculation as a column
                text(f"MATCH({self.fulltext_field}) AGAINST(:query IN NATURAL LANGUAGE MODE) as score").bindparams(
                    bindparam("query", query)
                )
            ]

            stmt = select(*columns)

            # Add where conditions
            for condition in where_conditions:
                stmt = stmt.where(condition)

            # Order by score DESC to get best matches first
            stmt = stmt.order_by(text('score DESC'))

            # Add limit
            if limit:
                stmt = stmt.limit(limit)

            # Execute the query with parameters - use direct parameter passing
            with self.obvector.engine.connect() as conn:
                with conn.begin():
                    logger.info(f"Executing FTS query with parameters: query={query}")
                    # Execute with parameter dictionary - the standard SQLAlchemy way
                    results = conn.execute(stmt)
                    rows = results.fetchall()

        except Exception as e:
            logger.warning(f"Full-text search failed, falling back to LIKE search: {e}")
            try:
                # Fallback to simple LIKE search with parameters
                like_query = f"%{query}%"
                like_condition = text(f"{self.fulltext_field} LIKE :like_query").bindparams(
                    bindparam("like_query", like_query)
                )

                fallback_conditions = [like_condition]
                if filter_where_clause:
                    fallback_conditions.extend(filter_where_clause)

                # Build fallback query with default score
                columns = [
                    self.table.c[self.text_field],
                    self.table.c[self.metadata_field],
                    self.table.c[self.primary_field],
                    self.table.c["user_id"],
                    self.table.c["agent_id"],
                    self.table.c["run_id"],
                    self.table.c["actor_id"],
                    self.table.c["hash"],
                    self.table.c["created_at"],
                    self.table.c["updated_at"],
                    self.table.c["category"],
                    # Default score for LIKE search
                    '1.0 as score'
                ]

                stmt = select(*columns)

                for condition in fallback_conditions:
                    stmt = stmt.where(condition)

                if limit:
                    stmt = stmt.limit(limit)

                # Execute fallback query with parameters
                with self.obvector.engine.connect() as conn:
                    with conn.begin():
                        logger.info(f"Executing LIKE fallback query with parameters: like_query={like_query}")
                        # Execute with parameter dictionary - the standard SQLAlchemy way
                        results = conn.execute(stmt)
                        rows = results.fetchall()
            except Exception as fallback_error:
                logger.error(f"Both full-text search and LIKE fallback failed: {fallback_error}")
                return []

        # Convert results to OutputData objects
        fts_results = []
        for row in rows:
            # Parse the row data including score as the last column
            (text_content, metadata_json, vector_id, user_id, agent_id, run_id, actor_id, hash_val,
             created_at, updated_at, category, fts_score) = self._parse_row(row)

            # Build standard metadata
            metadata = self._build_standard_metadata(
                user_id, agent_id, run_id, actor_id, hash_val,
                created_at, updated_at, category, metadata_json
            )

            # Use the actual FTS score from the query
            fts_results.append(self._create_output_data(vector_id, text_content, float(fts_score), metadata))

        logger.info(f"_fulltext_search results, len : {len(fts_results)}, fts_results : {fts_results}")
        return fts_results

    def _hybrid_search(self, query: str, vectors: List[List[float]], limit: int = 5, filters: Optional[Dict] = None,
                       fusion_method: str = "rrf", k: int = 60):
        """Perform hybrid search combining vector and full-text search with optional reranking."""
        # Determine candidate limit for reranking
        candidate_limit = limit * 3 if self.reranker else limit

        # Perform vector search and full-text search in parallel for better performance
        with ThreadPoolExecutor(max_workers=2) as executor:
            # Submit both searches concurrently
            vector_future = executor.submit(self._vector_search, query, vectors, candidate_limit, filters)
            fts_future = executor.submit(self._fulltext_search, query, candidate_limit, filters)
            # Wait for both to complete and get results
            vector_results = vector_future.result()
            fts_results = fts_future.result()

        # Step 1: Coarse ranking - Combine results using RRF or weighted fusion
        coarse_ranked_results = self._combine_search_results(
            vector_results, fts_results, candidate_limit, fusion_method, k
        )
        logger.debug(f"Coarse ranking completed, candidates: {len(coarse_ranked_results)}")
        
        # Step 2: Fine ranking - Use Rerank model for precision sorting (if enabled)
        if self.reranker and query and coarse_ranked_results:
            try:
                final_results = self._apply_rerank(query, coarse_ranked_results, limit)
                logger.debug(f"Rerank applied, final results: {len(final_results)}")
                return final_results
            except Exception as e:
                logger.warning(f"Rerank failed, falling back to coarse ranking: {e}")
                return coarse_ranked_results[:limit]
        else:
            # No reranker, return coarse ranking results
            return coarse_ranked_results[:limit]

    def _apply_rerank(self, query: str, candidates: List[OutputData], limit: int) -> List[OutputData]:
        """
        Apply Rerank model for precision sorting.
        
        Args:
            query: Search query text
            candidates: Candidate results from coarse ranking
            limit: Number of final results to return
            
        Returns:
            List of reranked OutputData objects
        """
        if not candidates:
            return []
        
        # Extract document texts from candidates
        documents = [result.payload.get('data', '') for result in candidates]

        # Call reranker to get reranked indices and scores
        reranked_indices = self.reranker.rerank(query, documents, top_n=limit)
        
        # Reconstruct results with rerank scores
        final_results = []
        for idx, rerank_score in reranked_indices:
            result = candidates[idx]
            # Preserve original scores in payload
            result.payload['_fusion_score'] = result.score
            # Update score to rerank score
            result.score = rerank_score
            result.payload['_rerank_score'] = rerank_score
            final_results.append(result)
        
        # Reorder results: high scores on both ends, low scores in the middle
        if len(final_results) > 1:
            reordered = [None] * len(final_results)
            left = 0
            right = len(final_results) - 1
            
            for i, result in enumerate(final_results):
                if i % 2 == 0:
                    # Even indices go to the left side
                    reordered[left] = result
                    left += 1
                else:
                    # Odd indices go to the right side
                    reordered[right] = result
                    right -= 1
            
            final_results = reordered
        
        logger.debug(f"Rerank completed: {len(final_results)} results")

        return final_results

    def _combine_search_results(self, vector_results: List[OutputData], fts_results: List[OutputData],
                                limit: int, fusion_method: str = "rrf", k: int = 60):
        """Combine and rerank vector and full-text search results using RRF or weighted fusion."""
        if fusion_method == "rrf":
            return self._rrf_fusion(vector_results, fts_results, limit, k)
        else:
            return self._weighted_fusion(vector_results, fts_results, limit)

    def _rrf_fusion(self, vector_results: List[OutputData], fts_results: List[OutputData],
                    limit: int, k: int = 60):
        """
        Reciprocal Rank Fusion (RRF) for combining search results.
        
        Uses weights configured at initialization (self.vector_weight and self.fts_weight)
        to control the contribution of vector search vs full-text search.
        """
        # Create mapping of document ID to result data
        all_docs = {}

        # Process vector search results (rank-based scoring with weight)
        for rank, result in enumerate(vector_results, 1):
            rrf_score = self.vector_weight * (1.0 / (k + rank))
            all_docs[result.id] = {
                'result': result,
                'vector_rank': rank,
                'fts_rank': None,
                'rrf_score': rrf_score
            }

        # Process FTS results (add or update RRF scores with weight)
        for rank, result in enumerate(fts_results, 1):
            fts_rrf_score = self.fts_weight * (1.0 / (k + rank))

            if result.id in all_docs:
                # Document found in both searches - combine RRF scores
                all_docs[result.id]['fts_rank'] = rank
                all_docs[result.id]['rrf_score'] += fts_rrf_score
            else:
                # Document only in FTS results
                all_docs[result.id] = {
                    'result': result,
                    'vector_rank': None,
                    'fts_rank': rank,
                    'rrf_score': fts_rrf_score
                }

        # Convert to final results and sort by RRF score
        heap = []
        for doc_id, doc_data in all_docs.items():
            # Use document ID as tiebreaker to avoid dict comparison when rrf_scores are equal
            if len(heap) < limit:
                heapq.heappush(heap, (doc_data['rrf_score'], doc_id, doc_data))
            elif doc_data['rrf_score'] > heap[0][0]:
                heapq.heapreplace(heap, (doc_data['rrf_score'], doc_id, doc_data))

        final_results = []
        for score, _, doc_data in sorted(heap, key=lambda x: x[0], reverse=True):
            result = doc_data['result']
            result.score = score
            # Add ranking information to metadata for debugging
            result.payload['_fusion_info'] = {
                'vector_rank': doc_data['vector_rank'],
                'fts_rank': doc_data['fts_rank'],
                'rrf_score': score,
                'fusion_method': 'rrf',
                'vector_weight': self.vector_weight,
                'fts_weight': self.fts_weight
            }
            final_results.append(result)

        return final_results

    def _weighted_fusion(self, vector_results: List[OutputData], fts_results: List[OutputData],
                         limit: int, vector_weight: float = 0.7, text_weight: float = 0.3):
        """Traditional weighted score fusion (fallback method)."""
        # Create a mapping of id to results for deduplication
        combined_results = {}

        # Normalize vector scores to 0-1 range
        if vector_results:
            vector_scores = [result.score for result in vector_results]
            min_vector_score = min(vector_scores)
            max_vector_score = max(vector_scores)
            vector_score_range = max_vector_score - min_vector_score

            for result in vector_results:
                if vector_score_range > 0:
                    # For distance metrics, lower is better, so we invert the normalized score
                    if self.vidx_metric_type in ["l2", "cosine"]:
                        normalized_score = 1.0 - (result.score - min_vector_score) / vector_score_range
                    else:  # inner_product
                        normalized_score = (result.score - min_vector_score) / vector_score_range
                else:
                    normalized_score = 1.0

                combined_results[result.id] = {
                    'result': result,
                    'vector_score': normalized_score,
                    'fts_score': 0.0
                }

        # Add FTS results (FTS scores are already normalized to 0-1)
        for result in fts_results:
            if result.id in combined_results:
                # Update existing result with FTS score
                combined_results[result.id]['fts_score'] = result.score
            else:
                # Add new FTS-only result
                combined_results[result.id] = {
                    'result': result,
                    'vector_score': 0.0,
                    'fts_score': result.score
                }

        # Calculate combined scores and create final results
        heap = []
        for doc_id, doc_data in combined_results.items():
            combined_score = (vector_weight * doc_data['vector_score'] +
                              text_weight * doc_data['fts_score'])

            if len(heap) < limit:
                heapq.heappush(heap, (combined_score, doc_id, doc_data))
            elif combined_score > heap[0][0]:
                heapq.heapreplace(heap, (combined_score, doc_id, doc_data))

        final_results = []
        for score, _, doc_data in sorted(heap, key=lambda x: x[0], reverse=True):
            result = doc_data['result']
            result.score = score
            # Add fusion info for debugging
            result.payload['_fusion_info'] = {
                'vector_score': doc_data['vector_score'],
                'fts_score': doc_data['fts_score'],
                'combined_score': score,
                'fusion_method': 'weighted'
            }
            final_results.append(result)

        # Return top results
        return final_results

    def delete(self, vector_id: int):
        """Delete a vector by ID."""
        try:
            self.obvector.delete(
                table_name=self.collection_name,
                ids=[vector_id],
            )
            logger.debug(f"Successfully deleted vector with ID: {vector_id} from collection '{self.collection_name}'")
        except Exception as e:
            logger.error(f"Failed to delete vector with ID {vector_id} from collection '{self.collection_name}': {e}", exc_info=True)
            raise

    def update(self, vector_id: int, vector: Optional[List[float]] = None, payload: Optional[Dict] = None):
        """Update a vector and its payload."""
        try:
            # Get existing record to preserve fields not being updated
            existing_result = self.obvector.get(
                table_name=self.collection_name,
                ids=[vector_id],
                output_column_name=[self.vector_field]  # Get the existing vector
            )

            existing_rows = existing_result.fetchall()
            if not existing_rows:
                logger.warning(f"Vector with ID {vector_id} not found in collection '{self.collection_name}'")
                return

            # Prepare update data
            update_data: Dict[str, Any] = {
                self.primary_field: vector_id,
            }

            if vector is not None:
                update_data[self.vector_field] = (
                    vector if not self.normalize else self._normalize(vector)
                )
            else:
                # Preserve the existing vector to avoid it being cleared by upsert
                existing_vector = existing_rows[0][0] if existing_rows[0] else None
                if existing_vector is not None:
                    update_data[self.vector_field] = existing_vector
                    logger.debug(f"Preserving existing vector for ID {vector_id}")

            if payload is not None:
                # Use the helper method to build fields, then merge with update_data
                temp_record = self._build_record_for_insert(vector or [], payload)

                # Copy relevant fields from temp_record (excluding primary key and vector if not updating)
                for key, value in temp_record.items():
                    if key != self.primary_field and (vector is not None or key != self.vector_field):
                        update_data[key] = value

            # Update record
            self.obvector.upsert(
                table_name=self.collection_name,
                data=[update_data],
            )
            logger.debug(f"Successfully updated vector with ID: {vector_id} in collection '{self.collection_name}'")
            
        except Exception as e:
            logger.error(f"Failed to update vector with ID {vector_id} in collection '{self.collection_name}': {e}", exc_info=True)
            raise

    def get(self, vector_id: int):
        """Retrieve a vector by ID."""
        try:
            results = self.obvector.get(
                table_name=self.collection_name,
                ids=[vector_id],
                output_column_name=[
                    self.vector_field,
                    self.text_field,
                    self.metadata_field,
                    "user_id",
                    "agent_id",
                    "run_id",
                    "actor_id",
                    "hash",
                    "created_at",
                    "updated_at",
                    "category",
                ],
            )

            rows = results.fetchall()
            if not rows:
                logger.debug(f"Vector with ID {vector_id} not found in collection '{self.collection_name}'")
                return None

            (vector, text_content, metadata_json, user_id, agent_id,
             run_id, actor_id, hash_val, created_at, updated_at, category, _) = self._parse_row(rows[0])

            # Build standard metadata
            metadata = self._build_standard_metadata(
                user_id, agent_id, run_id, actor_id, hash_val,
                created_at, updated_at, category, metadata_json
            )

            logger.debug(f"Successfully retrieved vector with ID: {vector_id} from collection '{self.collection_name}'")
            return self._create_output_data(vector_id, text_content, 0.0, metadata)
            
        except Exception as e:
            logger.error(f"Failed to get vector with ID {vector_id} from collection '{self.collection_name}': {e}", exc_info=True)
            raise

    def list_cols(self):
        """List all collections."""
        try:
            # Get all tables from the database using the correct SQLAlchemy API
            with self.obvector.engine.connect() as conn:
                result = conn.execute(text("SHOW TABLES"))
                tables = [row[0] for row in result.fetchall()]
                logger.debug(f"Successfully listed {len(tables)} collections")
                return tables
        except Exception as e:
            logger.error(f"Failed to list collections: {e}", exc_info=True)
            raise

    def delete_col(self):
        """Delete the collection."""
        try:
            if self.obvector.check_table_exists(self.collection_name):
                self.obvector.drop_table_if_exist(self.collection_name)
                logger.info(f"Successfully deleted collection '{self.collection_name}'")
            else:
                logger.warning(f"Collection '{self.collection_name}' does not exist, skipping deletion")
        except Exception as e:
            logger.error(f"Failed to delete collection '{self.collection_name}': {e}", exc_info=True)
            raise

    def _get_existing_vector_dimension(self) -> Optional[int]:
        """Get the dimension of the existing vector field in the table."""
        if not self.obvector.check_table_exists(self.collection_name):
            return None

        try:
            # Get table schema information using the correct SQLAlchemy API
            with self.obvector.engine.connect() as conn:
                result = conn.execute(text(f"DESCRIBE {self.collection_name}"))
                columns = result.fetchall()

            # Find the vector field and extract its dimension
            for col in columns:
                if col[0] == self.vector_field:
                    # Parse vector type like "VECTOR(1536)" to extract dimension
                    col_type = col[1]
                    if col_type.startswith("VECTOR(") and col_type.endswith(")"):
                        dim_str = col_type[7:-1]  # Extract dimension from "VECTOR(1536)"
                        return int(dim_str)
            return None
        except Exception as e:
            logger.warning(f"Failed to get vector dimension for table {self.collection_name}: {e}")
            return None

    def col_info(self):
        """Get information about the collection."""
        try:
            if not self.obvector.check_table_exists(self.collection_name):
                logger.debug(f"Collection '{self.collection_name}' does not exist")
                return None

            # Get table schema information using the correct SQLAlchemy API
            with self.obvector.engine.connect() as conn:
                result = conn.execute(text(f"DESCRIBE {self.collection_name}"))
                columns = result.fetchall()

            logger.debug(f"Successfully retrieved info for collection '{self.collection_name}'")
            return {
                "name": self.collection_name,
                "columns": [{"name": col[0], "type": col[1]} for col in columns],
                "index_type": self.index_type,
                "metric_type": self.vidx_metric_type,
            }
            
        except Exception as e:
            logger.error(f"Failed to get collection info for '{self.collection_name}': {e}", exc_info=True)
            raise

    def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None):
        """List all memories."""
        try:
            # Build where clause from filters
            where_clause = self._generate_where_clause(filters)

            # Get all records
            results = self.obvector.get(
                table_name=self.collection_name,
                ids=None,
                output_column_name=[
                    self.primary_field,
                    self.vector_field,
                    self.text_field,
                    self.metadata_field,
                    "user_id",
                    "agent_id",
                    "run_id",
                    "actor_id",
                    "hash",
                    "created_at",
                    "updated_at",
                    "category",
                ],
                where_clause=where_clause
            )

            memories = []
            for row in results.fetchall():
                (vector_id, vector, text_content, metadata_json, user_id, agent_id, run_id,
                 actor_id, hash_val, created_at, updated_at, category) = self._parse_row(row)

                # Build standard metadata
                metadata = self._build_standard_metadata(
                    user_id, agent_id, run_id, actor_id, hash_val,
                    created_at, updated_at, category, metadata_json
                )

                memories.append(self._create_output_data(vector_id, text_content, 0.0, metadata))

            if limit:
                memories = memories[:limit]

            logger.debug(f"Successfully listed {len(memories)} memories from collection '{self.collection_name}'")
            return [memories]
            
        except Exception as e:
            logger.error(f"Failed to list memories from collection '{self.collection_name}': {e}", exc_info=True)
            raise

    def reset(self):
        """Reset by deleting the collection and recreating it."""
        try:
            logger.info(f"Resetting collection '{self.collection_name}'")
            self.delete_col()
            if self.embedding_model_dims is not None:
                self._create_table_with_index_by_embedding_model_dims()

            if self.hybrid_search:
                self._check_and_create_fulltext_index()
                
            logger.info(f"Successfully reset collection '{self.collection_name}'")
            
        except Exception as e:
            logger.error(f"Failed to reset collection '{self.collection_name}': {e}", exc_info=True)
            raise

    def _check_and_create_fulltext_index(self):
        # Check whether the full-text index exists, if not, create it
        if not self._check_fulltext_index_exists():
            self._create_fulltext_index()

    def _check_fulltext_index_exists(self) -> bool:
        """
        Check if the full-text index of the specified table exists.
        """
        try:
            with self.obvector.engine.connect() as conn:
                result = conn.execute(text(f"SHOW INDEX FROM {self.collection_name}"))
                indexes = result.fetchall()

                for index in indexes:
                    # Index [2] is the index name, index [4] is the column name, and index [10] is the index type
                    if len(index) > 10 and index[10] == 'FULLTEXT':
                        if self.fulltext_field in str(index[4]):
                            return True

            return False

        except Exception as e:
            logger.error(f"An error occurred while checking the full-text index: {e}")
            return False

    def _create_fulltext_index(self):
        try:
            logger.debug(
                "About to create fulltext index for collection '%s' using parser '%s'",
                self.collection_name,
                self.fulltext_parser,
            )

            # Create fulltext index with the specified parser using SQL
            with self.obvector.engine.connect() as conn:
                sql_command = text(f"""ALTER TABLE {self.collection_name}
                    ADD FULLTEXT INDEX fulltext_index_for_col_text ({self.fulltext_field}) WITH PARSER {self.fulltext_parser}""")

                logger.debug("DEBUG: Executing SQL: %s", sql_command)
                conn.execute(sql_command)
                logger.debug("DEBUG: Fulltext index created successfully for '%s'", self.collection_name)

        except Exception as e:
            logger.exception("Exception occurred while creating fulltext index")
            raise Exception(
                "Failed to add fulltext index to the target table, your OceanBase version must be "
                "4.3.5.1 or above to support fulltext index and vector index in the same table"
            ) from e

        # Refresh metadata
        self.obvector.refresh_metadata([self.collection_name])
    
    def execute_sql(self, sql: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """
        Execute a raw SQL statement and return results.
        
        This method is used by SubStoreMigrationManager to manage migration status table.
        
        Args:
            sql: SQL statement to execute
            params: Optional parameters for the SQL statement
            
        Returns:
            List of result rows as dictionaries
        """
        try:
            with self.obvector.engine.connect() as conn:
                if params:
                    result = conn.execute(text(sql), params)
                else:
                    result = conn.execute(text(sql))
                
                # Commit for DDL/DML statements
                conn.commit()
                
                # Try to fetch results (for SELECT queries)
                try:
                    rows = result.fetchall()
                    # Convert rows to dictionaries
                    if rows and result.keys():
                        return [dict(zip(result.keys(), row)) for row in rows]
                    return []
                except Exception:
                    # No results to fetch (for INSERT/UPDATE/DELETE/CREATE)
                    return []
                    
        except Exception as e:
            logger.error(f"Failed to execute SQL: {e}")
            logger.debug(f"SQL statement: {sql}")
            raise