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

"""Native batch and streaming planning behind the public scan facade."""

import logging
from dataclasses import replace
from typing import List, Optional, Tuple

from pypaimon.common.identifier import UNKNOWN_DATABASE
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.predicate import Predicate
from pypaimon.execution.native import runtime
from pypaimon.execution.native import table as native_table
from pypaimon.execution.native.read import builder as native_builder
from pypaimon.read.plan import Plan
from pypaimon.read.query_auth_split import QueryAuthSplit
from pypaimon.read.scan_request import ScanRequest
from pypaimon.read import split_serializer

logger = logging.getLogger(__name__)

_NATIVE_TIME_TRAVEL_OPTIONS = frozenset({
    CoreOptions.SCAN_VERSION.key(), CoreOptions.SCAN_SNAPSHOT_ID.key(),
    CoreOptions.SCAN_TAG_NAME.key(), CoreOptions.SCAN_TIMESTAMP.key(),
    CoreOptions.SCAN_TIMESTAMP_MILLIS.key(), CoreOptions.SCAN_WATERMARK.key(),
})


def native_plan(
        table,
        predicate: Optional[Predicate] = None,
        limit: Optional[int] = None,
        projection: Optional[List[str]] = None,
        row_ranges: Optional[List[Tuple[int, int]]] = None,
        incremental_range: Optional[Tuple[int, int]] = None,
        incremental_mode: str = 'delta',
        row_position_slice: Optional[Tuple[int, int]] = None,
        row_position_shard: Optional[Tuple[int, int]] = None,
        chunk_shuffle: Optional[Tuple[int, int]] = None,
        shard: Optional[Tuple[int, int]] = None,
        partition_predicate: Optional[Predicate] = None,
        global_index_result=None,
        file_filters: Optional[Tuple[bool, bool]] = None,
        all_files: bool = False) -> Plan:
    """Plan with pypaimon_rust, preserving snapshot metadata.

    Native conversion or planning failures are handled by TableScan, which
    falls back to the Python planner.
    """
    if all_files and incremental_range is not None:
        raise ValueError('planning all files does not accept incremental ranges')
    if incremental_range is None and incremental_mode != 'delta':
        raise ValueError('incremental_mode requires incremental_range')
    if not runtime.native_runtime_available():
        raise RuntimeError(
            "scan.native-plan.enabled needs the pypaimon-rust extension")
    builder = native_builder._configure_native_read_builder(
        native_builder._native_read_builder(table), predicate, limit, projection)
    if partition_predicate is not None:
        builder = builder.with_partition_filter(native_builder._predicate_to_native(partition_predicate))
    if row_ranges is not None:
        builder = builder.with_row_ranges(row_ranges)
    from pypaimon.table.source.global_index_split_result import GlobalIndexSplitResult
    supplied_pk_splits = (isinstance(global_index_result, GlobalIndexSplitResult)
                          and not table.options.query_auth_enabled)
    if supplied_pk_splits or incremental_range is None:
        scan = builder.new_scan()
    elif incremental_mode == 'delta':
        scan = builder.new_incremental_scan(*incremental_range)
    else:
        scan = builder.new_incremental_scan(*incremental_range, incremental_mode)
    if file_filters is not None:
        scan = scan.with_file_filters(*file_filters)
    if global_index_result is not None and not table.options.query_auth_enabled:
        # A PK result already selected a snapshot and physical files. Java's
        # PrimaryKeyBatchScan returns those selections before ordinary planning.
        from pypaimon_rust.datafusion import GlobalIndexResult, Split as NativeSplit
        from pypaimon.globalindex.vector_search_result import ScoredGlobalIndexResult
        from pypaimon.table.source.primary_key_scored_result import PrimaryKeyScoredResult
        if isinstance(global_index_result, GlobalIndexSplitResult):
            # Searches and scans share the same core GlobalIndexResult.
            result = (getattr(global_index_result, '_native_result', None)
                      if type(global_index_result) is PrimaryKeyScoredResult else None)
            # Java's PK result uses snapshot 0 for an uncommitted table,
            # whereas the core search uses None. Keep the Python result's ID.
            if result is not None and result.snapshot_id() is None:
                result = None
            if result is None:
                result = GlobalIndexResult.from_splits(
                    global_index_result.snapshot_id,
                    [NativeSplit.deserialize(split_serializer.serialize_split_v1(split))
                     for split in global_index_result.splits])
        elif isinstance(global_index_result, ScoredGlobalIndexResult):
            score = global_index_result.score_getter()
            result = GlobalIndexResult.from_scored_row_ids({
                row_id: score(row_id) for row_id in global_index_result.results()})
        else:
            result = GlobalIndexResult.from_row_ranges([
                (range_.from_, range_.to) for range_ in global_index_result.results().to_range_list()])
        scan = scan.with_global_index_result(result)
    if row_position_slice is not None:
        scan = scan.with_row_position_slice(*row_position_slice)
    if row_position_shard is not None:
        scan = scan.with_row_position_shard(*row_position_shard)
    if chunk_shuffle is not None:
        seed, chunk_size = chunk_shuffle
        scan = scan.with_chunk_shuffle(str(seed), chunk_size)
    if shard is not None:
        scan = scan.with_shard(*shard)
    rust_plan = scan.plan_all() if all_files else scan.plan()
    return _from_native_plan(table, rust_plan)


