# 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.

from typing import Any, Callable, Dict, List, Optional, Tuple

import pyarrow as pa

from pypaimon.common.where_parser import parse_where_clause
from pypaimon.multimodal.blob_read import fetch_blob_bodies
from pypaimon.schema.data_types import is_blob_type
from pypaimon.table.special_fields import SpecialFields


def _select_arrow_columns(batch, columns):
    return batch.select(columns)


class ScanQuery:
    """Chainable scan wrapper for MultimodalTable."""

    def __init__(
            self,
            table,
            result_factory: Optional[Callable] = None):
        self._table = table
        self._predicate = None
        self._projection = None
        self._limit = None
        self._include_row_id = False
        self._result_factory = result_factory

    def where(self, predicate):
        predicate = self._coerce_predicate(predicate, "where()")
        if predicate is not None:
            self._predicate = self._and_predicate(self._predicate, predicate)
        return self

    def select(self, columns):
        if isinstance(columns, str):
            columns = [columns]
        self._projection = list(columns)
        return self

    def with_row_id(self):
        self._include_row_id = True
        return self

    def limit(self, limit: int):
        self._limit = limit
        return self

    def to_arrow(self):
        if self._result_factory is not None:
            return self._read_global_index_result(self._result_factory(self))

        read_builder = self._configured_read_builder()
        scan = read_builder.new_scan()
        plan = scan.plan()
        return read_builder.new_read().to_arrow(plan.splits())

    def _configured_read_builder(self):
        read_builder = self._table.new_read_builder()
        if self._predicate is not None:
            read_builder = read_builder.with_filter(self._predicate)
        projection = self._effective_projection()
        if projection is not None:
            read_builder = read_builder.with_projection(projection)
        if self._limit is not None:
            read_builder = read_builder.with_limit(self._limit)
        return read_builder

    def _effective_projection(self):
        if self._projection is None:
            if not self._include_row_id:
                return None
            projection = [field.name for field in self._table.fields]
        else:
            projection = list(self._projection)
        if self._include_row_id and SpecialFields.ROW_ID.name not in projection:
            projection.append(SpecialFields.ROW_ID.name)
        return projection

    def _read_global_index_result(self, result):
        read_builder = self._configured_read_builder()
        scan = read_builder.new_scan().with_global_index_result(result)
        plan = scan.plan()
        return read_builder.new_read().to_arrow(plan.splits())

    def to_pandas(self):
        return self.to_arrow().to_pandas()

    def to_list(self) -> List[dict]:
        return self.to_arrow().to_pylist()

    def to_ray(
            self,
            *,
            ray_remote_args: Optional[Dict[str, Any]] = None,
            concurrency: Optional[int] = None,
            override_num_blocks: Optional[int] = None,
            **read_args):
        """Read this scan as a Ray Dataset.

        BLOB columns are read as serialized descriptors. Use
        ``table.map_with_blobs(...)`` to fetch payload bytes on Ray workers.
        """
        if self._result_factory is not None:
            raise TypeError("to_ray is only supported on scan(), not search queries.")

        read_builder, file_io, visible_columns = self._blob_descriptor_query_read_builder()
        plan = read_builder.new_scan().plan()
        ds = read_builder.new_read().to_ray(
            plan.splits(),
            ray_remote_args=ray_remote_args,
            concurrency=concurrency,
            override_num_blocks=override_num_blocks,
            **read_args)
        if visible_columns is not None:
            ds = ds.map_batches(
                _select_arrow_columns,
                fn_kwargs={"columns": visible_columns},
                batch_format="pyarrow")
        setattr(ds, "_paimon_blob_file_io", file_io)
        setattr(ds, "_paimon_blob_columns", self._all_blob_columns())
        return ds

    def read_blobs(
            self, columns=None, *, parallelism: int = 64
    ) -> Tuple[pa.Table, Dict[str, List[Optional[bytes]]]]:
        """Materialise BLOB column(s) for the filtered rows with concurrent,
        coalesced ranged reads. Reads via blob-as-descriptor to skip the slow
        row-by-row blob resolution on multi-group data-evolution splits.

        ``columns`` picks the BLOB column(s) (default: all, intersected with
        ``select(...)``). Returns ``(scalar_arrow_table, {column: [bytes|None]})``,
        row-aligned. Use :meth:`stream_blobs` for a memory-bounded read.

        Unresolved blob-view columns are not supported and raise ``ValueError``.
        """
        blob_cols = self._resolve_blob_columns(columns)
        read_builder, file_io = self._blob_descriptor_read_builder(blob_cols)
        arrow = read_builder.new_read().to_arrow(
            read_builder.new_scan().plan().splits())
        bodies = self._fetch_bodies(
            file_io, arrow.select(blob_cols).to_pydict(), blob_cols, parallelism)
        scalar = arrow.select(self._scalar_columns(arrow.column_names))
        return scalar, bodies

    def stream_blobs(self, columns=None, *, parallelism: int = 64):
        """Memory-bounded streaming variant of :meth:`read_blobs`: yield
        ``(scalar_batch, {column: [bytes|None]})`` per Arrow batch, so peak memory
        is one batch rather than the whole result.
        """
        # Validate eagerly so a bad column raises here, not on the first next().
        blob_cols = self._resolve_blob_columns(columns)
        read_builder, file_io = self._blob_descriptor_read_builder(blob_cols)
        return self._iter_blobs(read_builder, file_io, blob_cols, parallelism)

    def _iter_blobs(self, read_builder, file_io, blob_cols, parallelism):
        reader = read_builder.new_read().to_arrow_batch_reader(
            read_builder.new_scan().plan().splits())
        try:
            for batch in reader:
                bodies = self._fetch_bodies(
                    file_io, batch.select(blob_cols).to_pydict(), blob_cols, parallelism)
                scalar = batch.select(self._scalar_columns(batch.schema.names))
                yield scalar, bodies
        finally:
            # Close the reader even if the caller breaks out early.
            if hasattr(reader, "close"):
                reader.close()

    def _blob_descriptor_query_read_builder(self):
        from pypaimon.common.options.core_options import CoreOptions
        read_table = self._table.copy({
            CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"
        })
        read_builder = read_table.new_read_builder()
        if self._predicate is not None:
            read_builder = read_builder.with_filter(self._predicate)
        projection = self._effective_projection()
        if projection is not None:
            internal_projection = list(projection)
            for name in self._predicate_fields():
                if name not in internal_projection:
                    internal_projection.append(name)
            read_builder = read_builder.with_projection(internal_projection)
        if self._limit is not None:
            read_builder = read_builder.with_limit(self._limit)
        if projection is None:
            return read_builder, read_table.file_io, None
        internal_column_names = [field.name for field in read_builder.read_type()]
        visible_columns = self._projected_output_columns(read_table, projection)
        if visible_columns == internal_column_names:
            visible_columns = None
        return read_builder, read_table.file_io, visible_columns

    @staticmethod
    def _projected_output_columns(table, projection):
        builder = table.new_read_builder().with_projection(projection)
        return [field.name for field in builder.read_type()]

    def _blob_descriptor_read_builder(self, blob_cols: List[str]):
        """Blob-as-descriptor read builder with this query's filter/projection/limit;
        returns ``(read_builder, file_io)``."""
        from pypaimon.common.options.core_options import CoreOptions
        read_table = self._table.copy({
            CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"
        })
        read_builder = read_table.new_read_builder()
        if self._predicate is not None:
            read_builder = read_builder.with_filter(self._predicate)
        projection = self._blob_read_projection(blob_cols)
        if projection is not None:
            read_builder = read_builder.with_projection(projection)
        if self._limit is not None:
            read_builder = read_builder.with_limit(self._limit)
        return read_builder, read_table.file_io

    @staticmethod
    def _fetch_bodies(file_io, data, blob_cols, parallelism):
        # Decode each descriptor to a (uri, offset, length) range and read them all in
        # one coalesced pass on ``file_io`` -- the read table's FileIO, which already
        # carries the merged DLF/OSS token. Going through Blob.from_bytes here would
        # build a BlobRef whose uri_reader re-resolves the URI via
        # ``FileIO.get(uri, catalog_options)`` off the raw options (no merged token),
        # failing with "endpoint should be non-empty" / "Init credential failed" unless
        # the caller also passes fs.oss.* -- which users should not have to.
        return fetch_blob_bodies(file_io, data, blob_cols, parallelism)

    def _all_blob_columns(self) -> List[str]:
        return [
            field.name for field in self._table.fields
            if is_blob_type(field.type)
        ]

    def _resolve_blob_columns(self, columns) -> List[str]:
        all_blob = self._all_blob_columns()
        if columns is None:
            selected = all_blob
            if self._projection is not None:
                selected = [c for c in all_blob if c in self._projection]
            if not selected:
                raise ValueError("No BLOB column to read; pass columns=.")
            return selected
        if isinstance(columns, str):
            columns = [columns]
        columns = list(dict.fromkeys(columns))  # de-dup, keep order
        for name in columns:
            if name not in all_blob:
                raise ValueError(f"Column {name!r} is not a BLOB column.")
        return columns

    def _blob_read_projection(self, blob_cols: List[str]) -> List[str]:
        # Base on _effective_projection() (honours with_row_id), drop BLOB
        # descriptors, then append the requested BLOB columns and every predicate
        # column -- including predicate columns that are themselves BLOB (read as
        # descriptor) -- so SplitRead keeps the row-level filter for where().
        blob_set = set(self._all_blob_columns())
        effective = self._effective_projection()
        if effective is None:
            base = [f.name for f in self._table.fields if f.name not in blob_set]
        else:
            base = [name for name in effective if name not in blob_set]
        for name in self._predicate_fields():
            if name not in base and name not in blob_cols:
                base.append(name)
        return base + list(blob_cols)

    def _scalar_columns(self, available) -> List[str]:
        # Non-BLOB columns to expose, based on _effective_projection() so
        # with_row_id() is honoured; drop BLOBs and predicate-only helpers, and
        # skip unknown projected names to match to_arrow()'s silent drop.
        blob_set = set(self._all_blob_columns())
        effective = self._effective_projection()
        if effective is None:
            return [name for name in available if name not in blob_set]
        available = set(available)
        return [name for name in effective
                if name in available and name not in blob_set]

    def _predicate_fields(self):
        if self._predicate is None:
            return ()
        from pypaimon.read.push_down_utils import _get_all_fields
        return _get_all_fields(self._predicate)

    def _coerce_predicate(self, predicate, method):
        if predicate is None:
            return None
        if isinstance(predicate, str):
            return parse_where_clause(predicate, self._table.fields)
        raise ValueError("%s expects a SQL-like string." % method)

    def _and_predicate(self, left, right):
        if left is None:
            return right
        from pypaimon.common.predicate_builder import PredicateBuilder
        return PredicateBuilder.and_predicates([left, right])


