"""
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.
"""
from unittest.mock import patch

import pyarrow as pa

from pypaimon import Schema
from pypaimon.catalog.catalog_exception import (DatabaseAlreadyExistException,
                                                DatabaseNotExistException,
                                                TableAlreadyExistException,
                                                TableNotExistException)
from pypaimon.common.options import Options
from pypaimon.common.options.config import OssOptions
from pypaimon.filesystem.pyarrow_file_io import PyArrowFileIO
from pypaimon.tests.py36.pyarrow_compat import table_sort_by
from pypaimon.tests.rest.rest_base_test import RESTBaseTest


class AOSimpleTest(RESTBaseTest):
    def setUp(self):
        super().setUp()
        self.pa_schema = pa.schema([
            ('user_id', pa.int64()),
            ('item_id', pa.int64()),
            ('behavior', pa.string()),
            ('dt', pa.string()),
        ])
        self.data = {
            'user_id': [2, 4, 6, 8, 10],
            'item_id': [1001, 1002, 1003, 1004, 1005],
            'behavior': ['a', 'b', 'c', 'd', 'e'],
            'dt': ['2000-10-10', '2025-08-10', '2025-08-11', '2025-08-12', '2025-08-13']
        }
        self.expected = pa.Table.from_pydict(self.data, schema=self.pa_schema)

    def test_with_shard_ao_unaware_bucket(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.rest_catalog.create_table('default.test_with_shard_ao_unaware_bucket', schema, False)
        table = self.rest_catalog.get_table('default.test_with_shard_ao_unaware_bucket')
        write_builder = table.new_batch_write_builder()
        # first write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = {
            'user_id': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
            'item_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010, 1011, 1012, 1013, 1014],
            'behavior': ['a', 'b', 'c', None, 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm'],
            'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data1, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()
        # second write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data2 = {
            'user_id': [5, 6, 7, 8, 18],
            'item_id': [1005, 1006, 1007, 1008, 1018],
            'behavior': ['e', 'f', 'g', 'h', 'z'],
            'dt': ['p2', 'p1', 'p2', 'p2', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data2, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().with_shard(2, 3).plan().splits()
        actual = table_sort_by(table_read.to_arrow(splits), 'user_id')
        expected = pa.Table.from_pydict({
            'user_id': [5, 7, 8, 9, 11, 13],
            'item_id': [1005, 1007, 1008, 1009, 1011, 1013],
            'behavior': ['e', 'g', 'h', 'h', 'j', 'l'],
            'dt': ['p2', 'p2', 'p2', 'p2', 'p2', 'p2'],
        }, schema=self.pa_schema)
        self.assertEqual(actual, expected)

        # Get the three actual tables
        splits1 = read_builder.new_scan().with_shard(0, 3).plan().splits()
        actual1 = table_sort_by(table_read.to_arrow(splits1), 'user_id')
        splits2 = read_builder.new_scan().with_shard(1, 3).plan().splits()
        actual2 = table_sort_by(table_read.to_arrow(splits2), 'user_id')
        splits3 = read_builder.new_scan().with_shard(2, 3).plan().splits()
        actual3 = table_sort_by(table_read.to_arrow(splits3), 'user_id')

        # Concatenate the three tables
        actual = table_sort_by(pa.concat_tables([actual1, actual2, actual3]), 'user_id')
        expected = table_sort_by(self._read_test_table(read_builder), 'user_id')
        self.assertEqual(actual, expected)

    def test_with_shard_ao_fixed_bucket(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'],
                                            options={'bucket': '5', 'bucket-key': 'item_id'})
        self.rest_catalog.create_table('default.test_with_slice_ao_fixed_bucket', schema, False)
        table = self.rest_catalog.get_table('default.test_with_slice_ao_fixed_bucket')
        write_builder = table.new_batch_write_builder()
        # first write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = {
            'user_id': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
            'item_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010, 1011, 1012, 1013, 1014],
            'behavior': ['a', 'b', 'c', None, 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm'],
            'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data1, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()
        # second write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data2 = {
            'user_id': [5, 6, 7, 8],
            'item_id': [1005, 1006, 1007, 1008],
            'behavior': ['e', 'f', 'g', 'h'],
            'dt': ['p2', 'p1', 'p2', 'p2'],
        }
        pa_table = pa.Table.from_pydict(data2, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().with_shard(0, 3).plan().splits()
        actual = table_sort_by(table_read.to_arrow(splits), 'user_id')
        expected = pa.Table.from_pydict({
            'user_id': [1, 2, 3, 4, 5, 13],
            'item_id': [1001, 1002, 1003, 1004, 1005, 1013],
            'behavior': ['a', 'b', 'c', None, 'd', 'l'],
            'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p2'],
        }, schema=self.pa_schema)
        self.assertEqual(actual, expected)

        # Get the three actual tables
        splits1 = read_builder.new_scan().with_shard(0, 3).plan().splits()
        actual1 = table_sort_by(table_read.to_arrow(splits1), 'user_id')
        splits2 = read_builder.new_scan().with_shard(1, 3).plan().splits()
        actual2 = table_sort_by(table_read.to_arrow(splits2), 'user_id')
        splits3 = read_builder.new_scan().with_shard(2, 3).plan().splits()
        actual3 = table_sort_by(table_read.to_arrow(splits3), 'user_id')

        # Concatenate the three tables
        actual = table_sort_by(pa.concat_tables([actual1, actual2, actual3]), 'user_id')
        expected = table_sort_by(self._read_test_table(read_builder), 'user_id')
        self.assertEqual(actual, expected)

    def test_shard_single_partition(self):
        """Test sharding with single partition - tests _filter_by_shard with simple data"""
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.rest_catalog.create_table('default.test_shard_single_partition', schema, False)
        table = self.rest_catalog.get_table('default.test_shard_single_partition')
        write_builder = table.new_batch_write_builder()

        # Write data with single partition
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = {
            'user_id': [1, 2, 3, 4, 5, 6],
            'item_id': [1001, 1002, 1003, 1004, 1005, 1006],
            'behavior': ['a', 'b', 'c', 'd', 'e', 'f'],
            'dt': ['p1', 'p1', 'p1', 'p1', 'p1', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()

        # Test first shard (0, 2) - should get first 3 rows
        splits = read_builder.new_scan().with_shard(0, 2).plan().splits()
        actual = table_sort_by(table_read.to_arrow(splits), 'user_id')
        expected = pa.Table.from_pydict({
            'user_id': [1, 2, 3],
            'item_id': [1001, 1002, 1003],
            'behavior': ['a', 'b', 'c'],
            'dt': ['p1', 'p1', 'p1'],
        }, schema=self.pa_schema)
        self.assertEqual(actual, expected)

        # Test second shard (1, 2) - should get last 3 rows
        splits = read_builder.new_scan().with_shard(1, 2).plan().splits()
        actual = table_sort_by(table_read.to_arrow(splits), 'user_id')
        expected = pa.Table.from_pydict({
            'user_id': [4, 5, 6],
            'item_id': [1004, 1005, 1006],
            'behavior': ['d', 'e', 'f'],
            'dt': ['p1', 'p1', 'p1'],
        }, schema=self.pa_schema)
        self.assertEqual(actual, expected)

    def test_shard_uneven_distribution(self):
        """Test sharding with uneven row distribution across shards"""
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.rest_catalog.create_table('default.test_shard_uneven', schema, False)
        table = self.rest_catalog.get_table('default.test_shard_uneven')
        write_builder = table.new_batch_write_builder()

        # Write data with 7 rows (not evenly divisible by 3)
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = {
            'user_id': [1, 2, 3, 4, 5, 6, 7],
            'item_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007],
            'behavior': ['a', 'b', 'c', 'd', 'e', 'f', 'g'],
            'dt': ['p1', 'p1', 'p1', 'p1', 'p1', 'p1', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()

        # Test sharding into 3 parts: 2, 2, 3 rows
        splits = read_builder.new_scan().with_shard(0, 3).plan().splits()
        actual1 = table_sort_by(table_read.to_arrow(splits), 'user_id')
        expected1 = pa.Table.from_pydict({
            'user_id': [1, 2, 3],
            'item_id': [1001, 1002, 1003],
            'behavior': ['a', 'b', 'c'],
            'dt': ['p1', 'p1', 'p1'],
        }, schema=self.pa_schema)
        self.assertEqual(actual1, expected1)

        splits = read_builder.new_scan().with_shard(1, 3).plan().splits()
        actual2 = table_sort_by(table_read.to_arrow(splits), 'user_id')
        expected2 = pa.Table.from_pydict({
            'user_id': [4, 5],
            'item_id': [1004, 1005],
            'behavior': ['d', 'e'],
            'dt': ['p1', 'p1'],
        }, schema=self.pa_schema)
        self.assertEqual(actual2, expected2)

        splits = read_builder.new_scan().with_shard(2, 3).plan().splits()
        actual3 = table_sort_by(table_read.to_arrow(splits), 'user_id')
        expected3 = pa.Table.from_pydict({
            'user_id': [6, 7],
            'item_id': [1006, 1007],
            'behavior': ['f', 'g'],
            'dt': ['p1', 'p1'],
        }, schema=self.pa_schema)
        self.assertEqual(actual3, expected3)

    def test_shard_many_small_shards(self):
        """Test sharding with many small shards"""
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.rest_catalog.create_table('default.test_shard_many_small', schema, False)
        table = self.rest_catalog.get_table('default.test_shard_many_small')
        write_builder = table.new_batch_write_builder()

        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = {
            'user_id': [1, 2, 3, 4, 5, 6],
            'item_id': [1001, 1002, 1003, 1004, 1005, 1006],
            'behavior': ['a', 'b', 'c', 'd', 'e', 'f'],
            'dt': ['p1', 'p1', 'p1', 'p1', 'p1', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()

        # Test with 6 shards (one row per shard)
        for i in range(6):
            splits = read_builder.new_scan().with_shard(i, 6).plan().splits()
            actual = table_read.to_arrow(splits)
            self.assertEqual(len(actual), 1)
            self.assertEqual(actual['user_id'][0].as_py(), i + 1)

    def test_shard_boundary_conditions(self):
        """Test sharding boundary conditions with edge cases"""
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.rest_catalog.create_table('default.test_shard_boundary', schema, False)
        table = self.rest_catalog.get_table('default.test_shard_boundary')
        write_builder = table.new_batch_write_builder()

        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = {
            'user_id': [1, 2, 3, 4, 5],
            'item_id': [1001, 1002, 1003, 1004, 1005],
            'behavior': ['a', 'b', 'c', 'd', 'e'],
            'dt': ['p1', 'p1', 'p1', 'p1', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()

        # Test first shard (0, 4) - should get 1 row (5//4 = 1)
        splits = read_builder.new_scan().with_shard(0, 4).plan().splits()
        actual = table_read.to_arrow(splits)
        self.assertEqual(len(actual), 2)

        # Test middle shard (1, 4) - should get 1 row
        splits = read_builder.new_scan().with_shard(1, 4).plan().splits()
        actual = table_read.to_arrow(splits)
        self.assertEqual(len(actual), 1)

        # Test last shard (3, 4) - should get 2 rows (remainder goes to last shard)
        splits = read_builder.new_scan().with_shard(3, 4).plan().splits()
        actual = table_read.to_arrow(splits)
        self.assertEqual(len(actual), 1)

    def test_create_drop_database_table(self):
        # test create database
        self.rest_catalog.create_database("db1", False)

        with self.assertRaises(DatabaseAlreadyExistException) as context:
            self.rest_catalog.create_database("db1", False)

        self.assertEqual("db1", context.exception.database)

        try:
            self.rest_catalog.create_database("db1", True)
        except DatabaseAlreadyExistException:
            self.fail("create_database with ignore_if_exists=True should not raise DatabaseAlreadyExistException")

        # test create table
        self.rest_catalog.create_table("db1.tbl1",
                                       Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']),
                                       False)
        with self.assertRaises(TableAlreadyExistException) as context:
            self.rest_catalog.create_table("db1.tbl1",
                                           Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']),
                                           False)
        self.assertEqual("db1.tbl1", context.exception.identifier.get_full_name())

        try:
            self.rest_catalog.create_table("db1.tbl1",
                                           Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']),
                                           True)
        except TableAlreadyExistException:
            self.fail("create_table with ignore_if_exists=True should not raise TableAlreadyExistException")

        # test drop table
        self.rest_catalog.drop_table("db1.tbl1", False)
        with self.assertRaises(TableNotExistException) as context:
            self.rest_catalog.drop_table("db1.tbl1", False)
        self.assertEqual("db1.tbl1", context.exception.identifier.get_full_name())

        try:
            self.rest_catalog.drop_table("db1.tbl1", True)
        except TableNotExistException:
            self.fail("drop_table with ignore_if_not_exists=True should not raise TableNotExistException")

        # test drop database
        self.rest_catalog.drop_database("db1", False)
        with self.assertRaises(DatabaseNotExistException) as context:
            self.rest_catalog.drop_database("db1", False)
        self.assertEqual("db1", context.exception.database)

        try:
            self.rest_catalog.drop_database("db1", True)
        except DatabaseNotExistException:
            self.fail("drop_database with ignore_if_not_exists=True should not raise DatabaseNotExistException")

    def test_initialize_oss_fs_pyarrow_lt_7(self):
        props = {
            OssOptions.OSS_ACCESS_KEY_ID.key(): "AKID",
            OssOptions.OSS_ACCESS_KEY_SECRET.key(): "SECRET",
            OssOptions.OSS_SECURITY_TOKEN.key(): "TOKEN",
            OssOptions.OSS_REGION.key(): "cn-hangzhou",
            OssOptions.OSS_ENDPOINT.key(): "oss-cn-hangzhou.aliyuncs.com",
        }

        with patch("pypaimon.common.file_io.pyarrow.__version__", "6.0.0"), \
                patch("pyarrow.fs.S3FileSystem") as mock_s3fs:
            PyArrowFileIO("oss://oss-bucket/paimon-database/paimon-table", Options(props))
            mock_s3fs.assert_called_once_with(access_key="AKID",
                                              secret_key="SECRET",
                                              session_token="TOKEN",
                                              region="cn-hangzhou",
                                              endpoint_override="oss-bucket." + props[OssOptions.OSS_ENDPOINT.key()])
            PyArrowFileIO("oss://oss-bucket.endpoint/paimon-database/paimon-table", Options(props))
            mock_s3fs.assert_called_with(access_key="AKID",
                                         secret_key="SECRET",
                                         session_token="TOKEN",
                                         region="cn-hangzhou",
                                         endpoint_override="oss-bucket." + props[OssOptions.OSS_ENDPOINT.key()])
            PyArrowFileIO("oss://access_id:secret_key@Endpoint/oss-bucket/paimon-database/paimon-table", Options(props))
            mock_s3fs.assert_called_with(access_key="AKID",
                                         secret_key="SECRET",
                                         session_token="TOKEN",
                                         region="cn-hangzhou",
                                         endpoint_override="oss-bucket." + props[OssOptions.OSS_ENDPOINT.key()])

    def test_multi_prepare_commit_ao(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.rest_catalog.create_table('default.test_append_only_parquet', schema, False)
        table = self.rest_catalog.get_table('default.test_append_only_parquet')
        write_builder = table.new_stream_write_builder()

        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        # write 1
        data1 = {
            'user_id': [1, 2, 3, 4],
            'item_id': [1001, 1002, 1003, 1004],
            'behavior': ['a', 'b', 'c', None],
            'dt': ['p1', 'p1', 'p2', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data1, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_write.prepare_commit(0)
        # write 2
        data2 = {
            'user_id': [5, 6, 7, 8],
            'item_id': [1005, 1006, 1007, 1008],
            'behavior': ['e', 'f', 'g', 'h'],
            'dt': ['p2', 'p1', 'p2', 'p2'],
        }
        pa_table = pa.Table.from_pydict(data2, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_write.prepare_commit(1)
        # write 3
        data3 = {
            'user_id': [9, 10],
            'item_id': [1009, 1010],
            'behavior': ['i', 'j'],
            'dt': ['p2', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data3, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        cm = table_write.prepare_commit(2)
        # commit
        table_commit.commit(cm, 2)
        table_write.close()
        table_commit.close()
        self.assertEqual(2, table_write.file_store_write.commit_identifier)

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()
        actual = table_sort_by(table_read.to_arrow(splits), 'user_id')
        expected = pa.Table.from_pydict({
            'user_id': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
            'item_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010],
            'behavior': ['a', 'b', 'c', None, 'e', 'f', 'g', 'h', 'i', 'j'],
            'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p2', 'p2', 'p1']
        }, schema=self.pa_schema)
        self.assertEqual(expected, actual)
