# 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 dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa

from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.manifest.schema.manifest_entry import ManifestEntry
from pypaimon.read.split import DataSplit
from pypaimon.read.table_read import TableRead
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.table.special_fields import SpecialFields
from pypaimon.utils.range import Range
from pypaimon.write.commit_message import CommitMessage
from pypaimon.write.table_update_by_row_id import TableUpdateByRowId


@dataclass
class RowIdRewriteResult:
    commit_entries: List[ManifestEntry]
    rewritten_file_count: int


class RowIdConflictRewriter:
    """Rebase stale partial-column files onto current row-id file ranges."""

    def __init__(
            self,
            table,
            commit_user: str,
            commit_identifier: int,
            max_rewrite_size: int):
        self.table = table
        self.commit_user = commit_user
        self.commit_identifier = commit_identifier
        self.max_rewrite_size = max_rewrite_size

    def rewrite(
            self,
            latest_snapshot,
            base_entries: List[ManifestEntry],
            delta_entries: List[ManifestEntry],
    ) -> Optional[RowIdRewriteResult]:
        if self.max_rewrite_size <= 0:
            return None
        if self.table.options.deletion_vectors_enabled(False):
            return None

        current_exact_ranges = {
            self._range_key(entry)
            for entry in base_entries
            if self._is_normal_row_id_file(entry.file)
        }
        candidates = [
            entry
            for entry in delta_entries
            if self._is_rewrite_candidate(
                entry, current_exact_ranges, latest_snapshot.next_row_id)
        ]
        if not candidates:
            return None

        candidate_ids = {id(entry) for entry in candidates}
        if any(
                self._is_dedicated_file(entry.file)
                and entry.file.first_row_id is not None
                and entry.file.first_row_id < latest_snapshot.next_row_id
                for entry in delta_entries
        ):
            return None

        if not self._ranges_are_still_covered(base_entries, candidates):
            return None

        affected_files = self._affected_current_files(base_entries, candidates)
        affected_size = sum(file.file_size for file in affected_files)
        if affected_size > self.max_rewrite_size:
            raise RuntimeError(
                "Automatic row-id conflict rewrite requires reading "
                "{} bytes across {} current data file(s), exceeding "
                "'data-evolution.row-id-conflict-rewrite.max-size'={} bytes."
                .format(
                    affected_size,
                    len(affected_files),
                    self.max_rewrite_size,
                )
            )

        new_messages = []
        try:
            files_info = TableUpdateByRowId._files_info_from_entries(
                self.table,
                latest_snapshot.id,
                base_entries,
            )
            groups = self._group_by_write_columns(candidates)
            for column_names, entries in groups:
                update_tables = [
                    self._read_staged_update(entry, column_names)
                    for entry in entries
                ]
                update_data = (
                    update_tables[0]
                    if len(update_tables) == 1
                    else pa.concat_tables(update_tables)
                )
                row_ids = update_data[SpecialFields.ROW_ID.name].to_pylist()
                if len(row_ids) != len(set(row_ids)):
                    raise RuntimeError(
                        "Automatic row-id conflict rewrite cannot merge "
                        "overlapping staged files for columns {}.".format(
                            column_names)
                    )
                updater = TableUpdateByRowId(
                    self.table,
                    self.commit_user,
                    self.commit_identifier,
                    _precomputed_files_info=files_info,
                )
                try:
                    new_messages.extend(
                        updater.update_columns(update_data, column_names)
                    )
                except Exception:
                    self._abort(updater.commit_messages)
                    raise
        except Exception:
            self._abort(new_messages)
            raise

        rewritten_entries = [
            entry for entry in delta_entries if id(entry) not in candidate_ids
        ]
        rewritten_entries.extend(self._to_manifest_entries(new_messages))
        return RowIdRewriteResult(
            commit_entries=rewritten_entries,
            rewritten_file_count=len(candidates),
        )

    def _is_rewrite_candidate(
            self, entry, current_exact_ranges, next_row_id):
        file = entry.file
        return (
            self._is_normal_row_id_file(file)
            and file.first_row_id < next_row_id
            and bool(file.write_cols)
            and not any(
                SpecialFields.is_system_field(name)
                for name in file.write_cols
            )
            and self._range_key(entry) not in current_exact_ranges
        )

    def _ranges_are_still_covered(self, base_entries, candidates):
        current_ranges = {}
        for entry in base_entries:
            if not self._is_normal_row_id_file(entry.file):
                continue
            key = (tuple(entry.partition.values), entry.bucket)
            current_ranges.setdefault(key, []).append(
                entry.file.row_id_range())
        current_ranges = {
            key: Range.sort_and_merge_overlap(ranges, True, True)
            for key, ranges in current_ranges.items()
        }

        for entry in candidates:
            key = (tuple(entry.partition.values), entry.bucket)
            if entry.file.row_id_range().exclude(
                    current_ranges.get(key, [])):
                return False
        return True

    def _affected_current_files(self, base_entries, candidates):
        affected = {}
        for base in base_entries:
            if not self._is_normal_row_id_file(base.file):
                continue
            base_key = (tuple(base.partition.values), base.bucket)
            for candidate in candidates:
                candidate_key = (
                    tuple(candidate.partition.values), candidate.bucket)
                if (
                    base_key == candidate_key
                    and base.file.row_id_range().overlaps(
                        candidate.file.row_id_range())
                ):
                    key = (
                        base_key,
                        base.file.external_path or base.file.file_path
                        or base.file.file_name,
                    )
                    affected[key] = base.file
                    break
        return list(affected.values())

    @staticmethod
    def _group_by_write_columns(candidates):
        groups = {}
        for entry in candidates:
            key = tuple(entry.file.write_cols)
            groups.setdefault(key, []).append(entry)
        return [
            (list(columns), entries)
            for columns, entries in groups.items()
        ]

    def _read_staged_update(self, entry, column_names):
        read_fields = [self.table.field_dict[name] for name in column_names]
        read_fields.append(SpecialFields.ROW_ID)
        split = DataSplit(
            files=[entry.file],
            partition=entry.partition,
            bucket=entry.bucket,
            raw_convertible=True,
        )
        result = TableRead(
            self.table,
            predicate=None,
            read_type=read_fields,
        ).to_arrow([split])
        if result.num_rows != entry.file.row_count:
            raise RuntimeError(
                "Automatic row-id conflict rewrite read {} rows from staged "
                "file '{}', expected {}.".format(
                    result.num_rows,
                    entry.file.file_name,
                    entry.file.row_count,
                )
            )
        return result

    def _to_manifest_entries(self, messages: List[CommitMessage]):
        entries = []
        for message in messages:
            partition = GenericRow(
                list(message.partition),
                self.table.partition_keys_fields,
            )
            for file in message.new_files:
                entries.append(ManifestEntry(
                    kind=0,
                    partition=partition,
                    bucket=message.bucket,
                    total_buckets=self.table.total_buckets,
                    file=file,
                ))
        return entries

    def _abort(self, messages):
        for message in messages:
            for file in message.new_files:
                path = file.external_path or file.file_path
                if path:
                    self.table.file_io.delete_quietly(path)

    @staticmethod
    def _range_key(entry):
        return (
            tuple(entry.partition.values),
            entry.bucket,
            entry.file.first_row_id,
            entry.file.row_count,
        )

    @staticmethod
    def _is_normal_row_id_file(file):
        return (
            file.first_row_id is not None
            and not RowIdConflictRewriter._is_dedicated_file(file)
        )

    @staticmethod
    def _is_dedicated_file(file):
        return (
            DataFileMeta.is_blob_file(file.file_name)
            or DataFileMeta.is_vector_file(file.file_name)
        )