class _PreFilterQuery(ScanQuery):

    def __init__(
            self,
            table,
            result_factory: Optional[Callable] = None,
            pre_filter=None):
        self._pre_filter = None
        super().__init__(table, result_factory=result_factory)
        if pre_filter is not None:
            self.pre_filter(pre_filter)

    def pre_filter(self, predicate):
        predicate = self._coerce_predicate(predicate, "pre_filter()")
        if predicate is not None:
            self._pre_filter = self._and_predicate(self._pre_filter, predicate)
        return self

    # Blob reads bypass the search result_factory, so reject them (else a plain scan).
    def read_blobs(self, *args, **kwargs):
        raise TypeError("read_blobs is only supported on scan(), not search queries.")

    def stream_blobs(self, *args, **kwargs):
        raise TypeError("stream_blobs is only supported on scan(), not search queries.")

    def to_ray(self, *args, **kwargs):
        raise TypeError("to_ray is only supported on scan(), not search queries.")


class VectorQuery(_PreFilterQuery):
    """Chainable query wrapper for vector global-index search."""

    def __init__(
            self,
            table,
            vector,
            vector_column,
            vector_options=None,
            pre_filter=None):
        self._vector = vector
        self._vector_column = vector_column
        self._vector_options = dict(vector_options or {})
        super().__init__(
            table, result_factory=self._execute_vector, pre_filter=pre_filter)

    def _execute_vector(self, query):
        limit = query._limit if query._limit is not None else 10
        builder = (
            self._table.new_vector_search_builder()
            .with_vector_column(self._vector_column)
            .with_query_vector(self._vector)
            .with_limit(limit)
            .with_options(self._vector_options)
        )
        if query._pre_filter is not None:
            builder = builder.with_filter(query._pre_filter)
        return builder.execute_local()