def native_stream_scan(table, predicate=None, read_type=None, bucket_filter=None, consumer_id=None):
    """Create Rust's stateful StreamTableScan with the canonical reader type."""
    builder = native_builder._configure_native_read_builder(
        native_builder._native_read_builder(table), predicate, None, None, read_type=read_type)
    scan = builder.new_stream_scan()
    if bucket_filter is not None:
        scan.with_bucket_filter(bucket_filter)
    if consumer_id is not None:
        scan.with_consumer_id(consumer_id)
    return scan


def _from_native_plan(table, rust_plan) -> Plan:
    """Keep one split/metadata bridge for batch and streaming planning."""
    rust_splits = rust_plan.splits()
    pfields = native_table._partition_fields(table)
    # Trimmed primary keys decode per-file min/max keys (PK merge-on-read).
    kfields = table.trimmed_primary_keys_fields
    splits = [
        split_serializer.deserialize_split_v1(split.serialize(), pfields, kfields)
        for split in rust_splits
    ]
    if table.options.native_read_enabled():
        # Retain the opaque Rust split next to the Python metadata view so an
        # unchanged native plan can be read without reserializing each split.
        # A caller that needs different metadata (such as an endpoint DV)
        # passes a new Python split, which the native reader converts at read
        # time from its current fields.
        for split, rust_split in zip(splits, rust_splits):
            split._native_split = rust_split
    snapshot_id = rust_plan.snapshot_id()
    return Plan(splits, snapshot_id=snapshot_id, _native_stats=rust_plan.stats())


