"""
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 unittest

import pyarrow as pa
from pypaimon import Schema, CatalogFactory


class DataEvolutionTest(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', False)

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

    def test_basic(self):
        simple_pa_schema = pa.schema([
            ('f0', pa.int8()),
            ('f1', pa.int16()),
        ])
        schema = Schema.from_pyarrow_schema(simple_pa_schema,
                                            options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'})
        self.catalog.create_table('default.test_row_tracking', schema, False)
        table = self.catalog.get_table('default.test_row_tracking')

        # write 1
        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        expect_data = pa.Table.from_pydict({
            'f0': [-1, 2],
            'f1': [-1001, 1002]
        }, schema=simple_pa_schema)
        table_write.write_arrow(expect_data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # write 2
        table_write = write_builder.new_write().with_write_type(['f0'])
        table_commit = write_builder.new_commit()
        data2 = pa.Table.from_pydict({
            'f0': [3, 4],
        }, schema=pa.schema([
            ('f0', pa.int8()),
        ]))
        table_write.write_arrow(data2)
        cmts = table_write.prepare_commit()
        cmts[0].new_files[0].first_row_id = 0
        table_commit.commit(cmts)
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual_data = table_read.to_arrow(table_scan.plan().splits())
        expect_data = pa.Table.from_pydict({
            'f0': [3, 4],
            'f1': [-1001, 1002]
        }, schema=pa.schema([
            ('f0', pa.int8()),
            ('f1', pa.int16()),
        ]))
        self.assertEqual(actual_data, expect_data)

    def test_multiple_appends(self):
        simple_pa_schema = pa.schema([
            ('f0', pa.int32()),
            ('f1', pa.string()),
            ('f2', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(
            simple_pa_schema,
            options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}
        )
        self.catalog.create_table('default.test_multiple_appends', schema, False)
        table = self.catalog.get_table('default.test_multiple_appends')

        write_builder = table.new_batch_write_builder()

        # write 100 rows: (1, "a", "b")
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        init_data = pa.Table.from_pydict({
            'f0': [1] * 100,
            'f1': ['a'] * 100,
            'f2': ['b'] * 100,
        }, schema=simple_pa_schema)
        table_write.write_arrow(init_data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()
        # append：set first_row_id = 100 to modify the row with columns write
        write0 = write_builder.new_write().with_write_type(['f0', 'f1'])
        write1 = write_builder.new_write().with_write_type(['f2'])
        commit = write_builder.new_commit()
        data0 = pa.Table.from_pydict({'f0': [2], 'f1': ['x']},
                                     schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())]))
        data1 = pa.Table.from_pydict({'f2': ['y']}, schema=pa.schema([('f2', pa.string())]))
        write0.write_arrow(data0)
        write1.write_arrow(data1)
        cmts = write0.prepare_commit() + write1.prepare_commit()
        for c in cmts:
            for nf in c.new_files:
                nf.first_row_id = 100
        commit.commit(cmts)
        write0.close()
        write1.close()
        commit.close()

        # append：write (3, "c") and ("d"), set first_row_id = 101
        write0 = write_builder.new_write().with_write_type(['f0', 'f1'])
        commit0 = write_builder.new_commit()
        data0 = pa.Table.from_pydict({'f0': [3], 'f1': ['c']},
                                     schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())]))
        write0.write_arrow(data0)
        cmts0 = write0.prepare_commit()
        for c in cmts0:
            for nf in c.new_files:
                nf.first_row_id = 101
        commit0.commit(cmts0)
        write0.close()
        commit0.close()

        write1 = write_builder.new_write().with_write_type(['f2'])
        commit1 = write_builder.new_commit()
        data1 = pa.Table.from_pydict({'f2': ['d']}, schema=pa.schema([('f2', pa.string())]))
        write1.write_arrow(data1)
        cmts1 = write1.prepare_commit()
        for c in cmts1:
            for nf in c.new_files:
                nf.first_row_id = 101
        commit1.commit(cmts1)
        write1.close()
        commit1.close()

        read_builder = table.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual = table_read.to_arrow(table_scan.plan().splits())

        self.assertEqual(actual.num_rows, 102)
        expect = pa.Table.from_pydict({
            'f0': [1] * 100 + [2] + [3],
            'f1': ['a'] * 100 + ['x'] + ['c'],
            'f2': ['b'] * 100 + ['y'] + ['d'],
        }, schema=simple_pa_schema)
        self.assertEqual(actual, expect)

    def test_disorder_cols_append(self):
        simple_pa_schema = pa.schema([
            ('f0', pa.int32()),
            ('f1', pa.string()),
            ('f2', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(
            simple_pa_schema,
            options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}
        )
        self.catalog.create_table('default.test_disorder_cols_append', schema, False)
        table = self.catalog.get_table('default.test_disorder_cols_append')

        write_builder = table.new_batch_write_builder()
        num_rows = 100
        # write 1 rows: (1, "a", "b")
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        init_data = pa.Table.from_pydict({
            'f0': [1] * num_rows,
            'f1': ['a'] * num_rows,
            'f2': ['b'] * num_rows,
        }, schema=simple_pa_schema)
        table_write.write_arrow(init_data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # append：set first_row_id = 0 to modify the row with columns write
        write0 = write_builder.new_write().with_write_type(['f0', 'f2'])
        write1 = write_builder.new_write().with_write_type(['f1'])
        commit = write_builder.new_commit()
        data0 = pa.Table.from_pydict({'f0': [2] * num_rows, 'f2': ['y'] * num_rows},
                                     schema=pa.schema([('f0', pa.int32()), ('f2', pa.string())]))
        data1 = pa.Table.from_pydict({'f1': ['x'] * num_rows}, schema=pa.schema([('f1', pa.string())]))
        write0.write_arrow(data0)
        write1.write_arrow(data1)
        cmts = write0.prepare_commit() + write1.prepare_commit()
        for c in cmts:
            for nf in c.new_files:
                nf.first_row_id = 0
        commit.commit(cmts)
        write0.close()
        write1.close()
        commit.close()

        read_builder = table.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual = table_read.to_arrow(table_scan.plan().splits())

        self.assertEqual(actual.num_rows, 100)
        expect = pa.Table.from_pydict({
            'f0': [2] * num_rows,
            'f1': ['x'] * num_rows,
            'f2': ['y'] * num_rows,
        }, schema=simple_pa_schema)
        self.assertEqual(actual, expect)

    def test_only_some_columns(self):
        simple_pa_schema = pa.schema([
            ('f0', pa.int32()),
            ('f1', pa.string()),
            ('f2', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(
            simple_pa_schema,
            options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}
        )
        self.catalog.create_table('default.test_only_some_columns', schema, False)
        table = self.catalog.get_table('default.test_only_some_columns')

        write_builder = table.new_batch_write_builder()

        # Commit 1: f0
        w0 = write_builder.new_write().with_write_type(['f0'])
        c0 = write_builder.new_commit()
        d0 = pa.Table.from_pydict({'f0': [1]}, schema=pa.schema([('f0', pa.int32())]))
        w0.write_arrow(d0)
        c0.commit(w0.prepare_commit())
        w0.close()
        c0.close()

        # Commit 2: f1, first_row_id = 0
        w1 = write_builder.new_write().with_write_type(['f1'])
        c1 = write_builder.new_commit()
        d1 = pa.Table.from_pydict({'f1': ['a']}, schema=pa.schema([('f1', pa.string())]))
        w1.write_arrow(d1)
        cmts1 = w1.prepare_commit()
        for c in cmts1:
            for nf in c.new_files:
                nf.first_row_id = 0
        c1.commit(cmts1)
        w1.close()
        c1.close()

        # Commit 3: f2, first_row_id = 0
        w2 = write_builder.new_write().with_write_type(['f2'])
        c2 = write_builder.new_commit()
        d2 = pa.Table.from_pydict({'f2': ['b']}, schema=pa.schema([('f2', pa.string())]))
        w2.write_arrow(d2)
        cmts2 = w2.prepare_commit()
        for c in cmts2:
            for nf in c.new_files:
                nf.first_row_id = 0
        c2.commit(cmts2)
        w2.close()
        c2.close()

        read_builder = table.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual = table_read.to_arrow(table_scan.plan().splits())

        expect = pa.Table.from_pydict({
            'f0': [1],
            'f1': ['a'],
            'f2': ['b'],
        }, schema=simple_pa_schema)
        self.assertEqual(actual, expect)

    def test_null_values(self):
        simple_pa_schema = pa.schema([
            ('f0', pa.int32()),
            ('f1', pa.string()),
            ('f2', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(
            simple_pa_schema,
            options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}
        )
        self.catalog.create_table('default.test_null_values', schema, False)
        table = self.catalog.get_table('default.test_null_values')

        write_builder = table.new_batch_write_builder()

        # Commit 1: some cols are null
        w0 = write_builder.new_write().with_write_type(['f0', 'f1'])
        w1 = write_builder.new_write().with_write_type(['f2'])
        c = write_builder.new_commit()

        d0 = pa.Table.from_pydict({'f0': [1], 'f1': [None]},
                                  schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())]))
        d1 = pa.Table.from_pydict({'f2': [None]}, schema=pa.schema([('f2', pa.string())]))
        w0.write_arrow(d0)
        w1.write_arrow(d1)
        cmts = w0.prepare_commit() + w1.prepare_commit()
        for msg in cmts:
            for nf in msg.new_files:
                nf.first_row_id = 0
        c.commit(cmts)
        w0.close()
        w1.close()
        c.close()

        # Commit 2
        w1 = write_builder.new_write().with_write_type(['f2'])
        c1 = write_builder.new_commit()
        d1 = pa.Table.from_pydict({'f2': ['c']}, schema=pa.schema([('f2', pa.string())]))
        w1.write_arrow(d1)
        cmts1 = w1.prepare_commit()
        for msg in cmts1:
            for nf in msg.new_files:
                nf.first_row_id = 0
        c1.commit(cmts1)
        w1.close()
        c1.close()

        read_builder = table.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual = table_read.to_arrow(table_scan.plan().splits())
        expect = pa.Table.from_pydict({
            'f0': [1],
            'f1': [None],
            'f2': ['c'],
        }, schema=simple_pa_schema)
        self.assertEqual(actual, expect)

    # different first_row_id append multiple times
    def test_multiple_appends_different_first_row_ids(self):
        simple_pa_schema = pa.schema([
            ('f0', pa.int32()),
            ('f1', pa.string()),
            ('f2', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(
            simple_pa_schema,
            options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}
        )
        self.catalog.create_table('default.test_multiple_appends_diff_rowid', schema, False)
        table = self.catalog.get_table('default.test_multiple_appends_diff_rowid')

        write_builder = table.new_batch_write_builder()

        # commit 1
        w0 = write_builder.new_write().with_write_type(['f0', 'f1'])
        w1 = write_builder.new_write().with_write_type(['f2'])
        c = write_builder.new_commit()
        d0 = pa.Table.from_pydict({'f0': [1], 'f1': ['a']},
                                  schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())]))
        d1 = pa.Table.from_pydict({'f2': ['b']}, schema=pa.schema([('f2', pa.string())]))
        w0.write_arrow(d0)
        w1.write_arrow(d1)
        cmts = w0.prepare_commit() + w1.prepare_commit()
        for msg in cmts:
            for nf in msg.new_files:
                nf.first_row_id = 0
        c.commit(cmts)
        w0.close()
        w1.close()
        c.close()

        # commit 2
        w0 = write_builder.new_write().with_write_type(['f0', 'f1'])
        c0 = write_builder.new_commit()
        d0 = pa.Table.from_pydict({'f0': [2], 'f1': ['c']},
                                  schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())]))
        w0.write_arrow(d0)
        cmts0 = w0.prepare_commit()
        for msg in cmts0:
            for nf in msg.new_files:
                nf.first_row_id = 1
        c0.commit(cmts0)
        w0.close()
        c0.close()

        # commit 3
        w1 = write_builder.new_write().with_write_type(['f2'])
        c1 = write_builder.new_commit()
        d1 = pa.Table.from_pydict({'f2': ['d']}, schema=pa.schema([('f2', pa.string())]))
        w1.write_arrow(d1)
        cmts1 = w1.prepare_commit()
        for msg in cmts1:
            for nf in msg.new_files:
                nf.first_row_id = 1
        c1.commit(cmts1)
        w1.close()
        c1.close()

        read_builder = table.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual = table_read.to_arrow(table_scan.plan().splits())

        expect = pa.Table.from_pydict({
            'f0': [1, 2],
            'f1': ['a', 'c'],
            'f2': ['b', 'd'],
        }, schema=simple_pa_schema)
        self.assertEqual(actual, expect)

    def test_more_data(self):
        simple_pa_schema = pa.schema([
            ('f0', pa.int32()),
            ('f1', pa.string()),
            ('f2', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(
            simple_pa_schema,
            options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}
        )
        self.catalog.create_table('default.test_more_data', schema, False)
        table = self.catalog.get_table('default.test_more_data')

        write_builder = table.new_batch_write_builder()

        # first commit：100k rows
        w0 = write_builder.new_write().with_write_type(['f0', 'f1'])
        w1 = write_builder.new_write().with_write_type(['f2'])
        c = write_builder.new_commit()
        size = 100000
        d0 = pa.Table.from_pydict({
            'f0': list(range(size)),
            'f1': [f'a{i}' for i in range(size)],
        }, schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())]))
        d1 = pa.Table.from_pydict({
            'f2': [f'b{i}' for i in range(size)],
        }, schema=pa.schema([('f2', pa.string())]))
        w0.write_arrow(d0)
        w1.write_arrow(d1)
        cmts = w0.prepare_commit() + w1.prepare_commit()
        for msg in cmts:
            for nf in msg.new_files:
                nf.first_row_id = 0
        c.commit(cmts)
        w0.close()
        w1.close()
        c.close()

        # second commit：overwrite f2 to 'c{i}'
        w1 = write_builder.new_write().with_write_type(['f2'])
        c1 = write_builder.new_commit()
        d1 = pa.Table.from_pydict({
            'f2': [f'c{i}' for i in range(size)],
        }, schema=pa.schema([('f2', pa.string())]))
        w1.write_arrow(d1)
        cmts1 = w1.prepare_commit()
        c1.commit(cmts1)
        w1.close()
        c1.close()

        read_builder = table.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual = table_read.to_arrow(table_scan.plan().splits())

        expect = pa.Table.from_pydict({
            'f0': list(range(size)),
            'f1': [f'a{i}' for i in range(size)],
            'f2': [f'c{i}' for i in range(size)],
        }, schema=simple_pa_schema)
        self.assertEqual(actual, expect)
