# 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 abc import ABC, abstractmethod
from typing import Callable, List, Optional, Dict, Tuple

from pypaimon.common.options.core_options import CoreOptions
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.manifest.schema.manifest_entry import ManifestEntry
from pypaimon.read.split import Split
from pypaimon.read.split import DataSplit
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.table.source.deletion_file import DeletionFile


class AbstractSplitGenerator(ABC):
    """
    Abstract base class for generating splits.
    """
    
    # Special key for tracking file end position in split file index map
    NEXT_POS_KEY = '_next_pos'

    def __init__(
        self,
        table,
        target_split_size: int,
        open_file_cost: int,
        deletion_files_map: Optional[Dict] = None
    ):
        self.table = table
        self.target_split_size = target_split_size
        self.open_file_cost = open_file_cost
        self.deletion_files_map = deletion_files_map or {}
        self.default_part_value = table.options.options.get(
            CoreOptions.PARTITION_DEFAULT_NAME, "__DEFAULT_PARTITION__")
        
        # Shard configuration
        self.idx_of_this_subtask = None
        self.number_of_para_subtasks = None
        self.start_pos_of_this_subtask = None
        self.end_pos_of_this_subtask = None

    def with_shard(self, idx_of_this_subtask: int, number_of_para_subtasks: int):
        """Configure sharding for parallel processing."""
        if idx_of_this_subtask >= number_of_para_subtasks:
            raise ValueError("idx_of_this_subtask must be less than number_of_para_subtasks")
        if self.start_pos_of_this_subtask is not None:
            raise ValueError("with_shard and with_slice cannot be used simultaneously")
        self.idx_of_this_subtask = idx_of_this_subtask
        self.number_of_para_subtasks = number_of_para_subtasks
        return self

    def with_slice(self, start_pos: int, end_pos: int):
        """Configure slice range for processing."""
        if start_pos >= end_pos:
            raise ValueError("start_pos must be less than end_pos")
        if self.idx_of_this_subtask is not None:
            raise ValueError("with_slice and with_shard cannot be used simultaneously")
        self.start_pos_of_this_subtask = start_pos
        self.end_pos_of_this_subtask = end_pos
        return self

    @abstractmethod
    def create_splits(self, file_entries: List[ManifestEntry]) -> List[Split]:
        """
        Create splits from manifest entries.
        """
        pass

    def _build_split_from_pack(
        self,
        packed_files: List[List[DataFileMeta]],
        file_entries: List[ManifestEntry],
        for_primary_key_split: bool,
        use_optimized_path: bool = False
    ) -> List[Split]:
        """
        Build splits from packed files.
        """
        splits = []
        for file_group in packed_files:
            if use_optimized_path:
                raw_convertible = True
            elif for_primary_key_split:
                raw_convertible = len(file_group) == 1 and self._without_delete_row(file_group[0])
            else:
                raw_convertible = True

            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 _get_deletion_files_for_split(
        self,
        data_files: List[DataFileMeta],
        partition: GenericRow,
        bucket: int
    ) -> Optional[List[DeletionFile]]:
        """Get deletion files for the given data files in a split."""
        if not self.deletion_files_map:
            return None

        partition_key = (tuple(partition.values), bucket)
        file_deletion_map = self.deletion_files_map.get(partition_key, {})

        if not file_deletion_map:
            return None

        deletion_files = []
        for data_file in data_files:
            deletion_file = file_deletion_map.get(data_file.file_name)
            if deletion_file:
                deletion_files.append(deletion_file)
            else:
                deletion_files.append(None)

        return deletion_files if any(df is not None for df in deletion_files) else None

    @staticmethod
    def _without_delete_row(data_file_meta: DataFileMeta) -> bool:
        """Check if a data file has no deleted rows."""
        if data_file_meta.delete_row_count is None:
            return True
        return data_file_meta.delete_row_count == 0

    @staticmethod
    def _pack_for_ordered(
        items: List,
        weight_func: Callable,
        target_weight: int
    ) -> List[List]:
        """Pack items into groups based on target weight."""
        packed = []
        bin_items = []
        bin_weight = 0

        for item in items:
            weight = weight_func(item)
            if bin_weight + weight > target_weight and len(bin_items) > 0:
                packed.append(list(bin_items))
                bin_items.clear()
                bin_weight = 0

            bin_weight += weight
            bin_items.append(item)

        if len(bin_items) > 0:
            packed.append(bin_items)

        return packed

    def _compute_shard_range(self, total_row: int) -> Tuple[int, int]:
        """
        Calculate start and end positions for this shard based on total rows.
        Uses balanced distribution to avoid last shard overload.
        """
        base_rows_per_shard = total_row // self.number_of_para_subtasks
        remainder = total_row % self.number_of_para_subtasks

        # Each of the first 'remainder' shards gets one extra row
        if self.idx_of_this_subtask < remainder:
            num_row = base_rows_per_shard + 1
            start_pos = self.idx_of_this_subtask * (base_rows_per_shard + 1)
        else:
            num_row = base_rows_per_shard
            start_pos = (
                remainder * (base_rows_per_shard + 1) +
                (self.idx_of_this_subtask - remainder) * base_rows_per_shard
            )

        end_pos = start_pos + num_row
        return start_pos, end_pos

    @staticmethod
    def _compute_file_range(
        plan_start_pos: int,
        plan_end_pos: int,
        file_begin_pos: int,
        file_row_count: int
    ) -> Optional[Tuple[int, int]]:
        """
        Compute the row range to read from a file given shard range and file position.
        Returns None if file is completely within shard range (no slicing needed).
        Returns (-1, -1) if file is completely outside shard range.
        """
        file_end_pos = file_begin_pos + file_row_count

        if file_begin_pos <= plan_start_pos < plan_end_pos <= file_end_pos:
            return plan_start_pos - file_begin_pos, plan_end_pos - file_begin_pos
        elif file_begin_pos < plan_start_pos < file_end_pos:
            return plan_start_pos - file_begin_pos, file_row_count
        elif file_begin_pos < plan_end_pos < file_end_pos:
            return 0, plan_end_pos - file_begin_pos
        elif file_end_pos <= plan_start_pos or file_begin_pos >= plan_end_pos:
            return -1, -1
        # File is completely within the shard range
        return None
