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

"""Global-index coverage for data-evolution tables."""

from typing import Collection, Dict, List, Optional, Union

from pypaimon.common.options.core_options import CoreOptions, GlobalIndexSearchMode
from pypaimon.common.options.options import Options
from pypaimon.common.predicate import Predicate
from pypaimon.read.push_down_utils import _get_all_fields
from pypaimon.schema.data_types import DataField
from pypaimon.utils.range import Range


class DataEvolutionGlobalIndexCoverage:
    """Computes global-index coverage by field id."""

    def __init__(
        self,
        table,
        snapshot,
        partition_filter,
        index_files: Collection['IndexFileMeta'],
    ):
        self._table = table
        self._snapshot = snapshot
        self._partition_filter = partition_filter
        self._coverage_by_field: Dict[int, List[Range]] = {}
        for index_file in index_files:
            meta = index_file.global_index_meta
            if meta is None:
                continue
            row_range = Range(meta.row_range_start, meta.row_range_end)
            self._add_coverage(meta.index_field_id, row_range)
            if meta.extra_field_ids is not None:
                for extra_field_id in meta.extra_field_ids:
                    self._add_coverage(extra_field_id, row_range)

    def unindexed_ranges(
        self,
        fields_or_field_id: Union[List[DataField], Collection[int], int],
        predicate: Optional[Predicate] = None,
        search_mode: Optional[GlobalIndexSearchMode] = None,
    ) -> List[Range]:
        if isinstance(fields_or_field_id, int):
            field_ids = {fields_or_field_id}
        elif _is_field_id_collection(fields_or_field_id):
            field_ids = set(fields_or_field_id)
        else:
            field_by_name = {f.name: f for f in fields_or_field_id}
            field_ids = set()
            for name in _get_all_fields(predicate):
                field = field_by_name.get(name)
                if field is not None:
                    field_ids.add(field.id)
        return self._unindexed_ranges(field_ids, search_mode)

    def _add_coverage(self, field_id: int, row_range: Range) -> None:
        self._coverage_by_field.setdefault(field_id, []).append(row_range)

    def _indexed_ranges(self, field_ids: Collection[int]) -> List[Range]:
        ranges = None
        for field_id in field_ids:
            field_ranges = self._coverage_by_field.get(field_id)
            if not field_ranges:
                return []
            field_ranges = Range.sort_and_merge_overlap(field_ranges, True)
            ranges = field_ranges if ranges is None else Range.and_(ranges, field_ranges)
        if ranges is None:
            return []
        return Range.sort_and_merge_overlap(ranges, True)

    def _unindexed_ranges(
        self,
        field_ids: Collection[int],
        search_mode: Optional[GlobalIndexSearchMode] = None,
    ) -> List[Range]:
        if search_mode is None:
            search_mode = _scalar_index_search_mode(self._table)
        if search_mode == GlobalIndexSearchMode.FAST:
            return []
        next_row_id = getattr(self._snapshot, "next_row_id", None)
        if self._snapshot is None or next_row_id is None:
            return []
        if next_row_id <= 0:
            return []

        if search_mode == GlobalIndexSearchMode.DETAIL:
            data_ranges = self._data_ranges_by_data_files()
        else:
            data_ranges = [Range(0, next_row_id - 1)]

        indexed_ranges = Range.sort_and_merge_overlap(
            self._indexed_ranges(field_ids), True)
        unindexed = []
        for data_range in Range.sort_and_merge_overlap(data_ranges, True):
            unindexed.extend(data_range.exclude(indexed_ranges))
        return Range.sort_and_merge_overlap(unindexed, True)

    def _data_ranges_by_data_files(self) -> List[Range]:
        if hasattr(self._table, "data_ranges_for_data_evolution_global_index_coverage"):
            return self._table.data_ranges_for_data_evolution_global_index_coverage(
                self._snapshot,
                self._partition_filter,
            )

        manifest_list_manager = getattr(self._table, "manifest_list_manager", None)
        if manifest_list_manager is None:
            from pypaimon.manifest.manifest_list_manager import ManifestListManager
            manifest_list_manager = ManifestListManager(self._table)
        manifest_files = manifest_list_manager.read_all(self._snapshot)
        from pypaimon.manifest.manifest_file_manager import ManifestFileManager
        manager = ManifestFileManager(self._table)
        entries = manager.read_entries_parallel(
            manifest_files,
            self._entry_matches_partition,
            max_workers=self._table.options.scan_manifest_parallelism(),
        )
        data_ranges = []
        for entry in entries:
            row_range = entry.file.row_id_range()
            if row_range is not None:
                data_ranges.append(row_range)
        return data_ranges

    def _entry_matches_partition(self, entry) -> bool:
        if self._partition_filter is None:
            return True
        return self._partition_filter.test(entry.partition)


def _scalar_index_search_mode(table):
    options = getattr(table, "options", None)
    if options is None:
        return GlobalIndexSearchMode.FULL
    if hasattr(options, "scalar_index_search_mode"):
        return options.scalar_index_search_mode()
    return CoreOptions(Options.from_none()).scalar_index_search_mode()


def _is_field_id_collection(fields_or_field_id):
    if isinstance(fields_or_field_id, (str, bytes)):
        return False
    try:
        iterator = iter(fields_or_field_id)
    except TypeError:
        return False
    return all(isinstance(field_id, int) for field_id in iterator)
