# 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 tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import pyarrow as pa

from pypaimon.execution.python.read.reader.concat_batch_reader import (DataEvolutionMergeReader, MergeAllBatchReader)
from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader


class _BatchReader(RecordBatchReader):

    def __init__(self, batches):
        self._batches = iter(batches)

    def read_arrow_batch(self):
        return next(self._batches, None)

    def close(self):
        pass


class ConcatBatchReaderTest(unittest.TestCase):

    def test_merge_all_resolves_blob_refs_and_keeps_row_payloads(self):
        from pypaimon.common.options import Options
        from pypaimon.common.options.core_options import CoreOptions
        from pypaimon.filesystem.local_file_io import LocalFileIO
        from pypaimon.execution.python.read.reader.blob_descriptor_convert_reader import (
            BlobDataBatchReader,
            BlobInlineConvertReader,
        )
        from pypaimon.table.row.blob import BlobDescriptor

        with tempfile.TemporaryDirectory() as root:
            # Both representations can have the same descriptor-shaped bytes.
            payload = BlobDescriptor(str(Path(root) / 'must-not-read'), 0, 1).serialize()
            source = Path(root) / 'source'
            source.write_bytes(payload)
            reference = BlobDescriptor(str(source), 0, len(payload)).serialize()
            table = SimpleNamespace(
                options=CoreOptions(Options({'blob-descriptor-field': 'payload'})),
                file_io=LocalFileIO.create(),
                catalog_environment=SimpleNamespace(catalog_loader=None))
            schema = pa.schema([pa.field(
                'payload', pa.large_binary(), metadata={b'description': b'keep'})])
            for reverse in (False, True):
                with self.subTest(reverse=reverse):
                    refs = _BatchReader([pa.RecordBatch.from_arrays(
                        [pa.array([reference], type=pa.large_binary())], schema=schema)])
                    data = BlobDataBatchReader(_BatchReader([pa.RecordBatch.from_arrays(
                        [pa.array([payload, None], type=pa.large_binary())], schema=schema)]), {'payload'})
                    sources = [data, refs] if reverse else [refs, data]
                    merged = MergeAllBatchReader(
                        [lambda reader=reader: reader for reader in sources], batch_size=3)
                    reader = BlobInlineConvertReader(merged, table)
                    values = []
                    sizes = []
                    try:
                        while True:
                            batch = reader.read_arrow_batch()
                            if batch is None:
                                break
                            values.extend(batch.column(0).to_pylist())
                            sizes.append(batch.num_rows)
                            metadata = batch.schema.field(0).metadata or {}
                            self.assertNotIn(b'paimon.row-sidecar.blob-data', metadata)
                            if batch.num_rows == 2:
                                self.assertEqual(metadata, {b'description': b'keep'})
                    finally:
                        reader.close()
                    self.assertEqual(values, [payload, None, payload] if reverse else [payload, payload, None])
                    self.assertEqual(sizes, [2, 1] if reverse else [1, 2])

    def test_merge_all_zero_column_batches_keep_row_counts(self):
        batches = [pa.record_batch([pa.array(range(size))], names=['id']).select([])
                   for size in (1, 2, 2)]
        reader = MergeAllBatchReader(
            [lambda batch=batch: _BatchReader([batch]) for batch in batches], batch_size=2)
        actual = []
        while True:
            batch = reader.read_arrow_batch()
            if batch is None:
                break
            self.assertEqual(batch.num_columns, 0)
            actual.append(batch.num_rows)
        self.assertEqual(actual, [2, 2, 1])

    def test_merge_all_coalesces_small_batches(self):
        batches = [
            pa.record_batch([pa.array(["a", "b"])], names=["value"]),
            pa.record_batch([pa.array(["c", "d", "e"])], names=["value"]),
        ]
        reader = MergeAllBatchReader(
            [lambda batch=batch: _BatchReader([batch]) for batch in batches],
            batch_size=2,
        )

        actual = []
        while True:
            batch = reader.read_arrow_batch()
            if batch is None:
                break
            actual.append(batch.column(0).to_pylist())

        self.assertEqual(actual, [["a", "b"], ["c", "d"], ["e"]])

    def test_merge_all_flushes_before_arrow_offset_limit(self):
        batches = [
            pa.record_batch([pa.array(["aaaa"])], names=["value"]),
            pa.record_batch([pa.array(["bbbb"])], names=["value"]),
        ]
        reader = MergeAllBatchReader(
            [lambda batch=batch: _BatchReader([batch]) for batch in batches],
            batch_size=2,
        )

        with patch(
                "pypaimon.execution.python.read.reader.concat_batch_reader._MAX_ARROW_OFFSET",
                4):
            actual = []
            while True:
                batch = reader.read_arrow_batch()
                if batch is None:
                    break
                actual.append(batch.column(0).to_pylist())

        self.assertEqual(actual, [["aaaa"], ["bbbb"]])

    def test_merge_all_flushes_before_list_offset_limit(self):
        child_count = 1100000000

        def batch():
            values = pa.nulls(child_count)
            offsets = pa.array([0, child_count], type=pa.int32())
            return pa.record_batch(
                [pa.ListArray.from_arrays(offsets, values)], names=["value"])

        batches = [batch(), batch()]
        self.assertTrue(all(item.nbytes < child_count for item in batches))
        reader = MergeAllBatchReader(
            [lambda item=item: _BatchReader([item]) for item in batches],
            batch_size=2,
        )

        actual = []
        while True:
            item = reader.read_arrow_batch()
            if item is None:
                break
            actual.append(item.column(0).value_lengths()[0].as_py())

        self.assertEqual(actual, [child_count, child_count])

    def test_merge_all_checks_nested_and_map_offsets(self):
        arrays = [
            pa.array([["aaa"]], type=pa.list_(pa.string())),
            pa.array(
                [[("a", 1), ("b", 2)]],
                type=pa.map_(pa.string(), pa.int32()),
            ),
        ]
        for array in arrays:
            with self.subTest(data_type=array.type):
                batches = [
                    pa.record_batch([array], names=["value"]),
                    pa.record_batch([array], names=["value"]),
                ]
                reader = MergeAllBatchReader(
                    [
                        lambda item=item: _BatchReader([item])
                        for item in batches
                    ],
                    batch_size=2,
                )
                with patch(
                        "pypaimon.execution.python.read.reader.concat_batch_reader."
                        "_MAX_ARROW_OFFSET",
                        3):
                    sizes = []
                    while True:
                        item = reader.read_arrow_batch()
                        if item is None:
                            break
                        sizes.append(item.num_rows)

                self.assertEqual(sizes, [1, 1])

    def test_misaligned_small_files_keep_bounded_batch_count(self):
        row_count = 10000
        left = MergeAllBatchReader([
            lambda value=value: _BatchReader([
                pa.record_batch([pa.array([value])], names=["left"])
            ])
            for value in range(row_count)
        ])
        right = _BatchReader([
            pa.record_batch(
                [pa.array(range(start, min(start + 1024, row_count)))],
                names=["right"],
            )
            for start in range(0, row_count, 1024)
        ])
        reader = DataEvolutionMergeReader(
            row_offsets=[0, 1],
            field_offsets=[0, 0],
            readers=[left, right],
            schema=pa.schema([("left", pa.int64()), ("right", pa.int64())]),
        )

        batch_sizes = []
        values = []
        while True:
            batch = reader.read_arrow_batch()
            if batch is None:
                break
            batch_sizes.append(batch.num_rows)
            values.extend(batch.column(0).to_pylist())

        self.assertEqual(batch_sizes, [1024] * 9 + [784])
        self.assertEqual(values, list(range(row_count)))

    def test_merge_reader_does_not_join_buffered_remainders(self):
        left = _BatchReader([
            pa.record_batch([pa.array([0])], names=["left"]),
            pa.record_batch([pa.array([1, 2])], names=["left"]),
        ])
        right = _BatchReader([
            pa.record_batch([pa.array([10, 11])], names=["right"]),
            pa.record_batch([pa.array([12])], names=["right"]),
        ])
        reader = DataEvolutionMergeReader(
            row_offsets=[0, 1],
            field_offsets=[0, 0],
            readers=[left, right],
            schema=pa.schema([("left", pa.int64()), ("right", pa.int64())]),
        )

        with patch.object(
                pa, "concat_arrays",
                side_effect=AssertionError("must preserve Arrow chunks")):
            actual = []
            while True:
                batch = reader.read_arrow_batch()
                if batch is None:
                    break
                actual.extend(
                    {"left": left, "right": right}
                    for left, right in zip(
                        batch.column(0).to_pylist(),
                        batch.column(1).to_pylist(),
                    )
                )

        self.assertEqual(actual, [
            {"left": 0, "right": 10},
            {"left": 1, "right": 11},
            {"left": 2, "right": 12},
        ])


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