# 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 os
import shutil
import tempfile
import threading
import unittest
from unittest import mock

import pyarrow as pa

from pypaimon import CatalogFactory, Schema
from pypaimon.read.table_read import TableRead, _RemainingRows


class RemainingRowsTest(unittest.TestCase):
    """Pure unit tests for the row-quota counter — no Paimon table needed."""

    def test_unlimited(self):
        rr = _RemainingRows(None)
        self.assertEqual(rr.try_consume(1_000_000), 1_000_000)
        self.assertEqual(rr.try_consume(1), 1)
        self.assertFalse(rr.exhausted())

    def test_basic_pre_debit(self):
        rr = _RemainingRows(100)
        self.assertEqual(rr.try_consume(40), 40)
        self.assertEqual(rr.try_consume(40), 40)
        # Only 20 left, asking for 30 returns 20.
        self.assertEqual(rr.try_consume(30), 20)
        self.assertTrue(rr.exhausted())
        self.assertEqual(rr.try_consume(1), 0)

    def test_zero_request(self):
        rr = _RemainingRows(100)
        self.assertEqual(rr.try_consume(0), 0)
        # Quota unchanged.
        self.assertEqual(rr.try_consume(100), 100)

    def test_concurrent_consume_never_overcommits(self):
        rr = _RemainingRows(10_000)
        granted = []
        granted_lock = threading.Lock()
        barrier = threading.Barrier(8)

        def worker():
            barrier.wait()
            for _ in range(2000):
                got = rr.try_consume(7)
                if got:
                    with granted_lock:
                        granted.append(got)

        threads = [threading.Thread(target=worker) for _ in range(8)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        self.assertEqual(sum(granted), 10_000)
        self.assertTrue(rr.exhausted())


class ParallelReaderAppendOnlyTest(unittest.TestCase):
    """Append-only multi-partition table — parallel must match serial exactly."""

    @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', False)

        cls.pa_schema = pa.schema([
            ('user_id', pa.int64()),
            ('item_id', pa.int64()),
            ('behavior', pa.string()),
            ('dt', pa.string()),
        ])

        # 8 partitions => 8 splits.
        rows_per_partition = 250
        user_ids, item_ids, behaviors, dts = [], [], [], []
        for p in range(8):
            for i in range(rows_per_partition):
                user_ids.append(p * rows_per_partition + i)
                item_ids.append(1000 + i)
                behaviors.append(f"act-{i % 5}")
                dts.append(f"p{p}")
        cls.expected_rows = len(user_ids)
        data = pa.Table.from_pydict({
            'user_id': user_ids,
            'item_id': item_ids,
            'behavior': behaviors,
            'dt': dts,
        }, schema=cls.pa_schema)

        # Default table — read.parallelism unset (defaults to 1, i.e. serial).
        cls.table = cls._build_table('append_parallel_default', None, data)
        # Option-set table — read.parallelism=4 baked into the table schema.
        cls.table_opt_4 = cls._build_table(
            'append_parallel_opt4', {'read.parallelism': '4'}, data)

    @classmethod
    def _build_table(cls, name, options, data):
        schema = Schema.from_pyarrow_schema(
            cls.pa_schema,
            partition_keys=['dt'],
            options=options,
        )
        cls.catalog.create_table(f'default.{name}', schema, False)
        table = cls.catalog.get_table(f'default.{name}')
        wb = table.new_batch_write_builder()
        w, c = wb.new_write(), wb.new_commit()
        w.write_arrow(data)
        c.commit(w.prepare_commit())
        w.close()
        c.close()
        return table

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

    def _scan_splits(self, read_builder):
        return read_builder.new_scan().plan().splits()

    def test_multi_partition_yields_multiple_splits(self):
        splits = self._scan_splits(self.table.new_read_builder())
        self.assertGreaterEqual(len(splits), 4,
                                f"expected multi-split fixture, got {len(splits)}")

    # ------------------------------------------------------------------
    # Result-parity tests for the two opt-in paths.
    # ------------------------------------------------------------------

    def test_parallel_via_method_arg_matches_serial(self):
        rb = self.table.new_read_builder()
        splits = self._scan_splits(rb)
        read = rb.new_read()
        serial = read.to_arrow(splits)
        parallel = read.to_arrow(splits, parallelism=4)
        # Same split order preserved => byte-identical tables.
        self.assertEqual(serial, parallel)

        df_serial = serial.to_pandas().sort_values('user_id').reset_index(drop=True)
        df_parallel = read.to_pandas(splits, parallelism=4) \
            .sort_values('user_id').reset_index(drop=True)
        self.assertTrue(df_serial.equals(df_parallel))
        self.assertEqual(len(df_parallel), self.expected_rows)

    def test_parallel_via_table_option_matches_serial(self):
        rb_serial = self.table.new_read_builder()
        rb_parallel = self.table_opt_4.new_read_builder()
        splits_serial = self._scan_splits(rb_serial)
        splits_parallel = self._scan_splits(rb_parallel)

        serial_df = rb_serial.new_read().to_pandas(splits_serial) \
            .sort_values('user_id').reset_index(drop=True)
        # No explicit parallelism — must pick up read.parallelism=4 from the table option.
        parallel_df = rb_parallel.new_read().to_pandas(splits_parallel) \
            .sort_values('user_id').reset_index(drop=True)
        self.assertTrue(serial_df.equals(parallel_df))

    # ------------------------------------------------------------------
    # Priority: method arg > table option > built-in default.
    # ------------------------------------------------------------------

    def test_method_arg_overrides_option_to_serial(self):
        # option=4 but caller passes 1: should disable parallelism.
        read = self.table_opt_4.new_read_builder().new_read()
        with mock.patch.object(read, '_to_arrow_parallel') as patched:
            patched.side_effect = AssertionError(
                "_to_arrow_parallel should not be called when arg=1 overrides option")
            splits = self._scan_splits(self.table_opt_4.new_read_builder())
            read.to_arrow(splits, parallelism=1)

    def test_method_arg_overrides_option_to_parallel(self):
        # option=1 (default) but caller passes 4: should enable parallelism.
        read = self.table.new_read_builder().new_read()
        splits = self._scan_splits(self.table.new_read_builder())
        with mock.patch.object(
            read, '_to_arrow_parallel', wraps=read._to_arrow_parallel
        ) as spy:
            result = read.to_arrow(splits, parallelism=4)
        spy.assert_called_once()
        self.assertEqual(result.num_rows, self.expected_rows)

    # ------------------------------------------------------------------
    # Boundary / invalid value handling.
    # ------------------------------------------------------------------

    def test_parallelism_one_equals_serial(self):
        rb = self.table.new_read_builder()
        splits = self._scan_splits(rb)
        read = rb.new_read()
        self.assertEqual(read.to_arrow(splits),
                         read.to_arrow(splits, parallelism=1))

    def test_parallelism_exceeds_split_count(self):
        rb = self.table.new_read_builder()
        splits = self._scan_splits(rb)
        read = rb.new_read()
        # 64 workers but only ~8 splits — should clamp internally, no error.
        result = read.to_arrow(splits, parallelism=64)
        self.assertEqual(result.num_rows, self.expected_rows)

    def test_invalid_method_arg_raises(self):
        rb = self.table.new_read_builder()
        splits = self._scan_splits(rb)
        read = rb.new_read()
        with self.assertRaises(ValueError) as ctx:
            read.to_arrow(splits, parallelism=0)
        self.assertIn("parallelism", str(ctx.exception))
        self.assertNotIn("read.parallelism", str(ctx.exception))
        with self.assertRaises(ValueError):
            read.to_pandas(splits, parallelism=-1)

    def test_invalid_option_value_raises(self):
        # Build a fresh table with an invalid option value.
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            partition_keys=['dt'],
            options={'read.parallelism': '0'},
        )
        self.catalog.create_table('default.bad_option', schema, False)
        bad = self.catalog.get_table('default.bad_option')
        # Reproduce the data so split planning yields a non-trivial plan.
        wb = bad.new_batch_write_builder()
        w, c = wb.new_write(), wb.new_commit()
        w.write_arrow(pa.Table.from_pydict({
            'user_id': [1, 2], 'item_id': [10, 20],
            'behavior': ['a', 'b'], 'dt': ['p1', 'p2'],
        }, schema=self.pa_schema))
        c.commit(w.prepare_commit())
        w.close()
        c.close()

        read = bad.new_read_builder().new_read()
        splits = bad.new_read_builder().new_scan().plan().splits()
        with self.assertRaises(ValueError) as ctx:
            read.to_arrow(splits)
        self.assertIn("read.parallelism", str(ctx.exception))

    def test_empty_splits_with_parallel_arg(self):
        rb = self.table.new_read_builder()
        read = rb.new_read()
        result = read.to_arrow([], parallelism=4)
        self.assertEqual(result.num_rows, 0)
        self.assertEqual([f.name for f in result.schema],
                         [f.name for f in self.pa_schema])

    def test_parallel_with_limit_soft_stop(self):
        # 10 calls with limit=600 should all return exactly 600 rows.
        limit = 600
        for _ in range(10):
            rb = self.table.new_read_builder().with_limit(limit)
            splits = self._scan_splits(rb)
            df = rb.new_read().to_pandas(splits, parallelism=4)
            self.assertEqual(len(df), limit)

    def test_parallel_reader_error_propagates(self):
        rb = self.table.new_read_builder()
        splits = self._scan_splits(rb)
        self.assertGreaterEqual(len(splits), 2)

        original_create = TableRead._create_split_read
        call_counter = {'n': 0}
        lock = threading.Lock()

        def flaky(self_, split):
            with lock:
                call_counter['n'] += 1
                idx = call_counter['n']
            if idx == 2:
                raise RuntimeError("simulated reader failure")
            return original_create(self_, split)

        with mock.patch.object(TableRead, '_create_split_read', flaky):
            read = rb.new_read()
            with self.assertRaises(RuntimeError) as ctx:
                read.to_pandas(splits, parallelism=4)
            self.assertIn("simulated reader failure", str(ctx.exception))


