################################################################################
#  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.
################################################################################
import bisect
from typing import Dict, List, Optional

import pyarrow as pa
import pyarrow.compute as pc

from pypaimon.read.split import DataSplit
from pypaimon.read.table_read import TableRead
from pypaimon.schema.data_types import DataField
from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.table.special_fields import SpecialFields
from pypaimon.write.file_store_write import FileStoreWrite


class TableUpdateByRowId:
    """
    Table update for partial column updates (data evolution).

    This update is designed for adding/updating specific columns in existing tables.
    Input data should contain _ROW_ID column.
    """

    FIRST_ROW_ID_COLUMN = '_FIRST_ROW_ID'

    def __init__(self, table, commit_user: str):
        from pypaimon.table.file_store_table import FileStoreTable

        self.table: FileStoreTable = table
        self.commit_user = commit_user

        # Load existing first_row_ids and build partition map
        (self.snapshot_id,
         self.first_row_ids,
         self.first_row_id_to_partition_map,
         self.first_row_id_to_row_count_map,
         self.total_row_count,
         self.splits) = self._load_existing_files_info()

        # Collect commit messages
        self.commit_messages = []

    def _load_existing_files_info(self):
        """Load existing first_row_ids and build partition map for efficient lookup."""
        first_row_ids = []
        first_row_id_to_partition_map: Dict[int, GenericRow] = {}
        first_row_id_to_row_count_map: Dict[int, int] = {}

        read_builder = self.table.new_read_builder()
        scan = read_builder.new_scan()
        plan = scan.plan()
        splits = plan.splits()

        for split in splits:
            for file in split.files:
                if file.first_row_id is not None and not file.file_name.endswith('.blob'):
                    first_row_id = file.first_row_id
                    first_row_ids.append(first_row_id)
                    first_row_id_to_partition_map[first_row_id] = split.partition
                    first_row_id_to_row_count_map[first_row_id] = file.row_count

        total_row_count = sum(first_row_id_to_row_count_map.values())

        snapshot_id = plan.snapshot_id if plan.snapshot_id is not None else -1
        return (snapshot_id,
                sorted(list(set(first_row_ids))),
                first_row_id_to_partition_map,
                first_row_id_to_row_count_map,
                total_row_count,
                splits)

    def update_columns(self, data: pa.Table, column_names: List[str]) -> List:
        """
        Add or update columns in the table.

        Args:
            data: Input data containing row_id and columns to update
            column_names: Names of columns to update (excluding row_id)

        Returns:
            List of commit messages
        """

        # Validate column_names is not empty
        if not column_names:
            raise ValueError("column_names cannot be empty")

        # Validate input data has row_id column
        if SpecialFields.ROW_ID.name not in data.column_names:
            raise ValueError(f"Input data must contain {SpecialFields.ROW_ID.name} column")

        # Validate all update columns exist in the schema
        for col_name in column_names:
            if col_name not in self.table.field_names:
                raise ValueError(f"Column {col_name} not found in table schema")

        # Sort data by _ROW_ID column
        sorted_data = data.sort_by([(SpecialFields.ROW_ID.name, "ascending")])

        # Calculate first_row_id for each row
        data_with_first_row_id = self._calculate_first_row_id(sorted_data)

        # Group by first_row_id and write each group
        self._write_by_first_row_id(data_with_first_row_id, column_names)

        return self.commit_messages

    def _calculate_first_row_id(self, data: pa.Table) -> pa.Table:
        """Calculate _first_row_id for each row based on _ROW_ID.

        Supports partial row updates - row_ids don't need to be consecutive.
        """
        row_ids = data[SpecialFields.ROW_ID.name].to_pylist()

        # Validate row_ids have no duplicates
        if len(row_ids) != len(set(row_ids)):
            raise ValueError("Input data contains duplicate _ROW_ID values")

        # Validate row_ids are within valid range
        for row_id in row_ids:
            if row_id < 0 or row_id >= self.total_row_count:
                raise ValueError(f"Row ID {row_id} is out of valid range [0, {self.total_row_count})")

        # Calculate first_row_id for each row_id
        first_row_id_values = []
        for row_id in row_ids:
            first_row_id = self._floor_binary_search(self.first_row_ids, row_id)
            first_row_id_values.append(first_row_id)

        # Add first_row_id column to the table
        first_row_id_array = pa.array(first_row_id_values, type=pa.int64())
        return data.append_column(self.FIRST_ROW_ID_COLUMN, first_row_id_array)

    def _floor_binary_search(self, sorted_seq: List[int], value: int) -> int:
        """Binary search to find the floor value in sorted sequence."""
        if not sorted_seq:
            raise ValueError("The input sorted sequence is empty.")

        idx = bisect.bisect_right(sorted_seq, value) - 1
        if idx < 0:
            raise ValueError(f"Value {value} is less than the first element in the sorted sequence.")

        return sorted_seq[idx]

    def _write_by_first_row_id(self, data: pa.Table, column_names: List[str]):
        """Write data grouped by first_row_id."""
        # Extract unique first_row_id values
        first_row_id_array = data[self.FIRST_ROW_ID_COLUMN]
        unique_first_row_ids = pc.unique(first_row_id_array).to_pylist()

        for first_row_id in unique_first_row_ids:
            # Filter rows for this first_row_id
            mask = pc.equal(first_row_id_array, first_row_id)
            group_data = data.filter(mask)

            # Get partition for this first_row_id
            partition = self._find_partition_by_first_row_id(first_row_id)

            if partition is None:
                raise ValueError(f"No existing file found for first_row_id {first_row_id}")

            # Write this group
            self._write_group(partition, first_row_id, group_data, column_names)

    def _find_partition_by_first_row_id(self, first_row_id: int) -> Optional[GenericRow]:
        """Find the partition for a given first_row_id using pre-built partition map."""
        return self.first_row_id_to_partition_map.get(first_row_id)

    def _read_original_file_data(self, first_row_id: int, column_names: List[str]) -> Optional[pa.Table]:
        """Read original file data for the given first_row_id.

        Only reads columns that exist in the original file and need to be updated.
        In Data Evolution mode, uses the table's read API to get the latest data
        for the specified columns, which handles merging multiple files correctly.

        Args:
            first_row_id: The first_row_id of the file to read
            column_names: The column names to update

        Returns:
            PyArrow Table containing the original data for columns that exist in the file,
            or None if no columns need to be read from the original file.
        """

        # Build read type for the columns we need to read
        read_fields: List[DataField] = []
        for field in self.table.fields:
            if field.name in column_names:
                read_fields.append(field)

        if not read_fields:
            return None

        # Find the split that contains files with this first_row_id
        target_split = None
        target_files = []
        for split in self.splits:
            for file_idx, file in enumerate(split.files):
                if file.first_row_id == first_row_id:
                    target_files.append(file)
                    if target_split is None:
                        target_split = split
            if target_split is not None:
                break

        if not target_files:
            raise ValueError(f"No file found for first_row_id {first_row_id}")

        # Create a DataSplit containing all files with this first_row_id
        origin_split = DataSplit(
            files=target_files,
            partition=target_split.partition,
            bucket=target_split.bucket,
            raw_convertible=True
        )

        # Create TableRead and read the data
        table_read = TableRead(self.table, predicate=None, read_type=read_fields)
        origin_data = table_read.to_arrow([origin_split])

        return origin_data

    def _merge_update_with_original(self, original_data: Optional[pa.Table], update_data: pa.Table,
                                    column_names: List[str], first_row_id: int) -> pa.Table:
        """Merge update data with original data, preserving row order.

        For rows that have updates, use the update values.
        For rows without updates, use the original values (if available).

        Args:
            original_data: Original data from the file (may be None if no columns need to be read)
            update_data: Update data (may contain only partial rows)
            column_names: Column names being updated
            first_row_id: The first_row_id of this file group

        Returns:
            Merged PyArrow Table with all rows
        """

        # Get the _ROW_ID values from update_data to determine which rows are being updated
        relative_indices = pc.subtract(
            update_data[SpecialFields.ROW_ID.name],
            pa.scalar(first_row_id, type=pa.int64())
        ).cast(pa.int64())

        # Build a boolean mask: True at positions that need to be updated
        all_indices = pa.array(range(original_data.num_rows), type=pa.int64())
        mask = pc.is_in(all_indices, relative_indices)

        # Build the merged table column by column
        merged_columns = {}
        for col_name in column_names:
            update_col = update_data[col_name].combine_chunks()
            original_col = original_data[col_name].combine_chunks()
            # replace_with_mask fills mask=True positions with update values in order
            merged_columns[col_name] = pc.replace_with_mask(
                original_col, mask, update_col.cast(original_col.type)
            )

        # Create the merged table
        merged_table = pa.table(merged_columns)

        return merged_table

    def _write_group(self, partition: GenericRow, first_row_id: int,
                     data: pa.Table, column_names: List[str]):
        """Write a group of data with the same first_row_id.

        This method reads the original file data, merges it with the update data,
        and writes out the complete merged data. Supports partial row updates.
        """

        # Read original file data for the columns being updated
        original_data = self._read_original_file_data(first_row_id, column_names)

        # Merge update data with original data
        merged_data = self._merge_update_with_original(original_data, data, column_names, first_row_id)

        # Create a file store write for this partition
        # Disable rolling to ensure one output file per first_row_id group,
        file_store_write = FileStoreWrite(self.table, self.commit_user)
        file_store_write.disable_rolling()

        # Set write columns to only update specific columns
        write_cols = column_names
        file_store_write.write_cols = write_cols

        # Convert partition to tuple for hashing
        partition_tuple = tuple(partition.values)

        # Write merged data - convert Table to RecordBatch
        for batch in merged_data.to_batches():
            file_store_write.write(partition_tuple, 0, batch)

        # Prepare commit and assign first_row_id
        commit_messages = file_store_write.prepare_commit(BATCH_COMMIT_IDENTIFIER)

        # Assign first_row_id to the new files
        for msg in commit_messages:
            msg.check_from_snapshot = self.snapshot_id
            for file in msg.new_files:
                file.first_row_id = first_row_id
                file.write_cols = write_cols

        self.commit_messages.extend(commit_messages)

        # Close the writer
        file_store_write.close()
