# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""Builder to build vector search."""

from abc import ABC, abstractmethod

from pypaimon.common.predicate_builder import PredicateBuilder
from pypaimon.table.source.vector_search_read import VectorSearchReadImpl
from pypaimon.table.source.vector_search_scan import VectorSearchScanImpl


class VectorSearchBuilder(ABC):
    """Builder to build vector search."""

    @abstractmethod
    def with_limit(self, limit):
        # type: (int) -> VectorSearchBuilder
        """The top k results to return."""
        pass

    @abstractmethod
    def with_vector_column(self, name):
        # type: (str) -> VectorSearchBuilder
        """The vector column to search."""
        pass

    @abstractmethod
    def with_query_vector(self, vector):
        # type: (list) -> VectorSearchBuilder
        """The query vector (list of floats)."""
        pass

    @abstractmethod
    def with_filter(self, predicate):
        # type: (Predicate) -> VectorSearchBuilder
        """Scalar predicate used to pre-filter rows before vector search."""
        pass

    @abstractmethod
    def with_partition_filter(self, partition_filter):
        # type: (Predicate) -> VectorSearchBuilder
        """Partition predicate used to prune index manifest entries."""
        pass

    @abstractmethod
    def new_vector_search_scan(self):
        # type: () -> VectorSearchScan
        """Create vector search scan to scan index files."""
        pass

    @abstractmethod
    def new_vector_search_read(self):
        # type: () -> VectorSearchRead
        """Create vector search read to read index files."""
        pass

    def execute_local(self):
        # type: () -> GlobalIndexResult
        """Execute vector search locally."""
        return self.new_vector_search_read().read_plan(
            self.new_vector_search_scan().scan()
        )


class VectorSearchBuilderImpl(VectorSearchBuilder):
    """Implementation for VectorSearchBuilder."""

    def __init__(self, table):
        self._table = table
        self._limit = 0
        self._vector_column = None
        self._query_vector = None
        self._filter = None
        self._partition_filter = None

    def with_limit(self, limit):
        # type: (int) -> VectorSearchBuilder
        self._limit = limit
        return self

    def with_vector_column(self, name):
        # type: (str) -> VectorSearchBuilder
        field_dict = {f.name: f for f in self._table.fields}
        if name not in field_dict:
            raise ValueError("Vector column '%s' not found in table schema" % name)
        self._vector_column = field_dict[name]
        return self

    def with_query_vector(self, vector):
        # type: (list) -> VectorSearchBuilder
        self._query_vector = vector
        return self

    def with_filter(self, predicate):
        # type: (Predicate) -> VectorSearchBuilder
        if predicate is None:
            return self
        if self._filter is None:
            self._filter = predicate
        else:
            self._filter = PredicateBuilder.and_predicates([self._filter, predicate])
        # split out the partition-only conjuncts and store them as _partition_filter for
        # manifest pruning. Non-partition conjuncts remain in self._filter;
        # the silent drop of non-partition conjuncts *in the extracted copy*
        # is intentional — nothing is lost overall.
        extracted = self._extract_partition_only_conjuncts(predicate)
        if extracted is not None:
            if self._partition_filter is None:
                self._partition_filter = extracted
            else:
                self._partition_filter = PredicateBuilder.and_predicates(
                    [self._partition_filter, extracted])
        return self

    def with_partition_filter(self, partition_filter):
        # type: (Predicate) -> VectorSearchBuilder
        if partition_filter is None:
            self._partition_filter = None
            return self
        # Strict: every referenced field must be a partition key, otherwise a
        # non-partition conjunct would be silently dropped (with_filter has
        # the scalar fallback; with_partition_filter does not).
        partition_keys = list(self._table.partition_keys or [])
        if not partition_keys:
            raise ValueError(
                "with_partition_filter called on a non-partitioned table")
        from pypaimon.read.push_down_utils import _get_all_fields
        referenced = _get_all_fields(partition_filter)
        extras = referenced - set(partition_keys)
        if extras:
            raise ValueError(
                "Partition filter must reference only partition keys "
                "(%s); got non-partition field(s): %s"
                % (partition_keys, sorted(extras)))
        self._partition_filter = self._rebuild_leaf_indices_by_name(
            partition_filter,
            {name: idx for idx, name in enumerate(partition_keys)},
        )
        return self

    def _extract_partition_only_conjuncts(self, predicate):
        """AND-split ``predicate``, keep conjuncts that reference ONLY
        partition keys, and rebuild their leaf indices against the
        partition-only row by field name (so the caller's PredicateBuilder
        convention — full-row or partition-row — doesn't matter).
        """
        partition_keys = list(self._table.partition_keys or [])
        if not partition_keys:
            return None
        from pypaimon.read.push_down_utils import _split_and, _get_all_fields
        partition_key_set = set(partition_keys)
        pk_to_idx = {name: idx for idx, name in enumerate(partition_keys)}
        kept = [p for p in _split_and(predicate)
                if _get_all_fields(p).issubset(partition_key_set)]
        if not kept:
            return None
        rebuilt = [self._rebuild_leaf_indices_by_name(p, pk_to_idx)
                   for p in kept]
        return PredicateBuilder.and_predicates(rebuilt)

    @classmethod
    def _rebuild_leaf_indices_by_name(cls, predicate, pk_to_idx):
        """Return a copy of ``predicate`` with every leaf's ``index`` set to
        its position in ``pk_to_idx`` (field-name lookup). Input predicate may
        have been built against any schema — we key off ``Predicate.field``
        rather than ``Predicate.index`` so positional convention doesn't
        matter.
        """
        if predicate.method in ('and', 'or'):
            new_children = [cls._rebuild_leaf_indices_by_name(c, pk_to_idx)
                            for c in (predicate.literals or [])]
            return predicate.new_literals(new_children)
        return predicate.new_index(pk_to_idx[predicate.field])

    def new_vector_search_scan(self):
        # type: () -> VectorSearchScan
        if self._vector_column is None:
            raise ValueError("Vector column must be set via with_vector_column()")
        return VectorSearchScanImpl(
            self._table,
            self._vector_column,
            filter_=self._filter,
            partition_filter=self._partition_filter,
        )

    def new_vector_search_read(self):
        # type: () -> VectorSearchRead
        if self._limit <= 0:
            raise ValueError("Limit must be positive, set via with_limit()")
        if self._vector_column is None:
            raise ValueError("Vector column must be set via with_vector_column()")
        if self._query_vector is None:
            raise ValueError("Query vector must be set via with_query_vector()")
        return VectorSearchReadImpl(
            self._table,
            self._limit,
            self._vector_column,
            self._query_vector,
            filter_=self._filter,
        )