# Copyright 2025-present the zvec project
#
# Licensed 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.
from __future__ import annotations

from typing import Optional, Union

import numpy as np

from zvec._zvec import _Collection, _MultiQuery
from zvec._zvec.param import _Fts, _SearchQuery, _SubQuery

from ..extension import CallbackReRanker, ReRanker, RrfReRanker, WeightedReRanker
from ..model.convert import convert_to_py_doc
from ..model.doc import DocList
from ..model.param.query import Query
from ..model.schema import CollectionSchema
from ..typing import DataType

__all__ = [
    "QueryContext",
    "QueryExecutor",
]

DTYPE_MAP = {
    DataType.VECTOR_FP16.value: np.float16,
    DataType.VECTOR_FP32.value: np.float32,
    DataType.VECTOR_FP64.value: np.float64,
    DataType.VECTOR_INT8.value: np.int8,
}


def convert_to_numpy(vec: Union[list, np.ndarray], dtype: np.dtype) -> np.ndarray:
    if isinstance(vec, np.ndarray):
        if vec.dtype == dtype and vec.ndim == 1:
            return vec
        return np.asarray(vec, dtype=dtype).flatten()

    try:
        arr = np.asarray(vec, dtype=dtype)
        if arr.ndim != 1:
            arr = arr.flatten()
        return arr
    except (ValueError, TypeError) as e:
        raise TypeError(
            f"Cannot convert input to 1D numpy array with dtype={dtype}: {type(vec)}"
        ) from e


class QueryContext:
    def __init__(
        self,
        topk: int,
        filter: Optional[str] = None,
        include_vector: bool = False,
        queries: Optional[list[Query]] = None,
        output_fields: Optional[list[str]] = None,
        reranker: Optional[ReRanker] = None,
    ):
        # query param
        self._filter = filter
        self._queries = queries or []
        self._topk = topk
        self._include_vector = include_vector
        self._output_fields = output_fields

        # reranker
        self._reranker = reranker

    @property
    def topk(self):
        return self._topk

    @property
    def queries(self):
        return self._queries

    @property
    def filter(self):
        return self._filter

    @property
    def reranker(self):
        return self._reranker

    @property
    def output_fields(self):
        return self._output_fields

    @property
    def include_vector(self):
        return self._include_vector