class ParallelReaderPrimaryKeyTest(unittest.TestCase):
    """PK + multi-bucket merge-on-read parity between serial and parallel."""

    @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', False)

        cls.pa_schema = pa.schema([
            pa.field('user_id', pa.int32(), nullable=False),
            ('item_id', pa.int64()),
            ('behavior', pa.string()),
            pa.field('dt', pa.string(), nullable=False),
        ])

        schema = Schema.from_pyarrow_schema(
            cls.pa_schema,
            partition_keys=['dt'],
            primary_keys=['user_id', 'dt'],
            options={'bucket': '4'},
        )
        cls.catalog.create_table('default.pk_parallel', schema, False)
        cls.table = cls.catalog.get_table('default.pk_parallel')

        # First snapshot.
        v1 = pa.Table.from_pydict({
            'user_id': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] * 4,
            'item_id': list(range(1001, 1041)),
            'behavior': [f"v1-{i}" for i in range(40)],
            'dt': (['p1'] * 10 + ['p2'] * 10 + ['p3'] * 10 + ['p4'] * 10),
        }, schema=cls.pa_schema)
        wb = cls.table.new_batch_write_builder()
        w, c = wb.new_write(), wb.new_commit()
        w.write_arrow(v1)
        c.commit(w.prepare_commit())
        w.close()
        c.close()

        # Second snapshot — updates some rows for the same PK to exercise merge.
        v2 = pa.Table.from_pydict({
            'user_id': [1, 2, 3, 4, 5],
            'item_id': [9001, 9002, 9003, 9004, 9005],
            'behavior': ['v2-updated'] * 5,
            'dt': ['p1', 'p1', 'p2', 'p2', 'p3'],
        }, schema=cls.pa_schema)
        wb = cls.table.new_batch_write_builder()
        w, c = wb.new_write(), wb.new_commit()
        w.write_arrow(v2)
        c.commit(w.prepare_commit())
        w.close()
        c.close()

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

    def test_multi_bucket_yields_multiple_splits(self):
        splits = self.table.new_read_builder().new_scan().plan().splits()
        self.assertGreaterEqual(
            len(splits), 2,
            f"expected multi-bucket fixture to yield >= 2 splits, got {len(splits)}")

    def test_parallel_merge_matches_serial(self):
        rb = self.table.new_read_builder()
        splits = rb.new_scan().plan().splits()
        read = rb.new_read()
        serial = read.to_pandas(splits).sort_values(
            ['dt', 'user_id']).reset_index(drop=True)
        parallel = read.to_pandas(splits, parallelism=4).sort_values(
            ['dt', 'user_id']).reset_index(drop=True)
        self.assertTrue(serial.equals(parallel))
        # Ensure the merge actually picked the latest version.
        # user_id=1 / dt=p1 must have behavior='v2-updated', item_id=9001.
        updated = parallel[(parallel.user_id == 1) & (parallel.dt == 'p1')]
        self.assertEqual(len(updated), 1)
        self.assertEqual(updated.iloc[0].behavior, 'v2-updated')
        self.assertEqual(updated.iloc[0].item_id, 9001)

    def test_parallel_with_limit_pk(self):
        limit = 12
        rb = self.table.new_read_builder().with_limit(limit)
        splits = rb.new_scan().plan().splits()
        df = rb.new_read().to_pandas(splits, parallelism=4)
        self.assertLessEqual(len(df), limit)
        # PK table merge applies limit per-split internally; in addition the
        # parallel soft-stop caps the global total. The combined output must
        # never exceed the user-visible limit.

    def test_include_row_kind_parallel(self):
        rb = self.table.new_read_builder()
        splits = rb.new_scan().plan().splits()
        read = rb.new_read()
        read.include_row_kind = True
        serial = read.to_arrow(splits)
        parallel = read.to_arrow(splits, parallelism=4)
        self.assertEqual(serial.schema, parallel.schema)
        self.assertIn('_row_kind', serial.schema.names)
        df_s = serial.to_pandas().sort_values(['dt', 'user_id']).reset_index(drop=True)
        df_p = parallel.to_pandas().sort_values(['dt', 'user_id']).reset_index(drop=True)
        self.assertTrue(df_s.equals(df_p))


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