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

"""End-to-end tests for the ``partial-update`` merge engine.

Each test creates a PK table with ``merge-engine`` set to a particular
value, writes one or more batches, and reads back. Partial-update reads
must merge non-null fields across batches; ``deduplicate`` must keep
the latest row only; ``first-row`` must keep the earliest row.
``aggregation`` has its own engine-specific e2e coverage in
:mod:`test_aggregation_e2e`.
"""

import os
import shutil
import tempfile
import unittest

import pyarrow as pa

from pypaimon import CatalogFactory, Schema


class PartialUpdateMergeEngineE2ETest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.tempdir = tempfile.mkdtemp()
        cls.warehouse = os.path.join(cls.tempdir, 'warehouse')
        cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse})
        cls.catalog.create_database('default', True)

        cls.pa_schema = pa.schema([
            pa.field('id', pa.int64(), nullable=False),
            ('a', pa.string()),
            ('b', pa.string()),
            ('c', pa.string()),
        ])

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.tempdir, ignore_errors=True)

    def _create_pk_table(self, table_name, merge_engine='partial-update',
                         extra_options=None):
        # bucket=1 so all rows for any PK land in the same bucket; this is
        # what forces the read path through SortMergeReader instead of the
        # raw_convertible / single-file fast path. partial-update merging
        # only happens inside SortMergeReader.
        options = {
            'bucket': '1',
            'merge-engine': merge_engine,
        }
        if extra_options:
            options.update(extra_options)
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['id'],
            options=options,
        )
        full = 'default.{}'.format(table_name)
        self.catalog.create_table(full, schema, False)
        return self.catalog.get_table(full)

    def _write(self, table, rows):
        wb = table.new_batch_write_builder()
        w = wb.new_write()
        c = wb.new_commit()
        try:
            w.write_arrow(pa.Table.from_pylist(rows, schema=self.pa_schema))
            c.commit(w.prepare_commit())
        finally:
            w.close()
            c.close()

    def _write_many(self, table, batches):
        """Multiple ``write_arrow`` calls inside a single ``prepare_commit``.

        Mirrors the reviewer's question: rows that land in the same
        underlying data file must still go through the merge-engine
        dispatch; in-writer merging cannot silently degrade to dedupe.
        """
        wb = table.new_batch_write_builder()
        w = wb.new_write()
        c = wb.new_commit()
        try:
            for rows in batches:
                w.write_arrow(pa.Table.from_pylist(rows, schema=self.pa_schema))
            c.commit(w.prepare_commit())
        finally:
            w.close()
            c.close()

    def _read(self, table):
        rb = table.new_read_builder()
        splits = rb.new_scan().plan().splits()
        if not splits:
            return []
        return sorted(
            rb.new_read().to_arrow(splits).to_pylist(),
            key=lambda r: r['id'],
        )

    # -- partial-update happy path ---------------------------------------

    def test_partial_update_two_writes_merges_non_null(self):
        """Two writes against the same PK with disjoint non-null columns
        must merge into a single row that has both columns populated.
        """
        table = self._create_pk_table('two_writes')
        self._write(table, [{'id': 1, 'a': 'A', 'b': None, 'c': None}])
        self._write(table, [{'id': 1, 'a': None, 'b': 'B', 'c': None}])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'A', 'b': 'B', 'c': None}],
        )

    def test_partial_update_three_writes_merges_left_to_right(self):
        """Three overlapping writes — each filling in a different column —
        compose into the union of non-null fields.
        """
        table = self._create_pk_table('three_writes')
        self._write(table, [{'id': 1, 'a': 'A', 'b': None, 'c': None}])
        self._write(table, [{'id': 1, 'a': None, 'b': 'B', 'c': None}])
        self._write(table, [{'id': 1, 'a': None, 'b': None, 'c': 'C'}])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'A', 'b': 'B', 'c': 'C'}],
        )

    def test_partial_update_disjoint_keys_unaffected(self):
        """Three rows with disjoint PKs must all appear unchanged in the
        output — partial-update only merges rows that share a PK.
        """
        table = self._create_pk_table('disjoint_keys')
        self._write(table, [
            {'id': 1, 'a': 'A1', 'b': None, 'c': None},
            {'id': 2, 'a': None, 'b': 'B2', 'c': None},
            {'id': 3, 'a': None, 'b': None, 'c': 'C3'},
        ])

        self.assertEqual(
            self._read(table),
            [
                {'id': 1, 'a': 'A1', 'b': None, 'c': None},
                {'id': 2, 'a': None, 'b': 'B2', 'c': None},
                {'id': 3, 'a': None, 'b': None, 'c': 'C3'},
            ],
        )

    def test_partial_update_later_value_wins_over_earlier_non_null(self):
        """When two writes both supply a non-null value for the same
        column, the later value wins (latest non-null per field).
        """
        table = self._create_pk_table('later_wins')
        self._write(table, [{'id': 1, 'a': 'old', 'b': 'keep', 'c': None}])
        self._write(table, [{'id': 1, 'a': 'new', 'b': None, 'c': 'fill'}])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'new', 'b': 'keep', 'c': 'fill'}],
        )

    def test_partial_update_later_null_does_not_clobber_earlier_value(self):
        """A later write with NULL for a column does NOT overwrite an
        earlier non-null value for that column.
        """
        table = self._create_pk_table('null_no_clobber')
        self._write(table, [{'id': 1, 'a': 'A', 'b': 'B', 'c': 'C'}])
        self._write(table, [{'id': 1, 'a': None, 'b': None, 'c': None}])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'A', 'b': 'B', 'c': 'C'}],
        )

    # -- single-commit, multiple write_arrow calls -----------------------
    #
    # The in-memory merge buffer added to ``KeyValueDataWriter`` runs
    # the merge function on flush, so rows from multiple ``write_arrow``
    # calls that share a primary key are folded into a single row before
    # the data file is written. The flushed file therefore satisfies the
    # LSM "PK unique within a file" invariant the read-side
    # ``raw_convertible`` fast path relies on.

    def test_partial_update_two_write_arrows_single_commit(self):
        """Two ``write_arrow`` calls + one ``prepare_commit``: each
        carries a disjoint non-null field; result is the per-field merge.
        """
        table = self._create_pk_table('two_writes_single_commit')
        self._write_many(table, [
            [{'id': 1, 'a': 'A', 'b': None, 'c': None}],
            [{'id': 1, 'a': None, 'b': 'B', 'c': None}],
        ])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'A', 'b': 'B', 'c': None}],
        )

    def test_partial_update_three_write_arrows_single_commit(self):
        """Three ``write_arrow`` calls in a single commit compose into
        the union of non-null fields.
        """
        table = self._create_pk_table('three_writes_single_commit')
        self._write_many(table, [
            [{'id': 1, 'a': 'A', 'b': None, 'c': None}],
            [{'id': 1, 'a': None, 'b': 'B', 'c': None}],
            [{'id': 1, 'a': None, 'b': None, 'c': 'C'}],
        ])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'A', 'b': 'B', 'c': 'C'}],
        )

    # -- deduplicate (regression) ----------------------------------------

    def test_deduplicate_engine_unchanged(self):
        """The default ``deduplicate`` engine must keep the latest row
        intact, including its NULLs — exactly the pre-PR behaviour.
        """
        table = self._create_pk_table('dedupe', merge_engine='deduplicate')
        self._write(table, [{'id': 1, 'a': 'old', 'b': 'old-b', 'c': 'old-c'}])
        self._write(table, [{'id': 1, 'a': 'new', 'b': None, 'c': None}])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'new', 'b': None, 'c': None}],
        )

    def test_deduplicate_two_write_arrows_single_commit(self):
        """Pre-PR master silently returned both rows because the
        flushed file held two records sharing a primary key. With the
        in-memory merge buffer in place, ``deduplicate`` collapses
        same-PK rows in a single commit too -- LSM "PK unique within a
        file" invariant restored.
        """
        table = self._create_pk_table(
            'dedupe_two_writes_single_commit',
            merge_engine='deduplicate',
        )
        self._write_many(table, [
            [{'id': 1, 'a': 'first', 'b': 'old', 'c': None}],
            [{'id': 1, 'a': 'second', 'b': 'new', 'c': None}],
        ])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'second', 'b': 'new', 'c': None}],
        )

    # -- other supported engines (smoke) ---------------------------------

    def test_first_row_engine_keeps_first(self):
        """The ``first-row`` engine must keep the earliest row per PK.

        Both the writer-side merge buffer and the reader-side merge
        function go through ``merge_engine_dispatch``, so first-row is
        a real supported engine (no dedupe fallback / no NotImplemented
        raise) on both sides.
        """
        table = self._create_pk_table('first_row_supported',
                                      merge_engine='first-row')
        self._write(table, [{'id': 1, 'a': 'first', 'b': None, 'c': None}])
        self._write(table, [{'id': 1, 'a': 'second', 'b': 'B', 'c': 'C'}])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'first', 'b': None, 'c': None}],
        )

    def test_aggregation_engine_write_logs_fallback_warning(self):
        """The write-side fallback to deduplicate for unsupported engines
        is silent in terms of return value -- a ``logging.warning`` is
        the only signal that file contents will not match the table's
        declared semantics. Important when the same table is read back
        by a reader that honours the declared engine; the pypaimon
        read-side raise wouldn't fire there.
        """
        table = self._create_pk_table('agg_warning',
                                      merge_engine='aggregation')
        with self.assertLogs(
                'pypaimon.write.file_store_write', level='WARNING') as cm:
            self._write(table, [{'id': 1, 'a': 'x', 'b': None, 'c': None}])
        combined = '\n'.join(cm.output)
        self.assertIn('aggregation', combined)
        self.assertIn('deduplicate', combined)

    # -- partial-update + out-of-scope option combinations ---------------
    #
    # When a user pairs ``merge-engine: partial-update`` with any option
    # this port doesn't implement (sequence-group, per-field aggregator
    # override, ignore-delete, partial-update.remove-record-on-*), we
    # must raise rather than silently run the simple last-non-null merge
    # — otherwise we'd reproduce the same silent-corruption pattern this
    # PR exists to close.

    def _assert_partial_update_unsupported(self, table_name, extra_options,
                                           expected_keys):
        # Shared dispatch runs at write time too, so the unsupported-
        # option error surfaces inside the first ``write_arrow`` call
        # (when ``FileStoreWrite._create_data_writer`` first runs)
        # rather than waiting for read.
        table = self._create_pk_table(
            table_name, extra_options=extra_options)
        with self.assertRaises(NotImplementedError) as cm:
            self._write(table, [{'id': 1, 'a': 'A', 'b': None, 'c': None}])
        msg = str(cm.exception)
        self.assertIn("partial-update", msg)
        for key in expected_keys:
            self.assertIn(key, msg,
                          "expected option key '{}' in error: {}".format(key, msg))

    def test_partial_update_with_sequence_group_raises(self):
        self._assert_partial_update_unsupported(
            'pu_seq_group',
            {'fields.b.sequence-group': 'a'},
            ['fields.b.sequence-group'],
        )

    def test_partial_update_with_field_aggregate_function_raises(self):
        self._assert_partial_update_unsupported(
            'pu_field_agg',
            {'fields.a.aggregate-function': 'last_non_null_value'},
            ['fields.a.aggregate-function'],
        )

    def test_partial_update_with_default_aggregate_function_raises(self):
        self._assert_partial_update_unsupported(
            'pu_default_agg',
            {'fields.default-aggregate-function': 'last_non_null_value'},
            ['fields.default-aggregate-function'],
        )

    def test_partial_update_with_ignore_delete_raises(self):
        self._assert_partial_update_unsupported(
            'pu_ignore_delete',
            {'ignore-delete': 'true'},
            ['ignore-delete'],
        )

    def test_partial_update_with_remove_record_on_delete_raises(self):
        self._assert_partial_update_unsupported(
            'pu_rrod',
            {'partial-update.remove-record-on-delete': 'true'},
            ['partial-update.remove-record-on-delete'],
        )

    def test_partial_update_with_remove_record_on_sequence_group_raises(self):
        self._assert_partial_update_unsupported(
            'pu_rrosg',
            {'partial-update.remove-record-on-sequence-group': 'true'},
            ['partial-update.remove-record-on-sequence-group'],
        )

    def test_partial_update_unsupported_options_guard_covers_raw_convertible(self):
        """The read-side guard at ``TableRead.__init__`` must fire even
        when the scan would dispatch every split through
        ``RawFileSplitRead`` (single-snapshot, non-overlapping rows).

        Before the guard moved to ``TableRead.__init__`` this case
        silently bypassed validation because raw-convertible splits skip
        ``MergeFileSplitRead`` entirely -- the read path's
        ``_build_merge_function`` never ran, so an option like
        ``partial-update.remove-record-on-delete`` was ignored on read.

        The shared dispatch now also fires on the write path's first
        flush (see ``_assert_partial_update_unsupported``), so we skip
        ``_write`` here: the read-side guard runs at ``new_read()``
        construction time regardless of whether data exists.
        """
        table = self._create_pk_table(
            'pu_rrod_raw_convertible',
            extra_options={'partial-update.remove-record-on-delete': 'true'},
        )
        rb = table.new_read_builder()
        with self.assertRaises(NotImplementedError) as cm:
            rb.new_read()
        msg = str(cm.exception)
        self.assertIn('partial-update', msg)
        self.assertIn('partial-update.remove-record-on-delete', msg)

    def test_partial_update_with_explicit_ignore_delete_false_does_not_raise(self):
        """Explicitly setting ignore-delete=false is equivalent to leaving
        it unset and must not trip the guard."""
        table = self._create_pk_table(
            'pu_ignore_delete_false',
            extra_options={'ignore-delete': 'false'},
        )
        self._write(table, [{'id': 1, 'a': 'A', 'b': None, 'c': None}])
        self._write(table, [{'id': 1, 'a': None, 'b': 'B', 'c': None}])

        self.assertEqual(
            self._read(table),
            [{'id': 1, 'a': 'A', 'b': 'B', 'c': None}],
        )


if __name__ == '__main__':
    unittest.main()
