# 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 sys
import unittest
from types import ModuleType
from unittest.mock import Mock, patch

from pypaimon.catalog.catalog_context import CatalogContext
from pypaimon.catalog.filesystem_catalog_loader import FileSystemCatalogLoader
from pypaimon.catalog.jdbc_catalog_loader import JdbcCatalogLoader
from pypaimon.catalog.rest.rest_catalog_loader import RESTCatalogLoader
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.options.options import Options
from pypaimon.read.native_plan import _catalog_options, native_plan
from pypaimon.read.scan_stats import ScanStats
from pypaimon.read.table_scan import TableScan
from pypaimon.table.bucket_mode import BucketMode


def _scan(native_enabled, file_scanner):
    """Build a TableScan without running its heavy __init__."""
    scan = TableScan.__new__(TableScan)
    scan.table = Mock()
    scan.table.options.native_plan_enabled.return_value = native_enabled
    scan.table.options.options.contains_key.return_value = False   # no time-travel
    scan.table.options.options.contains.return_value = False       # no incremental
    scan.table.options.merge_engine.return_value = None            # not first-row
    scan.table.options.query_auth_enabled = False
    scan.table.current_branch.return_value = 'main'
    scan.table.is_primary_key_table = False        # not a pk table
    scan.table.trimmed_primary_keys = ['k']        # non-empty trimmed pk
    scan.table.bucket_mode.return_value = BucketMode.HASH_FIXED
    scan.table._applied_dynamic_options = {}       # no copy() overrides
    scan.table.partition_keys = []                 # not partitioned
    scan.table.table_schema.id = 1
    scan.table.schema_manager.latest.return_value.id = 1   # loaded schema is latest
    scan.table.identifier.get_database_name.return_value = 'default'
    scan.table.catalog_environment.catalog_loader = FileSystemCatalogLoader(
        CatalogContext.create_from_options(Options({})))           # filesystem catalog
    file_scanner.idx_of_this_subtask = None       # no shard
    file_scanner.start_pos_of_this_subtask = None  # no slice
    file_scanner.chunk_shuffle = None              # no chunk-shuffle
    file_scanner._global_index_result = None       # no global-index result
    file_scanner.deletion_vectors_enabled = False  # no deletion vectors
    file_scanner.data_evolution = False            # no data evolution
    file_scanner.only_read_real_buckets = False    # not postpone bucket
    scan.file_scanner = file_scanner
    scan._query_auth_fn = None      # no query-auth restrictions
    scan._read_type = None
    scan.limit = None               # no row limit
    return scan