class NativeTableScan:
    """Adapt the public scan configuration to Core planning."""

    _is_native = True

    def __init__(self, scan):
        self._scan = scan

    def plan(self, auth_result=None):
        # Core requests and captures query authorization with its own plan.
        return self._plan()

    def plan_all(self):
        return self._plan(all_files=True)

    def plan_files(self, skip_level0, only_read_real_buckets):
        """Plan complete file metadata before query result transformations."""
        try:
            request = self._scan._scan_request()
            builder = native_builder._configure_native_read_builder(
                native_builder._native_read_builder(self._scan.table),
                request.predicate, None, None)
            if request.partition_filter is not None:
                builder = builder.with_partition_filter(native_builder._predicate_to_native(
                    request.partition_key_predicate))
            plan = builder.new_scan().with_file_filters(
                skip_level0, only_read_real_buckets).plan_files()
            from pypaimon.execution.native.read.file_scan import from_native_file_plan
            return from_native_file_plan(self._scan.table, plan)
        except Exception as error:
            runtime._raise_if_native_fork_safety_error(error)
            logger.warning("Native file planning failed, falling back to Python: %s", error)
            return None

    def scan_with_stats(self, auth_result=None):
        # Core records counters in the same planning pass.
        plan = self.plan()
        return None if plan is None else (plan, plan.stats())

    def _native_plan_supported(self) -> bool:
        # Any probe failure (e.g. a remote schema/metadata read) must fall back, not fail the scan.
        try:
            return self._native_plan_supported_impl()
        except Exception as e:
            logger.warning("Native-plan capability probe failed, falling back: %s", e)
            return False

    def _native_plan_supported_impl(self) -> bool:
        """Fall back to the Python scanner for scans native can't carry:
        a primary-key table whose trimmed PK is empty (PK equals the partition
        key; Rust rejects this schema), unsupported time travel selectors,
        custom authorization, a missing pypaimon-rust, or a catalog / identifier Rust
        cannot reconstruct. Keep this capability gate in sync when adding scan
        features."""
        scan = self._scan
        from pypaimon.execution.native.table import _resolved_schema_file_io_options
        if not runtime.native_runtime_available():
            return False
        request = scan._scan_request()
        if not self._native_global_index_result_supported(request):
            return False
        if request.chunk_shuffle is not None:
            request.validate_chunk_shuffle_compat()
        if _resolved_schema_file_io_options(scan.table) is None:
            loader = getattr(
                getattr(scan.table, 'catalog_environment', None),
                'catalog_loader',
                None,
            )
            context_fn = getattr(loader, 'context', None)
            if not callable(context_fn):
                return False
            from pypaimon.execution.native.table import _catalog_metastore
            if _catalog_metastore(loader) is None:
                return False
            context = context_fn()
            catalog_options = getattr(context, 'options', None)
            if catalog_options is None:
                return False
            database_name = scan.table.identifier.get_database_name()
            if not database_name or database_name == UNKNOWN_DATABASE:
                return False
        if scan.table.options.query_auth_enabled:
            # The filesystem path has no catalog-issued authorization session.
            loader = getattr(getattr(scan.table, 'catalog_environment', None), 'catalog_loader', None)
            from pypaimon.execution.native.table import _catalog_metastore
            if loader is None or _catalog_metastore(loader) != 'rest':
                return False
            from pypaimon.execution.native.table import _native_query_auth_supported
            if not _native_query_auth_supported(scan.table):
                return False
        # Rust rejects schemas whose primary keys are all partition keys.
        if getattr(scan.table, 'is_primary_key_table', False) \
                and not scan.table.trimmed_primary_keys:
            return False
        options = scan.table.options.options
        from pypaimon.snapshot.time_travel_util import SCAN_KEYS
        unsupported_scan_keys = set(SCAN_KEYS) - _NATIVE_TIME_TRAVEL_OPTIONS
        # Snapshot/tag ranges and explicit incremental modes use the Python scanner.
        unsupported_scan_keys.update((
            CoreOptions.INCREMENTAL_BETWEEN.key(), CoreOptions.INCREMENTAL_BETWEEN_SCAN_MODE.key(),
            CoreOptions.INCREMENTAL_TO_AUTO_TAG.key(),
        ))
        return not any(options.contains_key(k) for k in unsupported_scan_keys)

    def _native_global_index_result_supported(self, request: Optional[ScanRequest] = None) -> bool:
        scan = self._scan
        request = scan._scan_request() if request is None else request
        result = request.global_index_result
        if result is None:
            return True
        if request.is_primary_key_table:
            from pypaimon.table.source.global_index_split_result import GlobalIndexSplitResult
            return isinstance(result, GlobalIndexSplitResult)
        if not request.data_evolution:
            return False
        from pypaimon.globalindex.global_index_result import GlobalIndexResult
        return isinstance(result, GlobalIndexResult)

    def _plan(self, all_files=False) -> Optional[Plan]:
        """Plan via pypaimon_rust, then drop partitions the predicate rejects.

        Predicate and limit are pushed into Rust planning and are still enforced
        by the reader. Snapshot metadata is preserved even when pruning removes
        every split.
        """

        scan = self._scan
        try:
            request = scan._scan_request()
            has_distribution = (request.idx_of_this_subtask is not None
                                or request.start_pos_of_this_subtask is not None)
            extra_options = {}
            # Core resolves timestamp options after authorization, including
            # empty windows. Never preselect snapshots with Python here.
            chunk_shuffle = request.chunk_shuffle
            if chunk_shuffle is not None:
                extra_options['chunk_shuffle'] = chunk_shuffle
            if (request.idx_of_this_subtask is not None
                    and (chunk_shuffle is not None or request.is_primary_key_table)):
                extra_options['shard'] = (
                    request.idx_of_this_subtask, request.number_of_para_subtasks)
            if (has_distribution and not request.is_primary_key_table
                    and chunk_shuffle is None):
                if request.idx_of_this_subtask is not None:
                    extra_options['row_position_shard'] = (
                        request.idx_of_this_subtask, request.number_of_para_subtasks)
                else:
                    extra_options['row_position_slice'] = (
                        request.start_pos_of_this_subtask, request.end_pos_of_this_subtask)
            row_ranges = list(request.row_ranges) if request.row_ranges is not None else None
            if request.global_index_result is not None:
                extra_options['global_index_result'] = request.global_index_result
            if request.file_filters is not None:
                extra_options['file_filters'] = request.file_filters
            if request.partition_filter is not None:
                # Explicit partition selectors are scan-only. Keep the ordinary
                # row predicate separate so core can apply Java's override and
                # masked-partition validation semantics.
                extra_options['partition_predicate'] = request.partition_key_predicate
            # Append positions still precede LIMIT. Rust owns PK shards and
            # source-backed sorted-index refinement before applying LIMIT.
            defer_limit = has_distribution and not request.is_primary_key_table
            if all_files:
                extra_options['all_files'] = True
            plan = native_plan(
                scan.table,
                predicate=request.read_predicate,
                limit=None if defer_limit else request.read_limit,
                projection=(
                    [field.name for field in request.read_type]
                    if request.read_type is not None and chunk_shuffle is None else None),
                row_ranges=row_ranges,
                **extra_options,
            )
            splits = plan.splits()
            partition_predicate = request.partition_key_predicate
            if (partition_predicate is not None and not (
                    request.is_primary_key_table and request.global_index_result is not None)
                    and not any(isinstance(split, QueryAuthSplit) for split in splits)):
                # Authorized user predicates observe masked partition values;
                # Rust has already applied only the safe raw partition tests.
                splits = [s for s in splits
                          if getattr(s, 'partition', None) is None
                          or partition_predicate.test(s.partition)]
            if chunk_shuffle is None and defer_limit:
                # A partial IndexedSplit plus a file-wide DV cardinality cannot
                # reveal how many deleted rows lie inside the selected range.
                # Keep every selected split and let the reader enforce LIMIT.
                if not (request.deletion_vectors_enabled
                        and not request.is_primary_key_table):
                    splits = request.apply_push_down_limit(splits)
            return replace(plan, _splits=splits, _stats_limit=request.read_limit if defer_limit else None)
        except Exception as e:
            runtime._raise_if_native_fork_safety_error(e)
            # Any native construction/planning/pruning failure -> fall back.
            logger.warning(
                "Native plan failed, falling back to the Python scanner: %s", e)
            return None