class QueryExecutor:
    """Unified query executor that routes based on query count and reranker type."""

    def __init__(self, schema: CollectionSchema):
        self._schema = schema

    def _build_queries(
        self, ctx: QueryContext, collection: _Collection
    ) -> list[_SearchQuery]:
        """Build query vector list (no validation, conversion only)."""
        if not ctx.queries:
            return [self._build_base_search_query(ctx)]
        return [
            self._build_search_query(ctx, query, collection) for query in ctx.queries
        ]

    def execute(self, ctx: QueryContext, collection: _Collection) -> DocList:
        """Execute a query, routing by query count.

        A single (or vector-less) query is sent to C++ as a ``_SearchQuery``;
        multiple queries are assembled into a ``_MultiQuery``.
        """
        queries = self._build_queries(ctx, collection)
        if not queries:
            raise ValueError("No query to execute")

        if len(queries) == 1:
            return self._execute_single_query(queries[0], collection)
        return self._execute_multi_query(ctx, queries, collection)

    def _execute_single_query(
        self, query: _SearchQuery, collection: _Collection
    ) -> DocList:
        """Single/vector-less query: send a ``_SearchQuery`` to C++."""
        docs = collection.Query(query)
        return [convert_to_py_doc(doc, self._schema) for doc in docs]

    def _execute_multi_query(
        self, ctx: QueryContext, queries: list[_SearchQuery], collection: _Collection
    ) -> DocList:
        """Multiple queries: send a ``_MultiQuery`` to C++.

        A Python-only reranker (e.g. a model/API-based one) cannot run inside
        the C++ MultiQuery, so each route is executed individually and merged by
        the reranker in Python. The built-in RRF/Weighted/Callback rerankers use
        the C++ variant-based fast path.
        """
        reranker = ctx.reranker
        if reranker is None:
            raise ValueError(
                "A reranker is required to merge results from multiple queries; "
                "specify the 'reranker' argument."
            )
        if not isinstance(reranker, (RrfReRanker, WeightedReRanker, CallbackReRanker)):
            docs_list = self._execute_python_pipeline(queries, collection)
            return self._merge_and_rerank(ctx, docs_list)

        multi_query = self._build_multi_query(ctx, queries)
        docs = collection.Query(multi_query)
        return [convert_to_py_doc(doc, self._schema) for doc in docs]

    def _build_multi_query(
        self, ctx: QueryContext, queries: list[_SearchQuery]
    ) -> _MultiQuery:
        """Assemble a C++ ``_MultiQuery`` from per-route ``_SearchQuery`` objects."""
        multi_query = _MultiQuery()
        multi_query.queries = [_SubQuery.from_search_query(query) for query in queries]
        # num_candidates controls per-sub-query candidate count for reranking pool.
        # It must NOT be limited to the final output topk; use at least the C++
        # SubQuery default of 10 to ensure sufficient candidates for reranking.
        _DEFAULT_NUM_CANDIDATES = 10
        for sub in multi_query.queries:
            sub.num_candidates = max(ctx.topk, _DEFAULT_NUM_CANDIDATES)
        multi_query.topk = ctx.topk
        if ctx.filter:
            multi_query.filter = ctx.filter
        multi_query.include_vector = ctx.include_vector
        if ctx.output_fields is not None:
            multi_query.output_fields = ctx.output_fields
        # Set rerank strategy via the C++ variant-based API.
        reranker = ctx.reranker
        if isinstance(reranker, RrfReRanker):
            multi_query.set_rerank_rrf(reranker.rank_constant)
        elif isinstance(reranker, WeightedReRanker):
            multi_query.set_rerank_weighted(reranker.weights)
        elif isinstance(reranker, CallbackReRanker):
            multi_query.set_rerank_callback(reranker._callback)
        return multi_query

    def _execute_python_pipeline(
        self, vectors: list[_SearchQuery], collection: _Collection
    ) -> list[DocList]:
        """Execute queries serially for the Python-only reranker path."""
        return [self._execute_single_query(query, collection) for query in vectors]

    def _merge_and_rerank(self, ctx: QueryContext, docs_list: list[DocList]) -> DocList:
        """Merge and rerank results from the Python pipeline path."""
        if not docs_list:
            raise ValueError("Query results is empty")
        if len(docs_list) == 1 and not ctx.reranker:
            return docs_list[0]
        return ctx.reranker.rerank(docs_list, ctx.topk)

    def _build_base_search_query(self, ctx: QueryContext) -> _SearchQuery:
        search_query = _SearchQuery()
        search_query.topk = ctx.topk
        search_query.include_vector = ctx.include_vector
        if ctx.filter:
            search_query.filter = ctx.filter
        if ctx.output_fields is not None:
            search_query.output_fields = ctx.output_fields
        return search_query

    def _apply_fts(self, query: Query, search_query: _SearchQuery) -> None:
        """Set FTS query on search_query if the query has FTS parameters."""
        if query.has_fts():
            fts = _Fts()
            fts.query_string = query.fts.query_string or ""
            fts.match_string = query.fts.match_string or ""
            search_query.fts = fts

    def _build_search_query(
        self, ctx: QueryContext, query: Query, collection: _Collection
    ) -> _SearchQuery:
        query._validate()
        search_query = self._build_base_search_query(ctx)
        search_query.field_name = query.field_name
        if query.param:
            search_query.query_params = query.param

        # set FTS query if provided
        self._apply_fts(query, search_query)

        vector_schema = None
        if query.has_vector() or query.has_id():
            vector_schema = (
                self._schema.vector(query.field_name)
                if query
                else self._schema.vectors[0]
            )

            if vector_schema is None:
                raise ValueError("No vector field found")

        # set vector
        if query.has_vector():
            vec_data = query.vector
        elif query.has_id():
            fetched = collection.Fetch([query.id])
            doc = next(iter(fetched.values()), None)
            if not doc:
                raise ValueError(f"Document with id '{query.id}' not found")
            vec_data = doc.get_any(vector_schema.name, vector_schema.data_type)
        else:
            return search_query

        target_dtype = DTYPE_MAP.get(vector_schema.data_type.value)
        search_query.set_vector(
            vector_schema._get_object(),
            convert_to_numpy(vec_data, target_dtype) if target_dtype else vec_data,
        )
        return search_query