class TextQuery(_PreFilterQuery):
    """Chainable query wrapper for full-text global-index search."""

    def __init__(self, table, text_query, pre_filter=None):
        self._text_query = text_query
        super().__init__(
            table, result_factory=self._execute_fts, pre_filter=pre_filter)

    def _execute_fts(self, query):
        limit = query._limit if query._limit is not None else 10
        builder = (
            self._table.new_full_text_search_builder()
            .with_query(self._text_query["column"], self._text_query["query"])
            .with_limit(limit)
        )
        if query._pre_filter is not None:
            builder = builder.with_partition_filter(query._pre_filter)
        return builder.execute_local()


class HybridQuery(_PreFilterQuery):
    """Chainable query wrapper for hybrid global-index search."""

    def __init__(
            self,
            table,
            vector_routes=None,
            text_routes=None,
            ranker="rrf",
            route_limit=None,
            pre_filter=None):
        self._vector_routes = list(vector_routes or [])
        self._text_routes = list(text_routes or [])
        self._ranker = ranker
        self._route_limit = route_limit
        super().__init__(
            table, result_factory=self._execute_hybrid, pre_filter=pre_filter)

    def rerank(self, ranker):
        self._ranker = ranker
        return self

    def _execute_hybrid(self, query):
        final_limit = query._limit if query._limit is not None else 10
        route_limit = self._route_limit or final_limit
        builder = (
            self._table.new_hybrid_search_builder()
            .with_limit(final_limit)
            .with_ranker(self._ranker)
        )
        for route in self._vector_routes:
            builder = builder.add_vector_route(
                route["column"],
                route["vector"],
                limit=route.get("limit") or route_limit,
                weight=route["weight"],
                options=route["options"],
            )
        for route in self._text_routes:
            builder = builder.add_full_text_route(
                route["column"],
                route["query"],
                limit=route.get("limit") or route_limit,
                weight=route["weight"],
                options=route["options"],
            )
        if query._pre_filter is not None:
            builder = builder.with_filter(query._pre_filter)
        return builder.execute_local()


class BatchVectorQuery(_PreFilterQuery):
    """Chainable query wrapper for batch vector global-index search."""

    def __init__(
            self,
            table,
            vectors,
            vector_column,
            vector_options=None,
            pre_filter=None):
        self._vectors = vectors
        self._vector_column = vector_column
        self._vector_options = dict(vector_options or {})
        super().__init__(table, pre_filter=pre_filter)

    def to_arrow(self):
        return [
            self._read_global_index_result(result)
            for result in self._execute_batch_vector(self)
        ]

    def to_pandas(self):
        return [table.to_pandas() for table in self.to_arrow()]

    def to_list(self) -> List[List[dict]]:
        return [table.to_pylist() for table in self.to_arrow()]

    def _execute_batch_vector(self, query):
        limit = query._limit if query._limit is not None else 10
        builder = (
            self._table.new_batch_vector_search_builder()
            .with_vector_column(self._vector_column)
            .with_query_vectors(self._vectors)
            .with_limit(limit)
            .with_options(self._vector_options)
        )
        if query._pre_filter is not None:
            builder = builder.with_filter(query._pre_filter)
        return builder.execute_batch_local()
