# 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 unittest
import uuid
from datetime import datetime
from unittest.mock import MagicMock, Mock, patch

from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.manifest.schema.manifest_entry import ManifestEntry
from pypaimon.snapshot.snapshot_commit import PartitionStatistics
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.write.commit.row_id_conflict_rewriter import RowIdRewriteResult
from pypaimon.write.commit_message import CommitMessage
from pypaimon.write.file_store_commit import (
    FileStoreCommit,
    RewriteResult,
)


@patch('pypaimon.write.file_store_commit.ManifestFileManager')
@patch('pypaimon.write.file_store_commit.ManifestListManager')
class TestFileStoreCommit(unittest.TestCase):
    """Test cases for FileStoreCommit class."""

    def setUp(self):
        """Set up test fixtures."""
        # Mock table with required attributes
        self.mock_table = Mock()
        self.mock_table.partition_keys = ['dt', 'region']
        self.mock_table.current_branch.return_value = 'main'
        self.mock_table.table_path = '/test/table/path'
        self.mock_table.file_io = Mock()
        self.mock_table.options.manifest_target_size.return_value = 8 * 1024 * 1024
        self.mock_table.options.manifest_merge_min_count.return_value = 30

        # Mock snapshot commit
        self.mock_snapshot_commit = Mock()

    def _create_file_store_commit(self):
        """Helper method to create FileStoreCommit instance."""
        return FileStoreCommit(
            snapshot_commit=self.mock_snapshot_commit,
            table=self.mock_table,
            commit_user='test_user'
        )

    def test_generate_partition_statistics_single_partition_single_file(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        """Test partition statistics generation with single partition and single file."""
        # Create FileStoreCommit instance
        file_store_commit = self._create_file_store_commit()

        # Create test data
        creation_time_dt = datetime(2024, 1, 15, 10, 30, 0)
        from pypaimon.data.timestamp import Timestamp
        from pypaimon.table.row.generic_row import GenericRow
        from pypaimon.manifest.schema.simple_stats import SimpleStats
        creation_time = Timestamp.from_local_date_time(creation_time_dt)
        file_meta = DataFileMeta.create(
            file_name="test_file_1.parquet",
            file_size=1024 * 1024,  # 1MB
            row_count=10000,
            min_key=GenericRow([], []),
            max_key=GenericRow([], []),
            key_stats=SimpleStats.empty_stats(),
            value_stats=SimpleStats.empty_stats(),
            min_sequence_number=1,
            max_sequence_number=100,
            schema_id=0,
            level=0,
            extra_files=[],
            creation_time=creation_time,
            external_path=None,
            first_row_id=None,
            write_cols=None
        )

        commit_message = CommitMessage(
            partition=('2024-01-15', 'us-east-1'),
            bucket=0,
            new_files=[file_meta]
        )

        # Test method
        statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message]))

        # Verify results
        self.assertEqual(len(statistics), 1)

        stat = statistics[0]
        self.assertIsInstance(stat, PartitionStatistics)
        self.assertEqual(stat.spec, {'dt': '2024-01-15', 'region': 'us-east-1'})
        self.assertEqual(stat.record_count, 10000)
        self.assertEqual(stat.file_count, 1)
        self.assertEqual(stat.file_size_in_bytes, 1024 * 1024)
        expected_time = file_meta.creation_time_epoch_millis()
        self.assertEqual(stat.last_file_creation_time, expected_time)

    def test_generate_partition_statistics_multiple_files_same_partition(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        """Test partition statistics generation with multiple files in same partition."""
        # Create FileStoreCommit instance
        file_store_commit = self._create_file_store_commit()

        from pypaimon.data.timestamp import Timestamp
        from pypaimon.table.row.generic_row import GenericRow
        from pypaimon.manifest.schema.simple_stats import SimpleStats
        creation_time_1 = Timestamp.from_local_date_time(datetime(2024, 1, 15, 10, 30, 0))
        creation_time_2 = Timestamp.from_local_date_time(datetime(2024, 1, 15, 11, 30, 0))  # Later time

        file_meta_1 = DataFileMeta.create(
            file_name="test_file_1.parquet",
            file_size=1024 * 1024,  # 1MB
            row_count=10000,
            min_key=GenericRow([], []),
            max_key=GenericRow([], []),
            key_stats=SimpleStats.empty_stats(),
            value_stats=SimpleStats.empty_stats(),
            min_sequence_number=1,
            max_sequence_number=100,
            schema_id=0,
            level=0,
            extra_files=[],
            creation_time=creation_time_1
        )

        file_meta_2 = DataFileMeta.create(
            file_name="test_file_2.parquet",
            file_size=2 * 1024 * 1024,  # 2MB
            row_count=15000,
            min_key=GenericRow([], []),
            max_key=GenericRow([], []),
            key_stats=SimpleStats.empty_stats(),
            value_stats=SimpleStats.empty_stats(),
            min_sequence_number=101,
            max_sequence_number=200,
            schema_id=0,
            level=0,
            extra_files=[],
            creation_time=creation_time_2
        )

        commit_message = CommitMessage(
            partition=('2024-01-15', 'us-east-1'),
            bucket=0,
            new_files=[file_meta_1, file_meta_2]
        )

        # Test method
        statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message]))

        # Verify results
        self.assertEqual(len(statistics), 1)

        stat = statistics[0]
        self.assertEqual(stat.spec, {'dt': '2024-01-15', 'region': 'us-east-1'})
        self.assertEqual(stat.record_count, 25000)  # 10000 + 15000
        self.assertEqual(stat.file_count, 2)
        self.assertEqual(stat.file_size_in_bytes, 3 * 1024 * 1024)  # 1MB + 2MB
        expected_time = file_meta_2.creation_time_epoch_millis()
        self.assertEqual(stat.last_file_creation_time, expected_time)

    def test_generate_partition_statistics_multiple_partitions(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        """Test partition statistics generation with multiple different partitions."""
        # Create FileStoreCommit instance
        file_store_commit = self._create_file_store_commit()

        creation_time_dt = datetime(2024, 1, 15, 10, 30, 0)
        from pypaimon.data.timestamp import Timestamp
        from pypaimon.table.row.generic_row import GenericRow
        from pypaimon.manifest.schema.simple_stats import SimpleStats
        creation_time = Timestamp.from_local_date_time(creation_time_dt)

        # File for partition 1
        file_meta_1 = DataFileMeta.create(
            file_name="test_file_1.parquet",
            file_size=1024 * 1024,
            row_count=10000,
            min_key=GenericRow([], []),
            max_key=GenericRow([], []),
            key_stats=SimpleStats.empty_stats(),
            value_stats=SimpleStats.empty_stats(),
            min_sequence_number=1,
            max_sequence_number=100,
            schema_id=0,
            level=0,
            extra_files=[],
            creation_time=creation_time,
            external_path=None,
            first_row_id=None,
            write_cols=None
        )

        # File for partition 2
        file_meta_2 = DataFileMeta.create(
            file_name="test_file_2.parquet",
            file_size=2 * 1024 * 1024,
            row_count=20000,
            min_key=GenericRow([], []),
            max_key=GenericRow([], []),
            key_stats=SimpleStats.empty_stats(),
            value_stats=SimpleStats.empty_stats(),
            min_sequence_number=101,
            max_sequence_number=200,
            schema_id=0,
            level=0,
            extra_files=[],
            creation_time=creation_time,
            external_path=None,
            first_row_id=None,
            write_cols=None
        )

        commit_message_1 = CommitMessage(
            partition=('2024-01-15', 'us-east-1'),
            bucket=0,
            new_files=[file_meta_1]
        )

        commit_message_2 = CommitMessage(
            partition=('2024-01-15', 'us-west-2'),
            bucket=0,
            new_files=[file_meta_2]
        )

        # Test method
        statistics = file_store_commit._generate_partition_statistics(
            self._to_entries([commit_message_1, commit_message_2]))

        # Verify results
        self.assertEqual(len(statistics), 2)

        # Sort statistics by partition spec for consistent testing
        statistics.sort(key=lambda s: s.spec['region'])

        # Check first partition (us-east-1)
        stat_1 = statistics[0]
        self.assertEqual(stat_1.spec, {'dt': '2024-01-15', 'region': 'us-east-1'})
        self.assertEqual(stat_1.record_count, 10000)
        self.assertEqual(stat_1.file_count, 1)
        self.assertEqual(stat_1.file_size_in_bytes, 1024 * 1024)

        # Check second partition (us-west-2)
        stat_2 = statistics[1]
        self.assertEqual(stat_2.spec, {'dt': '2024-01-15', 'region': 'us-west-2'})
        self.assertEqual(stat_2.record_count, 20000)
        self.assertEqual(stat_2.file_count, 1)
        self.assertEqual(stat_2.file_size_in_bytes, 2 * 1024 * 1024)

    def test_generate_partition_statistics_unpartitioned_table(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        """Test partition statistics generation for unpartitioned table."""
        # Update mock table to have no partition keys
        self.mock_table.partition_keys = []

        # Create FileStoreCommit instance
        file_store_commit = self._create_file_store_commit()

        creation_time_dt = datetime(2024, 1, 15, 10, 30, 0)
        from pypaimon.data.timestamp import Timestamp
        from pypaimon.table.row.generic_row import GenericRow
        from pypaimon.manifest.schema.simple_stats import SimpleStats
        creation_time = Timestamp.from_local_date_time(creation_time_dt)
        file_meta = DataFileMeta.create(
            file_name="test_file_1.parquet",
            file_size=1024 * 1024,
            row_count=10000,
            min_key=GenericRow([], []),
            max_key=GenericRow([], []),
            key_stats=SimpleStats.empty_stats(),
            value_stats=SimpleStats.empty_stats(),
            min_sequence_number=1,
            max_sequence_number=100,
            schema_id=0,
            level=0,
            extra_files=[],
            creation_time=creation_time,
            external_path=None,
            first_row_id=None,
            write_cols=None
        )

        commit_message = CommitMessage(
            partition=(),  # Empty partition for unpartitioned table
            bucket=0,
            new_files=[file_meta]
        )

        # Test method
        statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message]))

        # Verify results
        self.assertEqual(len(statistics), 1)

        stat = statistics[0]
        self.assertEqual(stat.spec, {})  # Empty spec for unpartitioned table
        self.assertEqual(stat.record_count, 10000)
        self.assertEqual(stat.file_count, 1)
        self.assertEqual(stat.file_size_in_bytes, 1024 * 1024)

    def test_generate_partition_statistics_no_creation_time(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        """Test partition statistics generation when file has no creation time."""
        # Create FileStoreCommit instance
        file_store_commit = self._create_file_store_commit()

        file_meta = DataFileMeta(
            file_name="test_file_1.parquet",
            file_size=1024 * 1024,
            row_count=10000,
            min_key=None,
            max_key=None,
            key_stats=None,
            value_stats=None,
            min_sequence_number=1,
            max_sequence_number=100,
            schema_id=0,
            level=0,
            extra_files=None,
        )

        commit_message = CommitMessage(
            partition=('2024-01-15', 'us-east-1'),
            bucket=0,
            new_files=[file_meta]
        )

        # Test method
        statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message]))

        # Verify results
        self.assertEqual(len(statistics), 1)

        stat = statistics[0]
        # Should have a valid timestamp (current time)
        self.assertGreater(stat.last_file_creation_time, 0)

    def test_generate_partition_statistics_mismatched_partition_keys(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        """Test partition statistics generation when partition tuple doesn't match partition keys."""
        # Create FileStoreCommit instance
        file_store_commit = self._create_file_store_commit()

        # Table has 2 partition keys but partition tuple has 3 values
        from pypaimon.data.timestamp import Timestamp
        from pypaimon.table.row.generic_row import GenericRow
        from pypaimon.manifest.schema.simple_stats import SimpleStats
        file_meta = DataFileMeta.create(
            file_name="test_file_1.parquet",
            file_size=1024 * 1024,
            row_count=10000,
            min_key=GenericRow([], []),
            max_key=GenericRow([], []),
            key_stats=SimpleStats.empty_stats(),
            value_stats=SimpleStats.empty_stats(),
            min_sequence_number=1,
            max_sequence_number=100,
            schema_id=0,
            level=0,
            extra_files=[],
            creation_time=Timestamp.from_local_date_time(datetime(2024, 1, 15, 10, 30, 0))
        )

        commit_message = CommitMessage(
            partition=('2024-01-15', 'us-east-1', 'extra-value'),  # 3 values but table has 2 keys
            bucket=0,
            new_files=[file_meta]
        )

        # Test method
        statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message]))

        # Verify results - should fallback to index-based naming
        self.assertEqual(len(statistics), 1)

        stat = statistics[0]
        expected_spec = {
            'partition_0': '2024-01-15',
            'partition_1': 'us-east-1',
            'partition_2': 'extra-value'
        }
        self.assertEqual(stat.spec, expected_spec)

    def test_generate_partition_statistics_empty_commit_messages(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        """Test partition statistics generation with empty commit messages list."""
        # Create FileStoreCommit instance
        file_store_commit = self._create_file_store_commit()

        # Test method
        statistics = file_store_commit._generate_partition_statistics([])

        # Verify results
        self.assertEqual(len(statistics), 0)

    def test_append_commit_inherits_index_manifest(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        file_store_commit = self._create_file_store_commit()

        self.mock_table.identifier = 'default.test_table'
        self.mock_table.table_schema = Mock()
        self.mock_table.table_schema.id = 7
        self.mock_table.options.row_tracking_enabled.return_value = False

        snapshot_commit = MagicMock()
        snapshot_commit.__enter__.return_value = snapshot_commit
        snapshot_commit.__exit__.return_value = False
        snapshot_commit.commit.return_value = True
        file_store_commit.snapshot_commit = snapshot_commit

        file_store_commit._write_manifest_files = Mock(return_value=[Mock()])
        file_store_commit._generate_partition_statistics = Mock(return_value=[])
        file_store_commit.manifest_list_manager.read_all.return_value = []

        latest_snapshot = Mock()
        latest_snapshot.id = 3
        latest_snapshot.uuid = "base-snapshot-uuid"
        latest_snapshot.total_record_count = 10
        latest_snapshot.index_manifest = "index-manifest-existing"

        commit_entry = Mock()
        commit_entry.kind = 0
        commit_entry.file = Mock()
        commit_entry.file.row_count = 2

        result = file_store_commit._try_commit_once(
            retry_result=None,
            commit_kind="APPEND",
            commit_entries=[commit_entry],
            changelog_entries=[],
            commit_identifier=11,
            latest_snapshot=latest_snapshot
        )

        self.assertTrue(result.is_success())
        self.assertEqual(
            "base-snapshot-uuid",
            snapshot_commit.commit.call_args[0][0],
        )
        committed_snapshot = snapshot_commit.commit.call_args[0][1]
        self.assertEqual(
            "index-manifest-existing",
            committed_snapshot.index_manifest
        )
        self.assertEqual(str(uuid.UUID(committed_snapshot.uuid)), committed_snapshot.uuid)

    def test_null_partition_value(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        from pypaimon.data.timestamp import Timestamp
        from pypaimon.manifest.schema.simple_stats import SimpleStats
        from pypaimon.schema.data_types import DataField, AtomicType

        file_store_commit = self._create_file_store_commit()
        self.mock_table.partition_keys = ['dt']
        self.mock_table.partition_keys_fields = [
            DataField(0, 'dt', AtomicType('STRING'))
        ]
        self.mock_table.table_schema = Mock()
        self.mock_table.table_schema.id = 0

        file_store_commit.manifest_file_manager = Mock()
        file_store_commit.manifest_file_manager.manifest_path = '/test/manifest'
        self.mock_table.file_io.get_file_size.return_value = 1024

        creation_time = Timestamp.from_local_date_time(datetime(2024, 1, 15, 10, 30, 0))

        def make_file(name):
            return DataFileMeta.create(
                file_name=name,
                file_size=1024,
                row_count=100,
                min_key=GenericRow([], []),
                max_key=GenericRow([], []),
                key_stats=SimpleStats.empty_stats(),
                value_stats=SimpleStats.empty_stats(),
                min_sequence_number=1,
                max_sequence_number=10,
                schema_id=0,
                level=0,
                extra_files=[],
                creation_time=creation_time,
            )

        entries = [
            ManifestEntry(kind=0, partition=GenericRow([None], None), bucket=0, total_buckets=None,
                          file=make_file("f1.parquet")),
            ManifestEntry(kind=0, partition=GenericRow(['2024-01-15'], None), bucket=0, total_buckets=None,
                          file=make_file("f2.parquet")),
        ]

        result = file_store_commit._write_manifest_files(entries, "manifest-test")
        self.assertIsNotNone(result)

    def test_row_id_rewrite_respects_commit_retry_limit(
            self, mock_manifest_list_manager, mock_manifest_file_manager):
        file_store_commit = self._create_file_store_commit()
        file_store_commit.commit_max_retries = 1
        file_store_commit.commit_timeout = 10 ** 9
        file_store_commit._commit_retry_wait = Mock()

        latest_snapshot = Mock()
        latest_snapshot.id = 7
        file_store_commit.snapshot_manager.get_latest_snapshot.return_value = (
            latest_snapshot
        )

        commit_entry = Mock()
        rewrite = RewriteResult(RowIdRewriteResult([commit_entry], 1))
        file_store_commit._try_commit_once = Mock(side_effect=[
            rewrite,
            rewrite,
            AssertionError("rewrite retry budget was not enforced"),
        ])

        with self.assertRaises(RuntimeError) as ctx:
            file_store_commit._try_commit(
                commit_kind="APPEND",
                commit_identifier=11,
                commit_entries_plan=lambda snapshot: [commit_entry],
            )

        self.assertIn("with 1 retries", str(ctx.exception))
        self.assertEqual(2, file_store_commit._try_commit_once.call_count)
        file_store_commit._commit_retry_wait.assert_called_once_with(0)

    @staticmethod
    def _to_entries(commit_messages):
        commit_entries = []
        for msg in commit_messages:
            partition = GenericRow(list(msg.partition), None)
            for file in msg.new_files:
                commit_entries.append(ManifestEntry(
                    kind=0,
                    partition=partition,
                    bucket=msg.bucket,
                    total_buckets=None,
                    file=file
                ))
        return commit_entries
