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

"""Regression test for caching of OVERWRITE changes across commit retries.

On every OVERWRITE retry, pypaimon re-scanned the full target partitions to
recompute the files to delete. Under concurrent writers this made each retry as
expensive as the first attempt. OverwriteChangesProvider caches the existing
files of the target partitions and, on retry, reuses them when the snapshots in
between can be applied from target-partition DELTA manifests, instead of a full
re-scan.

The test deterministically forces ``K`` conflicts, each advancing the latest
snapshot with an append to an unrelated partition, and asserts the full scan
runs once (not ``K + 1``) and the cache is advanced by a delta probe per retry.
"""

import os
import shutil
import tempfile
import unittest

import pandas as pd
import pyarrow as pa

from pypaimon import CatalogFactory, Schema
from pypaimon.write.commit.overwrite_changes_provider import OverwriteChangesProvider


class OverwriteChangesCacheTest(unittest.TestCase):

    def setUp(self):
        self.temp_dir = tempfile.mkdtemp(prefix="ow_cache_")
        self.warehouse = os.path.join(self.temp_dir, 'wh')
        self.catalog = CatalogFactory.create({"warehouse": self.warehouse})
        self.catalog.create_database("test_db", True)

        pa_schema = pa.schema([('f0', pa.int32()), ('f1', pa.string())])
        schema = Schema.from_pyarrow_schema(
            pa_schema,
            partition_keys=['f0'],
            options={'dynamic-partition-overwrite': 'false'},
        )
        self.catalog.create_table('test_db.t', schema, False)
        self.table = self.catalog.get_table('test_db.t')

        # Seed: f0=1 is the overwrite target, f0=2 untouched.
        self._append(pd.DataFrame({'f0': [1, 1, 2], 'f1': ['a', 'b', 'c']}))

    def tearDown(self):
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def _append(self, df):
        wb = self.table.new_batch_write_builder()
        w = wb.new_write()
        c = wb.new_commit()
        w.write_pandas(df)
        c.commit(w.prepare_commit())
        w.close()
        c.close()

    def test_overwrite_scan_runs_once_not_once_per_retry(self):
        K = 3  # force 3 conflicts; the 4th attempt wins

        wb = self.table.new_batch_write_builder().overwrite({'f0': 1})
        w = wb.new_write()
        c = wb.new_commit()
        w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']}))
        messages = w.prepare_commit()

        fsc = c.file_store_commit

        # --- count provider full scans and delta probes (class-level spies) ---
        counts = {'full_scan': 0, 'probe': 0}
        orig_full = OverwriteChangesProvider._full_scan_manifest_entries
        orig_probe = OverwriteChangesProvider._read_delta_manifest_entries

        def spy_full(self, *a, **k):
            counts['full_scan'] += 1
            return orig_full(self, *a, **k)

        def spy_probe(self, *a, **k):
            counts['probe'] += 1
            return orig_probe(self, *a, **k)

        # --- inject K CAS conflicts; each advances latest via an APPEND to an
        #     unrelated partition (f0=99) so the delta probe finds the target
        #     (f0=1) untouched and the cache is reused.
        orig_cas = fsc.snapshot_commit.commit
        cas = {'fails': 0}

        def patched_cas(base_snapshot_uuid, snapshot, statistics):
            if snapshot.commit_kind == "OVERWRITE" and cas['fails'] < K:
                cas['fails'] += 1
                self._append(pd.DataFrame({'f0': [99], 'f1': [f'x{cas["fails"]}']}))
                return False
            return orig_cas(base_snapshot_uuid, snapshot, statistics)

        fsc.snapshot_commit.commit = patched_cas
        OverwriteChangesProvider._full_scan_manifest_entries = spy_full
        OverwriteChangesProvider._read_delta_manifest_entries = spy_probe
        try:
            c.commit(messages)
            c.close()
        finally:
            OverwriteChangesProvider._full_scan_manifest_entries = orig_full
            OverwriteChangesProvider._read_delta_manifest_entries = orig_probe

        # Harness sanity: we really did force K conflicts and then converged.
        self.assertEqual(cas['fails'], K, "expected exactly K forced conflicts")

        print(f"\n[overwrite-cache] K={K} conflicts -> "
              f"full_scan={counts['full_scan']}, probe={counts['probe']} "
              f"(target: full_scan=1, probe=K)")

        # #7894: the full overwrite scan runs once; each retry reuses the cache
        # after a cheap delta probe of the (unrelated) intervening append.
        self.assertEqual(
            counts['full_scan'], 1,
            f"full overwrite scan ran {counts['full_scan']}x; should run once and "
            f"reuse the cache on retries")
        self.assertEqual(
            counts['probe'], K,
            f"delta probe ran {counts['probe']}x; should probe once per retry "
            f"(= K = {K})")

        # Sanity: f0=1 overwritten, f0=2 preserved, f0=99 appended K times.
        read_builder = self.table.new_read_builder()
        actual = read_builder.new_read().to_pandas(
            read_builder.new_scan().plan().splits())
        self.assertEqual(sorted(actual[actual['f0'] == 1]['f1'].tolist()), ['new'])
        self.assertEqual(sorted(actual[actual['f0'] == 2]['f1'].tolist()), ['c'])
        self.assertEqual(len(actual[actual['f0'] == 99]), K)

    def test_cache_applies_delta_when_concurrent_append_hits_target_partition(self):
        # Concurrent appends hit the overwrite target; retry advances cached
        # target-partition state from APPEND deltas instead of rebuilding.
        K = 3

        wb = self.table.new_batch_write_builder().overwrite({'f0': 1})
        w = wb.new_write()
        c = wb.new_commit()
        w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']}))
        messages = w.prepare_commit()

        fsc = c.file_store_commit

        counts = {'full_scan': 0, 'probe': 0}
        orig_full = OverwriteChangesProvider._full_scan_manifest_entries
        orig_probe = OverwriteChangesProvider._read_delta_manifest_entries

        def spy_full(self, *a, **k):
            counts['full_scan'] += 1
            return orig_full(self, *a, **k)

        def spy_probe(self, *a, **k):
            counts['probe'] += 1
            return orig_probe(self, *a, **k)

        orig_cas = fsc.snapshot_commit.commit
        cas = {'fails': 0}

        def patched_cas(base_snapshot_uuid, snapshot, statistics):
            if snapshot.commit_kind == "OVERWRITE" and cas['fails'] < K:
                cas['fails'] += 1
                self._append(pd.DataFrame({'f0': [1], 'f1': [f'y{cas["fails"]}']}))
                return False
            return orig_cas(base_snapshot_uuid, snapshot, statistics)

        fsc.snapshot_commit.commit = patched_cas
        OverwriteChangesProvider._full_scan_manifest_entries = spy_full
        OverwriteChangesProvider._read_delta_manifest_entries = spy_probe
        try:
            c.commit(messages)
            c.close()
        finally:
            OverwriteChangesProvider._full_scan_manifest_entries = orig_full
            OverwriteChangesProvider._read_delta_manifest_entries = orig_probe

        self.assertEqual(cas['fails'], K, "expected exactly K forced conflicts")

        # Target APPEND deltas can be applied to cached state; no full rebuild.
        self.assertEqual(counts['full_scan'], 1,
                         f"full scan ran {counts['full_scan']}x; APPEND deltas "
                         f"should advance cached target state")
        self.assertEqual(counts['probe'], K,
                         f"delta probe ran {counts['probe']}x; once per retry (= K)")

        # Overwrite wins: f0=1 is just 'new', f0=2 untouched.
        read_builder = self.table.new_read_builder()
        actual = read_builder.new_read().to_pandas(
            read_builder.new_scan().plan().splits())
        self.assertEqual(sorted(actual[actual['f0'] == 1]['f1'].tolist()), ['new'])
        self.assertEqual(sorted(actual[actual['f0'] == 2]['f1'].tolist()), ['c'])

    def _overwrite_partition(self, part_val, f1_val):
        wb = self.table.new_batch_write_builder().overwrite({'f0': part_val})
        w = wb.new_write()
        c = wb.new_commit()
        w.write_pandas(pd.DataFrame({'f0': [part_val], 'f1': [f1_val]}))
        c.commit(w.prepare_commit())
        w.close()
        c.close()

    def _run_with_conflicts(self, c, messages, K, concurrent_fn):
        # Run an overwrite commit, forcing K CAS conflicts (each calls
        # concurrent_fn(i) to advance the latest snapshot). Returns this commit's
        # own OverwriteChangesProvider so the caller can read its counters
        # (only this provider, not any concurrent writer's).
        fsc = c.file_store_commit
        captured = {}
        orig_factory = fsc._overwrite_changes_provider

        def capturing_factory(*a, **k):
            captured['provider'] = orig_factory(*a, **k)
            return captured['provider']

        fsc._overwrite_changes_provider = capturing_factory

        orig_cas = fsc.snapshot_commit.commit
        cas = {'fails': 0}

        def patched_cas(base_snapshot_uuid, snapshot, statistics):
            if snapshot.commit_kind == "OVERWRITE" and cas['fails'] < K:
                cas['fails'] += 1
                concurrent_fn(cas['fails'])
                return False
            return orig_cas(base_snapshot_uuid, snapshot, statistics)

        fsc.snapshot_commit.commit = patched_cas

        c.commit(messages)
        c.close()

        self.assertEqual(cas['fails'], K, "expected exactly K forced conflicts")
        return captured['provider']

    def test_cache_applies_delta_when_concurrent_overwrite_hits_target_partition(self):
        # Concurrent overwrites produce DELETE+ADD deltas. They can be applied
        # to the cached target state just like APPEND deltas.
        K = 2
        wb = self.table.new_batch_write_builder().overwrite({'f0': 1})
        w = wb.new_write()
        c = wb.new_commit()
        w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']}))
        provider = self._run_with_conflicts(
            c, w.prepare_commit(), K,
            lambda i: self._overwrite_partition(1, f'z{i}'))

        self.assertEqual(provider.full_scan_count, 1)
        self.assertEqual(provider.delta_probe_count, K)
        self.assertEqual(provider.delta_apply_count, K)

        read_builder = self.table.new_read_builder()
        actual = read_builder.new_read().to_pandas(
            read_builder.new_scan().plan().splits())
        self.assertEqual(sorted(actual[actual['f0'] == 1]['f1'].tolist()), ['new'])
        self.assertEqual(sorted(actual[actual['f0'] == 2]['f1'].tolist()), ['c'])

    def test_whole_table_overwrite_advances_by_append_delta(self):
        # Whole-table overwrite (no partition filter) can still advance through
        # APPEND deltas because every appended file belongs to the target state.
        K = 2
        wb = self.table.new_batch_write_builder().overwrite()
        w = wb.new_write()
        c = wb.new_commit()
        w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']}))
        provider = self._run_with_conflicts(
            c, w.prepare_commit(), K,
            lambda i: self._append(pd.DataFrame({'f0': [99], 'f1': [f'x{i}']})))

        self.assertEqual(provider.full_scan_count, 1)
        self.assertEqual(provider.delta_probe_count, K)
        self.assertEqual(provider.delta_apply_count, K)

        read_builder = self.table.new_read_builder()
        actual = read_builder.new_read().to_pandas(
            read_builder.new_scan().plan().splits())
        self.assertEqual(sorted(actual['f1'].tolist()), ['new'])

    def test_dynamic_partition_overwrite_reuses_cache(self):
        # Dynamic-partition overwrite is scoped to the data's partitions, so an
        # unrelated concurrent append lets the cache be reused.
        pa_schema = pa.schema([('f0', pa.int32()), ('f1', pa.string())])
        schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['f0'])
        self.catalog.create_table('test_db.t_dyn', schema, False)
        table = self.catalog.get_table('test_db.t_dyn')

        def append(df):
            wb = table.new_batch_write_builder()
            w = wb.new_write()
            c = wb.new_commit()
            w.write_pandas(df)
            c.commit(w.prepare_commit())
            w.close()
            c.close()

        append(pd.DataFrame({'f0': [1, 1, 2], 'f1': ['a', 'b', 'c']}))

        K = 3
        wb = table.new_batch_write_builder().overwrite()  # dynamic: filter from data (f0=1)
        w = wb.new_write()
        c = wb.new_commit()
        w.write_pandas(pd.DataFrame({'f0': [1], 'f1': ['new']}))
        provider = self._run_with_conflicts(
            c, w.prepare_commit(), K,
            lambda i: append(pd.DataFrame({'f0': [99], 'f1': [f'x{i}']})))

        self.assertEqual(provider.full_scan_count, 1)   # target f0=1 untouched -> reuse
        self.assertEqual(provider.delta_probe_count, K)


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