# 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 collections import defaultdict
from typing import List, Optional, Tuple

from pypaimon.globalindex.indexed_split import IndexedSplit
from pypaimon.utils.range import Range
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.manifest.schema.manifest_entry import ManifestEntry
from pypaimon.read.scanner.split_generator import AbstractSplitGenerator
from pypaimon.read.split import DataSplit, Split


class DataEvolutionSplitGenerator(AbstractSplitGenerator):
    """
    Split generator for data evolution tables.
    """

    def __init__(
        self,
        table,
        target_split_size: int,
        open_file_cost: int,
        deletion_files_map=None,
        row_ranges: Optional[List] = None,
        score_getter=None
    ):
        super().__init__(table, target_split_size, open_file_cost, deletion_files_map)
        self.row_ranges = row_ranges
        self.score_getter = score_getter

    def create_splits(self, file_entries: List[ManifestEntry]) -> List[Split]:
        """
        Create splits for data evolution tables.
        """
        def sort_key(manifest_entry: ManifestEntry) -> tuple:
            first_row_id = (
                manifest_entry.file.first_row_id
                if manifest_entry.file.first_row_id is not None
                else float('-inf')
            )
            is_blob = 1 if DataFileMeta.is_blob_file(manifest_entry.file.file_name) else 0
            max_seq = manifest_entry.file.max_sequence_number
            return first_row_id, is_blob, -max_seq

        sorted_entries = sorted(file_entries, key=sort_key)

        partitioned_files = defaultdict(list)
        for entry in sorted_entries:
            partitioned_files[(tuple(entry.partition.values), entry.bucket)].append(entry)

        slice_row_ranges = None  # Row ID ranges for slice-based filtering

        if self.start_pos_of_this_subtask is not None or self.idx_of_this_subtask is not None:
            # Calculate Row ID range for slice-based filtering
            slice_row_ranges = self._calculate_slice_row_ranges(partitioned_files)
            # Filter files by Row ID range
            partitioned_files = self._filter_files_by_row_ranges(partitioned_files, slice_row_ranges)

        def weight_func(file_list: List[DataFileMeta]) -> int:
            return max(sum(f.file_size for f in file_list), self.open_file_cost)

        splits = []
        for key, sorted_entries_list in partitioned_files.items():
            if not sorted_entries_list:
                continue

            data_files: List[DataFileMeta] = [e.file for e in sorted_entries_list]

            # Split files by firstRowId for data evolution
            split_by_row_id = self._split_by_row_id(data_files)

            # Pack the split groups for optimal split sizes
            packed_files = self._pack_for_ordered(
                split_by_row_id, weight_func, self.target_split_size
            )

            # Flatten the packed files and build splits
            flatten_packed_files: List[List[DataFileMeta]] = [
                [file for sub_pack in pack for file in sub_pack]
                for pack in packed_files
            ]

            splits += self._build_split_from_pack_for_data_evolution(
                flatten_packed_files, packed_files, sorted_entries_list
            )

        # merge slice_row_ranges and self.row_ranges
        if slice_row_ranges is None:
            slice_row_ranges = self.row_ranges
        elif self.row_ranges is not None:
            slice_row_ranges = Range.and_(slice_row_ranges, self.row_ranges)

        # Wrap splits with IndexedSplit for slice-based filtering or row_ranges
        if slice_row_ranges:
            splits = self._wrap_to_indexed_splits(splits, slice_row_ranges)

        return splits

    def _build_split_from_pack_for_data_evolution(
        self,
        flatten_packed_files: List[List[DataFileMeta]],
        packed_files: List[List[List[DataFileMeta]]],
        file_entries: List[ManifestEntry]
    ) -> List[Split]:
        """
        Build splits from packed files for data evolution tables.
        raw_convertible is True only when each range (pack) contains exactly one file.
        """
        splits = []
        for i, file_group in enumerate(flatten_packed_files):
            # In Java: rawConvertible = f.stream().allMatch(file -> file.size() == 1)
            # This means raw_convertible is True only when each range contains exactly one file
            pack = packed_files[i] if i < len(packed_files) else []
            raw_convertible = all(len(sub_pack) == 1 for sub_pack in pack)

            for data_file in file_group:
                data_file.set_file_path(
                    self.table.table_path,
                    file_entries[0].partition,
                    file_entries[0].bucket,
                    self.default_part_value
                )

            if file_group:
                # Get deletion files for this split
                data_deletion_files = None
                if self.deletion_files_map:
                    data_deletion_files = self._get_deletion_files_for_split(
                        file_group,
                        file_entries[0].partition,
                        file_entries[0].bucket
                    )

                split = DataSplit(
                    files=file_group,
                    partition=file_entries[0].partition,
                    bucket=file_entries[0].bucket,
                    raw_convertible=raw_convertible,
                    data_deletion_files=data_deletion_files
                )
                splits.append(split)
        return splits

    def _calculate_slice_row_ranges(self, partitioned_files: defaultdict) -> List[Range]:
        """
        Calculate Row ID ranges for slice-based filtering based on start_pos and end_pos.
        """
        # Collect all Row ID ranges from files
        list_ranges = []
        for file_entries in partitioned_files.values():
            for entry in file_entries:
                list_ranges.append(entry.file.row_id_range())

        # Merge overlapping ranges
        sorted_ranges = Range.sort_and_merge_overlap(list_ranges, True, False)

        # Calculate the Row ID range for this slice
        start_range, end_range = self._divide_ranges_by_position(sorted_ranges)
        if start_range is None or end_range is None:
            return []

        # Return the range for this slice
        return [Range(start_range.from_, end_range.to)]

    def _divide_ranges_by_position(self, sorted_ranges: List[Range]) -> Tuple[Optional[Range], Optional[Range]]:
        """
        Divide ranges by position (start_pos, end_pos) to get the Row ID range for this slice.
        If idx_of_this_subtask exists, divide total rows by number_of_para_subtasks.
        """
        if not sorted_ranges:
            return None, None

        total_row_count = sum(r.count() for r in sorted_ranges)
        
        # If idx_of_this_subtask exists, calculate start_pos and end_pos based on number_of_para_subtasks
        if self.idx_of_this_subtask is not None:
            # Calculate shard boundaries based on total row count
            rows_per_task = total_row_count // self.number_of_para_subtasks
            remainder = total_row_count % self.number_of_para_subtasks
            
            start_pos = self.idx_of_this_subtask * rows_per_task
            # Distribute remainder rows across first 'remainder' tasks
            if self.idx_of_this_subtask < remainder:
                start_pos += self.idx_of_this_subtask
                end_pos = start_pos + rows_per_task + 1
            else:
                start_pos += remainder
                end_pos = start_pos + rows_per_task
        else:
            # Use existing start_pos and end_pos
            start_pos = self.start_pos_of_this_subtask
            end_pos = self.end_pos_of_this_subtask

        if start_pos >= total_row_count:
            return None, None

        # Find the start Row ID
        current_pos = 0
        start_row_id = None
        end_row_id = None

        for r in sorted_ranges:
            range_end_pos = current_pos + r.count()

            # Find start Row ID
            if start_row_id is None and start_pos < range_end_pos:
                offset = start_pos - current_pos
                start_row_id = r.from_ + offset

            # Find end Row ID
            if end_pos <= range_end_pos:
                offset = end_pos - current_pos
                end_row_id = r.from_ + offset - 1  # -1 because end_pos is exclusive
                break

            current_pos = range_end_pos

        if start_row_id is None:
            return None, None
        if end_row_id is None:
            end_row_id = sorted_ranges[-1].to

        return Range(start_row_id, start_row_id), Range(end_row_id, end_row_id)

    @staticmethod
    def _filter_files_by_row_ranges(partitioned_files: defaultdict, row_ranges: List[Range]) -> defaultdict:
        """
        Filter files by Row ID ranges. Keep files that overlap with the given ranges.
        Blob files are only included if they overlap with non-blob files that match the ranges.
        """
        filtered_partitioned_files = defaultdict(list)

        for key, file_entries in partitioned_files.items():
            # Separate blob and non-blob files
            non_blob_entries = []
            blob_entries = []
            
            for entry in file_entries:
                if DataFileMeta.is_blob_file(entry.file.file_name):
                    blob_entries.append(entry)
                else:
                    non_blob_entries.append(entry)
            
            # First, filter non-blob files based on row ranges
            filtered_non_blob_entries = []
            non_blob_ranges = []
            for entry in non_blob_entries:
                file_range = entry.file.row_id_range()

                # Check if file overlaps with any of the row ranges
                overlaps = False
                for r in row_ranges:
                    if r.overlaps(file_range):
                        overlaps = True
                        break
                
                if overlaps:
                    filtered_non_blob_entries.append(entry)
                    non_blob_ranges.append(file_range)
            
            # Then, filter blob files based on row ID range of non-blob files
            filtered_blob_entries = []
            non_blob_ranges = Range.sort_and_merge_overlap(non_blob_ranges, True, True)
            # Only keep blob files that overlap with merged non-blob ranges
            for entry in blob_entries:
                blob_range = entry.file.row_id_range()
                # Check if blob file overlaps with any merged range
                for merged_range in non_blob_ranges:
                    if merged_range.overlaps(blob_range):
                        filtered_blob_entries.append(entry)
                        break
            
            # Combine filtered non-blob and blob files
            filtered_entries = filtered_non_blob_entries + filtered_blob_entries
            
            if filtered_entries:
                filtered_partitioned_files[key] = filtered_entries

        return filtered_partitioned_files

    @staticmethod
    def _split_by_row_id(files: List[DataFileMeta]) -> List[List[DataFileMeta]]:
        """
        Split files by row ID for data evolution tables.
        Files are grouped by their overlapping row ID ranges.
        """
        list_ranges = [file.row_id_range() for file in files]

        if not list_ranges:
            return []

        sorted_ranges = Range.sort_and_merge_overlap(list_ranges, True, False)

        range_to_files = {}
        for file in files:
            file_range = file.row_id_range()
            for r in sorted_ranges:
                if r.overlaps(file_range):
                    range_to_files.setdefault(r, []).append(file)
                    break

        return list(range_to_files.values())

    def _wrap_to_indexed_splits(self, splits: List[Split], row_ranges: List[Range]) -> List[Split]:
        """
        Wrap splits with IndexedSplit for row range filtering.
        """
        indexed_splits = []
        for split in splits:
            # Calculate file ranges for this split
            file_ranges = [file.row_id_range() for file in split.files]

            if not file_ranges:
                # No row IDs, keep original split
                indexed_splits.append(split)
                continue

            # Merge file ranges
            file_ranges = Range.merge_sorted_as_possible(file_ranges)

            # Intersect with row_ranges from global index
            expected = Range.and_(file_ranges, row_ranges)

            if not expected:
                # No intersection, skip this split
                continue

            # Create scores array if score_getter is provided
            scores = None
            if self.score_getter is not None:
                scores = []
                for r in expected:
                    for row_id in range(r.from_, r.to + 1):
                        score = self.score_getter(row_id)
                        scores.append(score if score is not None else 0.0)

            indexed_splits.append(IndexedSplit(split, expected, scores))

        return indexed_splits