class NativePlanTest(unittest.TestCase):

    def setUp(self):
        # Make the real capability probe see a split-API-capable pypaimon-rust so
        # gate tests route natively. Tests that call native_plan() directly override
        # sys.modules within their own block.
        fake_df = ModuleType('pypaimon_rust.datafusion')
        fake_df.PaimonCatalog = type(
            'PaimonCatalog', (), {'get_table': lambda self, name: None})
        fake_df.Split = type('Split', (), {'serialize': lambda self: b''})
        fake_mod = ModuleType('pypaimon_rust')
        fake_mod.datafusion = fake_df
        patcher = patch.dict(
            sys.modules,
            {'pypaimon_rust': fake_mod, 'pypaimon_rust.datafusion': fake_df})
        patcher.start()
        self.addCleanup(patcher.stop)

    def test_switch_defaults_off(self):
        self.assertFalse(CoreOptions(Options({})).native_plan_enabled())
        self.assertTrue(
            CoreOptions(Options({"scan.native-plan.enabled": "true"})).native_plan_enabled())

    def test_plan_uses_file_scanner_when_switch_off(self):
        fs = Mock()
        sentinel = object()
        fs.scan.return_value = sentinel
        scan = _scan(native_enabled=False, file_scanner=fs)
        self.assertIs(scan.plan(), sentinel)
        fs.scan.assert_called_once_with()

    def test_plan_routes_to_native_and_prunes_partitions(self):
        # Native planner returns every partition; the predicate keeps only [2026, 7].
        keep = Mock(partition=Mock(values=[2026, 7]))
        drop = Mock(partition=Mock(values=[2025, 1]))
        pred = Mock()
        pred.test.side_effect = lambda part: part.values == [2026, 7]
        fs = Mock(partition_key_predicate=pred)
        scan = _scan(native_enabled=True, file_scanner=fs)

        with patch('pypaimon.read.native_plan.native_plan', return_value=[keep, drop]) as np:
            plan = scan.plan()

        np.assert_called_once_with(scan.table)
        fs.scan.assert_not_called()
        self.assertEqual(plan.splits(), [keep])

    def test_plan_falls_back_when_partition_prune_raises(self):
        pred = Mock()
        pred.test.side_effect = RuntimeError('predicate boom')
        fs = Mock(partition_key_predicate=pred)
        sentinel = object()
        fs.scan.return_value = sentinel
        scan = _scan(native_enabled=True, file_scanner=fs)
        split = Mock(partition=Mock(values=[2026, 7]), snapshot_id=1)
        with patch('pypaimon.read.native_plan.native_plan', return_value=[split]):
            self.assertIs(scan.plan(), sentinel)
        fs.scan.assert_called_once_with()

    def test_plan_native_no_partition_predicate_keeps_all(self):
        splits = [Mock(partition=Mock(values=[1])), Mock(partition=Mock(values=[2]))]
        fs = Mock(partition_key_predicate=None)
        scan = _scan(native_enabled=True, file_scanner=fs)
        with patch('pypaimon.read.native_plan.native_plan', return_value=splits):
            self.assertEqual(scan.plan().splits(), splits)

    def test_plan_falls_back_when_scan_is_not_plain(self):
        # Native planning does not carry shard/slice, global-index, or
        # time-travel/incremental scans -> must fall back to the file scanner.
        def check(setup):
            fs = Mock(partition_key_predicate=None)
            sentinel = object()
            fs.scan.return_value = sentinel
            scan = _scan(native_enabled=True, file_scanner=fs)
            setup(scan, fs)
            with patch('pypaimon.read.native_plan.native_plan') as np:
                self.assertIs(scan.plan(), sentinel)
            np.assert_not_called()
            fs.scan.assert_called_once_with()

        check(lambda s, fs: setattr(fs, 'idx_of_this_subtask', 0))
        check(lambda s, fs: setattr(fs, 'start_pos_of_this_subtask', 0))
        check(lambda s, fs: setattr(fs, 'chunk_shuffle', (1, 100)))
        check(lambda s, fs: setattr(fs, '_global_index_result', object()))
        check(lambda s, fs: setattr(fs, 'deletion_vectors_enabled', True))
        check(lambda s, fs: setattr(fs, 'data_evolution', True))
        check(lambda s, fs: setattr(fs, 'only_read_real_buckets', True))
        check(lambda s, fs: (setattr(s.table, 'is_primary_key_table', True),
                             setattr(s.table, 'trimmed_primary_keys', [])))
        check(lambda s, fs: s.table.bucket_mode.__setattr__(
            'return_value', BucketMode.HASH_DYNAMIC))
        check(lambda s, fs: s.table.bucket_mode.__setattr__(
            'return_value', BucketMode.CROSS_PARTITION))
        check(lambda s, fs: setattr(
            s.table, '_applied_dynamic_options', {'scan.snapshot-id': None}))
        check(lambda s, fs: setattr(s.table, 'partition_keys', ['dt']))
        check(lambda s, fs: setattr(s.table.schema_manager.latest.return_value, 'id', 2))
        check(lambda s, fs: s.table.schema_manager.latest.__setattr__(
            'side_effect', RuntimeError('metadata read failed')))
        check(lambda s, fs: setattr(s, 'limit', 5))
        check(lambda s, fs: s.table.options.options.contains_key.__setattr__(
            'side_effect', lambda k: k == 'scan.version'))
        check(lambda s, fs: s.table.options.merge_engine.__setattr__(
            'return_value', 'first-row'))
        check(lambda s, fs: setattr(s.table.options, 'query_auth_enabled', True))
        check(lambda s, fs: s.table.current_branch.__setattr__('return_value', 'b1'))
        check(lambda s, fs: s.table.identifier.get_database_name.__setattr__(
            'return_value', 'db.name'))
        check(lambda s, fs: s.table.identifier.get_database_name.__setattr__(
            'return_value', 'unknown'))
        check(lambda s, fs: setattr(
            s.table.catalog_environment, 'catalog_loader', object()))   # no context()
        for attr in ('hadoop_conf', 'prefer_io_loader', 'fallback_io_loader'):
            check(lambda s, fs, attr=attr: setattr(
                s.table.catalog_environment.catalog_loader.context(), attr, object()))
        check(lambda s, fs: s.table.options.options.contains_key.__setattr__(
            'return_value', True))          # time-travel
        check(lambda s, fs: s.table.options.options.contains.__setattr__(
            'return_value', True))          # incremental

    def test_plan_native_empty_falls_back(self):
        # Empty native result -> fall back for an atomic snapshot id.
        fs = Mock(partition_key_predicate=None)
        sentinel = object()
        fs.scan.return_value = sentinel
        scan = _scan(native_enabled=True, file_scanner=fs)
        with patch('pypaimon.read.native_plan.native_plan', return_value=[]):
            self.assertIs(scan.plan(), sentinel)
        fs.scan.assert_called_once_with()

    def test_plan_falls_back_when_rust_unavailable(self):
        # scan.native-plan.enabled but pypaimon-rust missing/old -> fall back,
        # not crash.
        fs = Mock(partition_key_predicate=None)
        sentinel = object()
        fs.scan.return_value = sentinel
        scan = _scan(native_enabled=True, file_scanner=fs)
        with patch('pypaimon.read.native_plan.native_runtime_available',
                   return_value=False), \
                patch('pypaimon.read.native_plan.native_plan') as np:
            self.assertIs(scan.plan(), sentinel)
        np.assert_not_called()
        fs.scan.assert_called_once_with()

    def test_plan_falls_back_when_native_plan_raises(self):
        # A native planning failure (e.g. unsupported scheme) must fall back, not crash.
        fs = Mock(partition_key_predicate=None)
        sentinel = object()
        fs.scan.return_value = sentinel
        scan = _scan(native_enabled=True, file_scanner=fs)
        with patch('pypaimon.read.native_plan.native_plan',
                   side_effect=RuntimeError('unsupported scheme viewfs://')):
            self.assertIs(scan.plan(), sentinel)
        fs.scan.assert_called_once_with()

    def test_plan_falls_back_for_jdbc_catalog_loader(self):
        fs = Mock(partition_key_predicate=None)
        sentinel = object()
        fs.scan.return_value = sentinel
        scan = _scan(native_enabled=True, file_scanner=fs)
        scan.table.catalog_environment.catalog_loader = JdbcCatalogLoader(
            CatalogContext.create_from_options(Options({})))

        with patch('pypaimon.read.native_plan.native_plan') as np:
            self.assertIs(scan.plan(), sentinel)

        np.assert_not_called()
        fs.scan.assert_called_once_with()

    def test_plan_falls_back_for_builtin_catalog_loader_subclasses(self):
        class RoutedFileSystemLoader(FileSystemCatalogLoader):
            def load(self):
                return object()

        class RoutedRESTLoader(RESTCatalogLoader):
            def load(self):
                return object()

        for loader_class in (RoutedFileSystemLoader, RoutedRESTLoader):
            with self.subTest(loader_class=loader_class.__name__):
                fs = Mock(partition_key_predicate=None)
                sentinel = object()
                fs.scan.return_value = sentinel
                scan = _scan(native_enabled=True, file_scanner=fs)
                scan.table.catalog_environment.catalog_loader = loader_class(
                    CatalogContext.create_from_options(Options({})))

                with patch('pypaimon.read.native_plan.native_plan') as np:
                    self.assertIs(scan.plan(), sentinel)

                np.assert_not_called()
                fs.scan.assert_called_once_with()

    def test_scan_with_stats_native_empty_uses_fallback_stats(self):
        fs = Mock(partition_key_predicate=None)
        fallback_plan = object()
        fallback_stats = ScanStats(manifest_files_total=7)
        fs.scan_with_stats.return_value = (fallback_plan, fallback_stats)
        scan = _scan(native_enabled=True, file_scanner=fs)

        with patch('pypaimon.read.native_plan.native_plan', return_value=[]) as np:
            plan, stats = scan.scan_with_stats()

        self.assertIs(plan, fallback_plan)
        self.assertIs(stats, fallback_stats)
        np.assert_called_once_with(scan.table)
        fs.scan_with_stats.assert_called_once_with()
        fs.scan.assert_not_called()

    def test_catalog_options_are_normalized_for_rust(self):
        table = Mock()
        table.catalog_environment.catalog_loader = FileSystemCatalogLoader(
            CatalogContext.create_from_options(Options({
                'warehouse': '/tmp/warehouse',
                'data-token.enabled': True,
                'retry-count': 3,
                'unset': None,
            })))

        self.assertEqual(_catalog_options(table), {
            'warehouse': '/tmp/warehouse',
            # Lowercase so Rust's case-sensitive bool parser accepts it.
            'data-token.enabled': 'true',
            'retry-count': '3',
            'metastore': 'filesystem',
        })

    def test_catalog_options_use_actual_rest_loader_type(self):
        table = Mock()
        table.catalog_environment.catalog_loader = RESTCatalogLoader(
            CatalogContext.create_from_options(Options({
                'uri': 'http://localhost:8181',
            })))

        self.assertEqual(_catalog_options(table), {
            'uri': 'http://localhost:8181',
            'metastore': 'rest',
        })

    def test_catalog_options_reject_loader_subclass(self):
        class RoutedFileSystemLoader(FileSystemCatalogLoader):
            pass

        table = Mock()
        table.catalog_environment.catalog_loader = RoutedFileSystemLoader(
            CatalogContext.create_from_options(Options({})))

        with self.assertRaisesRegex(ValueError, 'exact built-in catalog loader'):
            _catalog_options(table)

    def test_native_plan_threads_trimmed_keys_to_deserializer(self):
        # PK tables route through: the trimmed primary keys must reach the
        # deserializer so per-file min/max keys are decoded for merge-on-read.
        kfields = [object()]
        table = Mock(trimmed_primary_keys_fields=kfields)
        table.table_schema = Mock(fields=[], partition_keys=[])
        table.options.source_split_target_size.return_value = 1024
        table.options.source_split_open_file_cost.return_value = 128
        split = Mock()
        split.serialize.return_value = b'bytes'
        rt = Mock()
        rt.new_read_builder.return_value.new_scan.return_value.plan.return_value \
            .splits.return_value = [split]
        catalog = Mock()
        catalog.get_table.return_value = rt

        fake_df = ModuleType('pypaimon_rust.datafusion')
        fake_df.PaimonCatalog = Mock(return_value=catalog)
        fake_df.Split = type('Split', (), {'serialize': lambda self: b''})
        fake_mod = ModuleType('pypaimon_rust')
        fake_mod.datafusion = fake_df

        with patch.dict(sys.modules,
                        {'pypaimon_rust': fake_mod, 'pypaimon_rust.datafusion': fake_df}), \
                patch('pypaimon.read.native_plan._catalog_options', return_value={}), \
                patch('pypaimon.read.native_plan.deserialize_split_v1',
                      return_value='decoded') as des:
            result = native_plan(table)

        self.assertEqual(result, ['decoded'])
        rt.new_read_builder.assert_called_once_with({
            CoreOptions.SOURCE_SPLIT_TARGET_SIZE.key(): '1024',
            CoreOptions.SOURCE_SPLIT_OPEN_FILE_COST.key(): '128',
        })
        des.assert_called_once_with(b'bytes', [], kfields)

    def test_native_plan_requires_split_api(self):
        # An intermediate pypaimon-rust missing either get_table or Split.serialize
        # must raise a clear error, not an AttributeError mid-plan.
        split_ok = type('Split', (), {'serialize': lambda self: b''})
        catalog_ok = type('PaimonCatalog', (), {'get_table': lambda self, name: None})
        cases = {
            'no get_table': (type('PaimonCatalog', (), {}), split_ok),
            'no serialize': (catalog_ok, type('Split', (), {})),
        }
        for label, (catalog_cls, split_cls) in cases.items():
            with self.subTest(case=label):
                fake_df = ModuleType('pypaimon_rust.datafusion')
                fake_df.PaimonCatalog = catalog_cls
                fake_df.Split = split_cls
                fake_mod = ModuleType('pypaimon_rust')
                fake_mod.datafusion = fake_df
                with patch.dict(
                        sys.modules,
                        {'pypaimon_rust': fake_mod, 'pypaimon_rust.datafusion': fake_df}):
                    with self.assertRaisesRegex(RuntimeError, '0.3.0'):
                        native_plan(Mock())


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