#  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 CatalogFactory, Schema
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.schema.schema_manager import SchemaManager


class SimpleTableTest(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', True)

        cls.pa_schema = pa.schema([
            ('pt', pa.int32()),
            ('k', pa.int32()),
            ('v', pa.int64())
        ])
        cls.pk_pa_schema = pa.schema([
            pa.field('pt', pa.int32(), nullable=False),
            pa.field('k', pa.int32(), nullable=False),
            ('v', pa.int64())
        ])

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

    def test_tag_scan(self):
        """
        Test reading from a specific tag.

        1. Write data in 3 commits
        2. Create a tag at snapshot 2
        3. Read from the tag and verify only data from snapshots 1 and 2 is returned
        """
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_tag_scan', schema, False)
        table = self.catalog.get_table('default.test_tag_scan')

        write_builder = table.new_batch_write_builder()

        # First commit: pt=1, k=10, v=100 and pt=1, k=20, v=200
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = pa.Table.from_pydict({
            'pt': [1, 1],
            'k': [10, 20],
            'v': [100, 200]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data1)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Second commit: pt=2, k=30, v=101 and pt=2, k=40, v=201
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data2 = pa.Table.from_pydict({
            'pt': [2, 2],
            'k': [30, 40],
            'v': [101, 201]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data2)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Third commit: pt=3, k=50, v=500 and pt=3, k=60, v=600
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data3 = pa.Table.from_pydict({
            'pt': [3, 3],
            'k': [50, 60],
            'v': [500, 600]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data3)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Create tag at snapshot 2
        table.create_tag("tag2", snapshot_id=2)

        # Read from tag2 using scan.tag-name option
        table_with_tag = table.copy({CoreOptions.SCAN_TAG_NAME.key(): "tag2"})
        read_builder = table_with_tag.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        result = table_read.to_arrow(table_scan.plan().splits())

        # Verify: should only contain data from snapshot 1 and 2
        # (pt=1, k=10, v=100), (pt=1, k=20, v=200), (pt=2, k=30, v=101), (pt=2, k=40, v=201)
        result_sorted = result.sort_by([('pt', 'ascending'), ('k', 'ascending')])

        expected = pa.Table.from_pydict({
            'pt': [1, 1, 2, 2],
            'k': [10, 20, 30, 40],
            'v': [100, 200, 101, 201]
        }, schema=self.pk_pa_schema)

        self.assertEqual(result_sorted.num_rows, 4)
        self.assertEqual(result_sorted.column('pt').to_pylist(), expected.column('pt').to_pylist())
        self.assertEqual(result_sorted.column('k').to_pylist(), expected.column('k').to_pylist())
        self.assertEqual(result_sorted.column('v').to_pylist(), expected.column('v').to_pylist())

    def test_non_existing_tag(self):
        """
        Test that reading from a non-existing tag raises an error.
        """
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_non_existing_tag', schema, False)
        table = self.catalog.get_table('default.test_non_existing_tag')

        # Try to read from a non-existing tag
        table_with_tag = table.copy({CoreOptions.SCAN_TAG_NAME.key(): "non-existing"})
        read_builder = table_with_tag.new_read_builder()
        table_scan = read_builder.new_scan()

        with self.assertRaises(ValueError) as context:
            table_scan.plan()

        self.assertIn("non-existing", str(context.exception))
        self.assertIn("doesn't exist", str(context.exception))

    def test_tag_create_and_delete(self):
        """Test creating and deleting tags."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_tag_create_delete', schema, False)
        table = self.catalog.get_table('default.test_tag_create_delete')

        write_builder = table.new_batch_write_builder()

        # Write some data
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = pa.Table.from_pydict({
            'pt': [1, 1],
            'k': [10, 20],
            'v': [100, 200]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Create a tag
        table.create_tag("test_tag")

        # Verify tag exists
        tag_manager = table.tag_manager()
        self.assertTrue(tag_manager.tag_exists("test_tag"))

        # Get the tag
        tag = tag_manager.get("test_tag")
        self.assertIsNotNone(tag)
        self.assertEqual(tag.id, 1)

        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = pa.Table.from_pydict({
            'pt': [1, 1],
            'k': [10, 20],
            'v': [100, 200]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        table.create_tag("test_tag_2")
        all_tags = set(table.list_tags())
        self.assertEqual(all_tags, {"test_tag", "test_tag_2"})

        # Delete the tag
        result = table.delete_tag("test_tag")
        self.assertTrue(result)

        all_tags = table.list_tags()
        self.assertEqual(all_tags, ["test_tag_2"])

        # Verify tag no longer exists
        self.assertFalse(tag_manager.tag_exists("test_tag"))

    def test_tag_ignore_if_exists(self):
        """Test creating a tag with ignore_if_exists=True."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_tag_ignore_exists', schema, False)
        table = self.catalog.get_table('default.test_tag_ignore_exists')

        write_builder = table.new_batch_write_builder()

        # Write some data
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = pa.Table.from_pydict({
            'pt': [1],
            'k': [10],
            'v': [100]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Create a tag
        table.create_tag("duplicate_tag")

        # Try to create the same tag again without ignore_if_exists - should raise error
        with self.assertRaises(ValueError) as context:
            table.create_tag("duplicate_tag")
        self.assertIn("already exists", str(context.exception))

        # Create the same tag with ignore_if_exists=True - should not raise error
        table.create_tag("duplicate_tag", ignore_if_exists=True)

    def test_tag_rename(self):
        """Test renaming a tag."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_tag_rename', schema, False)
        table = self.catalog.get_table('default.test_tag_rename')

        write_builder = table.new_batch_write_builder()

        # Write some data - snapshot 1
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = pa.Table.from_pydict({
            'pt': [1],
            'k': [10],
            'v': [100]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Create a tag
        table.create_tag("old_tag")

        # Rename the tag
        table.rename_tag("old_tag", "new_tag")

        # Verify old tag no longer exists
        tags = table.list_tags()
        self.assertNotIn("old_tag", tags)
        self.assertIn("new_tag", tags)

        # Verify tag content is preserved
        tag_manager = table.tag_manager()
        tag = tag_manager.get("new_tag")
        self.assertIsNotNone(tag)
        self.assertEqual(tag.id, 1)

    def test_tag_rename_nonexistent(self):
        """Test renaming a nonexistent tag."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_tag_rename_nonexistent', schema, False)
        table = self.catalog.get_table('default.test_tag_rename_nonexistent')

        write_builder = table.new_batch_write_builder()

        # Write some data
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = pa.Table.from_pydict({
            'pt': [1],
            'k': [10],
            'v': [100]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Try to rename nonexistent tag
        with self.assertRaises(ValueError) as context:
            table.rename_tag("nonexistent", "new_tag")
        self.assertIn("doesn't exist", str(context.exception))

    def test_tag_rename_to_existing(self):
        """Test renaming a tag to a name that already exists."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_tag_rename_to_existing', schema, False)
        table = self.catalog.get_table('default.test_tag_rename_to_existing')

        write_builder = table.new_batch_write_builder()

        # Write some data - snapshot 1
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = pa.Table.from_pydict({
            'pt': [1],
            'k': [10],
            'v': [100]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Create two tags
        table.create_tag("tag1")
        table.create_tag("tag2")

        # Try to rename tag1 to tag2 (which already exists)
        with self.assertRaises(ValueError) as context:
            table.rename_tag("tag1", "tag2")
        self.assertIn("already exists", str(context.exception))

        # Verify both original tags still exist
        tags = table.list_tags()
        self.assertIn("tag1", tags)
        self.assertIn("tag2", tags)

    def test_tag_rename_and_scan(self):
        """Test that renaming a tag and then scanning by new name works correctly."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_tag_rename_scan', schema, False)
        table = self.catalog.get_table('default.test_tag_rename_scan')

        write_builder = table.new_batch_write_builder()

        # First commit
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = pa.Table.from_pydict({
            'pt': [1, 1],
            'k': [10, 20],
            'v': [100, 200]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data1)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Second commit
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data2 = pa.Table.from_pydict({
            'pt': [2, 2],
            'k': [30, 40],
            'v': [101, 201]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data2)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Create tag at snapshot 1
        table.create_tag("old_tag", snapshot_id=1)

        # Rename the tag
        table.rename_tag("old_tag", "new_tag")

        # Read from the renamed tag
        table_with_tag = table.copy({CoreOptions.SCAN_TAG_NAME.key(): "new_tag"})
        read_builder = table_with_tag.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        result = table_read.to_arrow(table_scan.plan().splits())

        # Verify: should only contain data from snapshot 1
        result_sorted = result.sort_by([('pt', 'ascending'), ('k', 'ascending')])

        expected = pa.Table.from_pydict({
            'pt': [1, 1],
            'k': [10, 20],
            'v': [100, 200]
        }, schema=self.pk_pa_schema)

        self.assertEqual(result_sorted.num_rows, 2)
        self.assertEqual(result_sorted.column('pt').to_pylist(), expected.column('pt').to_pylist())
        self.assertEqual(result_sorted.column('k').to_pylist(), expected.column('k').to_pylist())
        self.assertEqual(result_sorted.column('v').to_pylist(), expected.column('v').to_pylist())

    def test_schema_evolution_tag_read(self):
        # schema 0
        pa_schema = pa.schema([
            ('user_id', pa.int64()),
            ('item_id', pa.int64()),
            ('dt', pa.string())
        ])
        schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_0', schema, False)
        table1 = self.catalog.get_table('default.test_0')
        write_builder = table1.new_batch_write_builder()
        # write 1
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = {
            'user_id': [1, 2, 3, 4],
            'item_id': [1001, 1002, 1003, 1004],
            'dt': ['p1', 'p1', 'p2', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data1, schema=pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()
        # write 2
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = {
            'user_id': [11, 22, 33, 44],
            'item_id': [1001, 1002, 1003, 1004],
            'dt': ['p1', 'p1', 'p2', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data1, schema=pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()
        # schema 1  add behavior column
        pa_schema = pa.schema([
            ('user_id', pa.int64()),
            ('item_id', pa.int64()),
            ('dt', pa.string()),
            ('behavior', pa.string())
        ])
        schema2 = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_1', schema2, False)
        table2 = self.catalog.get_table('default.test_1')
        write_builder = table2.new_batch_write_builder()
        # write 1
        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],
            'dt': ['p2', 'p1', 'p2', 'p2'],
            'behavior': ['e', 'f', 'g', 'h'],
        }
        pa_table = pa.Table.from_pydict(data2, schema=pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()
        # write 2
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data2 = {
            'user_id': [55, 66, 77, 88],
            'item_id': [1005, 1006, 1007, 1008],
            'dt': ['p2', 'p1', 'p2', 'p2'],
            'behavior': ['e', 'f', 'g', 'h'],
        }
        pa_table = pa.Table.from_pydict(data2, schema=pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Create tag at snapshot 2
        table1.create_tag("tag2", snapshot_id=2)
        table2.create_tag("tag2", snapshot_id=2)
        # When table2 read tag2, it will access table1's schema
        table2.schema_manager = SchemaManager(table2.file_io, table1.table_path)
        table_with_tag = table2.copy({CoreOptions.SCAN_TAG_NAME.key(): "tag2"})
        read_builder = table_with_tag.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()

        result = table_read.to_arrow(table_scan.plan().splits())

        self.assertEqual(result.num_rows, 8)

        expected_schema = pa.schema([
            ('user_id', pa.int64()),
            ('item_id', pa.int64()),
            ('dt', pa.string())
        ])
        expected = pa.Table.from_pydict({
            'user_id': [5, 6, 7, 8, 55, 66, 77, 88],
            'item_id': [1005, 1006, 1007, 1008, 1005, 1006, 1007, 1008],
            'dt': ["p2", "p1", "p2", "p2", "p2", "p1", "p2", "p2"]
        }, schema=expected_schema)

        self.assertEqual(expected, result.sort_by('user_id'))

    def test_table_rollback_to_snapshot(self):
        """Test table-level rollback to a specific snapshot."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_rollback_snapshot', schema, False)
        table = self.catalog.get_table('default.test_rollback_snapshot')

        write_builder = table.new_batch_write_builder()

        # Write 5 commits
        for i in range(5):
            table_write = write_builder.new_write()
            table_commit = write_builder.new_commit()
            data = pa.Table.from_pydict({
                'pt': [1],
                'k': [i],
                'v': [i * 100]
            }, schema=self.pk_pa_schema)
            table_write.write_arrow(data)
            table_commit.commit(table_write.prepare_commit())
            table_write.close()
            table_commit.close()

        snapshot_mgr = table.snapshot_manager()
        self.assertEqual(snapshot_mgr.get_latest_snapshot().id, 5)

        # Rollback to snapshot 3
        table.rollback_to(3)

        self.assertEqual(snapshot_mgr.get_latest_snapshot().id, 3)
        self.assertIsNone(snapshot_mgr.get_snapshot_by_id(4))
        self.assertIsNone(snapshot_mgr.get_snapshot_by_id(5))

    def test_table_rollback_to_tag(self):
        """Test table-level rollback to a specific tag."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_rollback_tag', schema, False)
        table = self.catalog.get_table('default.test_rollback_tag')

        write_builder = table.new_batch_write_builder()

        # Write 5 commits and create tags
        for i in range(5):
            table_write = write_builder.new_write()
            table_commit = write_builder.new_commit()
            data = pa.Table.from_pydict({
                'pt': [1],
                'k': [i],
                'v': [i * 100]
            }, schema=self.pk_pa_schema)
            table_write.write_arrow(data)
            table_commit.commit(table_write.prepare_commit())
            table_write.close()
            table_commit.close()
            table.create_tag("v{}".format(i + 1))

        snapshot_mgr = table.snapshot_manager()
        tag_mgr = table.tag_manager()
        self.assertEqual(snapshot_mgr.get_latest_snapshot().id, 5)

        # Rollback to tag v3
        table.rollback_to("v3")

        self.assertEqual(snapshot_mgr.get_latest_snapshot().id, 3)
        # Tags with snapshot > 3 should be cleaned
        self.assertFalse(tag_mgr.tag_exists("v4"))
        self.assertFalse(tag_mgr.tag_exists("v5"))
        # Tag v3 should still exist
        self.assertTrue(tag_mgr.tag_exists("v3"))

    def test_table_rollback_to_nonexistent_snapshot(self):
        """Test that rollback to a non-existent snapshot raises ValueError."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_rollback_no_snap', schema, False)
        table = self.catalog.get_table('default.test_rollback_no_snap')

        write_builder = table.new_batch_write_builder()

        # Write 2 commits
        for i in range(2):
            table_write = write_builder.new_write()
            table_commit = write_builder.new_commit()
            data = pa.Table.from_pydict({
                'pt': [1],
                'k': [i],
                'v': [i * 100]
            }, schema=self.pk_pa_schema)
            table_write.write_arrow(data)
            table_commit.commit(table_write.prepare_commit())
            table_write.close()
            table_commit.close()

        with self.assertRaises(ValueError) as context:
            table.rollback_to(99)
        self.assertIn("Rollback snapshot", str(context.exception))
        self.assertIn("99", str(context.exception))
        self.assertIn("doesn't exist", str(context.exception))

    def test_table_rollback_to_nonexistent_tag(self):
        """Test that rollback to a non-existent tag raises ValueError."""
        schema = Schema.from_pyarrow_schema(
            self.pa_schema,
            primary_keys=['pt', 'k'],
            partition_keys=['pt'],
            options={'bucket': '3'}
        )
        self.catalog.create_table('default.test_rollback_no_tag', schema, False)
        table = self.catalog.get_table('default.test_rollback_no_tag')

        write_builder = table.new_batch_write_builder()

        # Write 1 commit
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data = pa.Table.from_pydict({
            'pt': [1],
            'k': [10],
            'v': [100]
        }, schema=self.pk_pa_schema)
        table_write.write_arrow(data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        with self.assertRaises(ValueError) as context:
            table.rollback_to("no-such-tag")
        self.assertIn("no-such-tag", str(context.exception))
        self.assertIn("doesn't exist", str(context.exception))
