################################################################################
#  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.
################################################################################

"""Tests for LocalDiskCacheManager, CachingInputStream, and CachingFileIO."""

import io
import os
import shutil
import tempfile
import threading
import unittest
from unittest.mock import MagicMock

from pypaimon.filesystem.caching_file_io import (
    LocalDiskCacheManager,
    CachingFileIO,
    CachingInputStream,
)


class LocalDiskCacheManagerTest(unittest.TestCase):

    def setUp(self):
        self.cache_dir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.cache_dir, ignore_errors=True)

    def test_put_and_get(self):
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        data = b"hello world block"
        cache.put_block("file1.index", 0, data)
        result = cache.get_block("file1.index", 0)
        self.assertEqual(data, result)

    def test_cache_miss(self):
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        result = cache.get_block("nonexistent", 0)
        self.assertIsNone(result)

    def test_different_keys(self):
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        cache.put_block("file1", 0, b"block0")
        cache.put_block("file1", 1, b"block1")
        cache.put_block("file2", 0, b"other0")

        self.assertEqual(b"block0", cache.get_block("file1", 0))
        self.assertEqual(b"block1", cache.get_block("file1", 1))
        self.assertEqual(b"other0", cache.get_block("file2", 0))

    def test_duplicate_put_is_noop(self):
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        cache.put_block("file1", 0, b"original")
        cache.put_block("file1", 0, b"duplicate")
        self.assertEqual(b"original", cache.get_block("file1", 0))

    def test_eviction(self):
        cache = LocalDiskCacheManager(self.cache_dir, max_size_bytes=100, block_size=64)
        cache.put_block("f", 0, b"a" * 60)
        cache.put_block("f", 1, b"b" * 60)
        # Total 120 > 100, at least one block should be evicted
        remaining_0 = cache.get_block("f", 0)
        remaining_1 = cache.get_block("f", 1)
        self.assertTrue(remaining_0 is None or remaining_1 is None)

    def test_scan_size_on_restart(self):
        cache1 = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        cache1.put_block("f", 0, b"x" * 100)
        cache1.put_block("f", 1, b"y" * 200)

        # Simulate restart: new cache instance on same directory
        cache2 = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        self.assertEqual(300, cache2._current_size)
        self.assertEqual(b"x" * 100, cache2.get_block("f", 0))
        self.assertEqual(b"y" * 200, cache2.get_block("f", 1))

    def test_cache_dir_created(self):
        new_dir = os.path.join(self.cache_dir, "sub", "deep")
        LocalDiskCacheManager(new_dir, 2 ** 63 - 1, block_size=64)
        self.assertTrue(os.path.isdir(new_dir))

    def test_concurrent_put_get(self):
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        errors = []

        def writer(idx):
            try:
                data = bytes([idx % 256]) * 100
                cache.put_block("concurrent", idx, data)
            except Exception as e:
                errors.append(e)

        def reader(idx):
            try:
                result = cache.get_block("concurrent", idx)
                if result is not None:
                    expected = bytes([idx % 256]) * 100
                    assert result == expected, f"Mismatch at {idx}"
            except Exception as e:
                errors.append(e)

        threads = []
        for i in range(20):
            threads.append(threading.Thread(target=writer, args=(i,)))
            threads.append(threading.Thread(target=reader, args=(i,)))
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        self.assertEqual([], errors)

    def test_unlimited_cache_skips_eviction(self):
        cache = LocalDiskCacheManager(self.cache_dir, max_size_bytes=2 ** 63 - 1, block_size=64)
        for i in range(50):
            cache.put_block("f", i, b"x" * 100)
        # All 50 blocks should still be present
        for i in range(50):
            self.assertIsNotNone(cache.get_block("f", i))


class CachingInputStreamTest(unittest.TestCase):

    def setUp(self):
        self.cache_dir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.cache_dir, ignore_errors=True)

    def _make_file_io(self, data, path="test"):
        file_io = MagicMock()
        file_io.new_input_stream.side_effect = lambda p: io.BytesIO(data)
        file_io.get_file_size.side_effect = lambda p: len(data)
        return file_io

    def test_read_entire_file(self):
        data = b"abcdefghijklmnop"
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)
        stream = CachingInputStream(self._make_file_io(data), "test", cache)
        result = stream.read()
        self.assertEqual(data, result)

    def test_read_with_size(self):
        data = b"abcdefghijklmnop"
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)
        stream = CachingInputStream(self._make_file_io(data), "test", cache)
        self.assertEqual(b"abcde", stream.read(5))
        self.assertEqual(b"fghij", stream.read(5))
        self.assertEqual(b"klmnop", stream.read(10))

    def test_seek_and_read(self):
        data = b"0123456789abcdef"
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)
        stream = CachingInputStream(self._make_file_io(data), "test", cache)
        stream.seek(10)
        self.assertEqual(b"abcdef", stream.read())

    def test_seek_whence_1(self):
        data = b"0123456789"
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)
        stream = CachingInputStream(self._make_file_io(data), "test", cache)
        stream.read(3)
        stream.seek(2, 1)  # relative
        self.assertEqual(5, stream.tell())
        self.assertEqual(b"56789", stream.read())

    def test_seek_whence_2(self):
        data = b"0123456789"
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)
        stream = CachingInputStream(self._make_file_io(data), "test", cache)
        stream.seek(-3, 2)  # from end
        self.assertEqual(b"789", stream.read())

    def test_read_spanning_multiple_blocks(self):
        data = bytes(range(256)) * 4  # 1024 bytes
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=100)
        stream = CachingInputStream(self._make_file_io(data), "test", cache)
        # Read across block boundary: block 0 ends at 100, block 1 starts at 100
        stream.seek(90)
        result = stream.read(30)  # 90..120, spans blocks 0 and 1
        self.assertEqual(data[90:120], result)

    def test_cache_hit_avoids_remote_read(self):
        data = b"0123456789abcdef"
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)

        # First read populates cache
        file_io1 = self._make_file_io(data)
        stream1 = CachingInputStream(file_io1, "test", cache)
        stream1.read()

        # Second stream: remote should not be called because all blocks are cached
        file_io2 = MagicMock()
        file_io2.get_file_size.side_effect = lambda p: len(data)
        stream2 = CachingInputStream(file_io2, "test", cache)
        result = stream2.read()
        self.assertEqual(data, result)
        file_io2.new_input_stream.assert_not_called()

    def test_read_empty(self):
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)
        stream = CachingInputStream(self._make_file_io(b""), "test", cache)
        self.assertEqual(b"", stream.read())

    def test_read_at_eof(self):
        data = b"abc"
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)
        stream = CachingInputStream(self._make_file_io(data), "test", cache)
        stream.read()
        self.assertEqual(b"", stream.read(10))

    def test_context_manager(self):
        data = b"test"
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)
        file_io = self._make_file_io(data)
        with CachingInputStream(file_io, "test", cache):
            pass

    def test_partial_last_block(self):
        data = b"abcdefghij"  # 10 bytes, block_size=8 -> block0=8, block1=2
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=8)
        stream = CachingInputStream(self._make_file_io(data), "test", cache)
        result = stream.read()
        self.assertEqual(data, result)
        # Verify block 1 is only 2 bytes
        self.assertEqual(b"ij", cache.get_block("test", 1))


class CachingFileIOTest(unittest.TestCase):

    def setUp(self):
        self.cache_dir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.cache_dir, ignore_errors=True)

    def _make_delegate(self, path_to_data):
        delegate = MagicMock()

        def new_input_stream(path):
            return io.BytesIO(path_to_data.get(path, b""))

        def get_file_size(path):
            return len(path_to_data.get(path, b""))

        delegate.new_input_stream.side_effect = new_input_stream
        delegate.get_file_size.side_effect = get_file_size
        return delegate

    def test_meta_file_is_cached(self):
        data = b"snapshot data"
        delegate = self._make_delegate({"snapshot-1": data})
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache)

        # First read
        with caching_io.new_input_stream("snapshot-1") as s:
            result1 = s.read()
        self.assertEqual(data, result1)

        # Second read should hit cache (returns CachingInputStream)
        with caching_io.new_input_stream("snapshot-1") as s:
            result2 = s.read()
        self.assertEqual(data, result2)
        # get_file_size called only once — second stream gets size from cache
        self.assertEqual(1, delegate.get_file_size.call_count)

    def test_manifest_file_is_cached(self):
        data = b"manifest data"
        delegate = self._make_delegate({"manifest-abc": data})
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache)

        with caching_io.new_input_stream("manifest-abc") as s:
            self.assertEqual(data, s.read())

    def test_global_index_file_is_cached(self):
        data = b"index data"
        delegate = self._make_delegate({"global-index-uuid.index": data})
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache)

        with caching_io.new_input_stream("global-index-uuid.index") as s:
            self.assertEqual(data, s.read())

    def test_bucket_index_file_not_cached_by_default(self):
        data = b"bucket index"
        delegate = self._make_delegate({"index-uuid-0": data})
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache)

        result = caching_io.new_input_stream("index-uuid-0")
        self.assertNotIsInstance(result, CachingInputStream)
        self.assertEqual(data, result.read())
        delegate.get_file_size.assert_not_called()

    def test_bucket_index_cached_when_in_whitelist(self):
        from pypaimon.utils.file_type import FileType
        data = b"bucket index"
        delegate = self._make_delegate({"index-uuid-0": data})
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        whitelist = {FileType.META, FileType.GLOBAL_INDEX, FileType.BUCKET_INDEX}
        caching_io = CachingFileIO(delegate, cache, whitelist)

        with caching_io.new_input_stream("index-uuid-0") as s:
            self.assertEqual(data, s.read())

    def test_custom_whitelist_meta_only(self):
        from pypaimon.utils.file_type import FileType
        delegate = self._make_delegate({
            "snapshot-1": b"snap",
            "global-index-uuid.index": b"idx",
        })
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache, {FileType.META})

        with caching_io.new_input_stream("snapshot-1") as s:
            self.assertIsInstance(s, CachingInputStream)

        result = caching_io.new_input_stream("global-index-uuid.index")
        self.assertNotIsInstance(result, CachingInputStream)

    def test_data_file_not_cached(self):
        data = b"data content"
        delegate = self._make_delegate({"data-abc.orc": data})
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache)

        result = caching_io.new_input_stream("data-abc.orc")
        # Should be raw BytesIO, not CachingInputStream
        self.assertNotIsInstance(result, CachingInputStream)
        self.assertEqual(data, result.read())
        # get_file_size should NOT be called for data files
        delegate.get_file_size.assert_not_called()

    def test_file_index_not_cached(self):
        data = b"file index content"
        delegate = self._make_delegate({"data-abc.orc.index": data})
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache)

        result = caching_io.new_input_stream("data-abc.orc.index")
        self.assertNotIsInstance(result, CachingInputStream)
        self.assertEqual(data, result.read())
        delegate.get_file_size.assert_not_called()

    def test_delegate_methods_forwarded(self):
        delegate = MagicMock()
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache)

        caching_io.exists("/some/path")
        delegate.exists.assert_called_once_with("/some/path")

        caching_io.mkdirs("/some/dir")
        delegate.mkdirs.assert_called_once_with("/some/dir")

        caching_io.delete("/some/file", recursive=True)
        delegate.delete.assert_called_once_with("/some/file", True)

        caching_io.rename("/src", "/dst")
        delegate.rename.assert_called_once_with("/src", "/dst")

        caching_io.get_file_status("/path")
        delegate.get_file_status.assert_called_once_with("/path")

        caching_io.list_status("/dir")
        delegate.list_status.assert_called_once_with("/dir")

        caching_io.new_output_stream("/out")
        delegate.new_output_stream.assert_called_once_with("/out")

    def test_to_filesystem_path_forwarded_to_delegate(self):
        delegate = MagicMock()
        delegate.to_filesystem_path.side_effect = lambda p: p.replace("oss://", "")
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache)

        self.assertEqual("bucket/key", caching_io.to_filesystem_path("oss://bucket/key"))
        delegate.to_filesystem_path.assert_called_once_with("oss://bucket/key")

    def test_write_parquet_when_enable_local_cache(self):
        import pyarrow as pa
        from pypaimon.filesystem.local_file_io import LocalFileIO

        delegate = LocalFileIO()
        cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1, block_size=64)
        caching_io = CachingFileIO(delegate, cache)

        out_path = os.path.join(self.cache_dir, "data.parquet")
        table = pa.table({"a": [1, 2, 3]})
        caching_io.write_parquet(out_path, table)

        self.assertTrue(os.path.exists(out_path))
        self.assertEqual(table, pa.parquet.read_table(out_path))

    def test_file_io_for_data_wraps_cache_when_data_token_enabled(self):
        from pypaimon.catalog.rest.rest_catalog import RESTCatalog
        from pypaimon.common.identifier import Identifier
        from pypaimon.common.options.options import Options

        catalog = MagicMock(spec=RESTCatalog)
        catalog.data_token_enabled = True
        catalog.fuse_enabled = False
        catalog._fuse_resolver = None
        catalog.context = MagicMock()
        catalog.context.options = Options({
            'local-cache.enabled': 'true',
            'local-cache.dir': self.cache_dir,
            'local-cache.whitelist': 'meta,global-index,data',
        })
        catalog._cache_manager = CachingFileIO.create_cache_manager(
            catalog.context.options)
        catalog.file_io_for_data = RESTCatalog.file_io_for_data.__get__(
            catalog, RESTCatalog)

        file_io = catalog.file_io_for_data(
            "oss://catalog/db1/table1", Identifier.create("db1", "table1"))

        self.assertIsInstance(
            file_io, CachingFileIO,
            msg="Cache wrap should apply even when data-token.enabled=true; "
                "currently bypassed in DLF mode (RESTTokenFileIO returned).")

    def test_default_memory_cache_max_size_capped(self):
        from pypaimon.common.options.options import Options
        from pypaimon.filesystem.caching_file_io import LocalMemoryCacheManager

        cache = CachingFileIO.create_cache_manager(Options({
            'local-cache.enabled': 'true',
        }))
        self.assertIsInstance(cache, LocalMemoryCacheManager)
        self.assertEqual(
            cache._max_size_bytes, 256 * 1024 * 1024,
            msg="Memory cache without explicit max-size should default to "
                "256 MB, not unlimited (OOM risk).")

    def test_default_disk_cache_max_size_capped(self):
        from pypaimon.common.options.options import Options
        from pypaimon.filesystem.caching_file_io import LocalDiskCacheManager

        cache = CachingFileIO.create_cache_manager(Options({
            'local-cache.enabled': 'true',
            'local-cache.dir': self.cache_dir,
        }))
        self.assertIsInstance(cache, LocalDiskCacheManager)
        self.assertEqual(
            cache._max_size_bytes, 10 * 1024 * 1024 * 1024,
            msg="Disk cache without explicit max-size should default to "
                "10 GB, not unlimited (disk-full risk).")


class ConfigOptionsTest(unittest.TestCase):

    def test_local_cache_options_defaults(self):
        from pypaimon.common.options import Options
        from pypaimon.common.options.core_options import CoreOptions

        opts = CoreOptions(Options({}))
        self.assertFalse(opts.local_cache_enabled())
        self.assertIsNone(opts.local_cache_dir())
        self.assertIsNone(opts.local_cache_max_size())
        self.assertEqual(1 * 1024 * 1024, opts.local_cache_block_size().get_bytes())
        self.assertEqual("meta,global-index", opts.local_cache_whitelist())

    def test_local_cache_options_custom(self):
        from pypaimon.common.options import Options
        from pypaimon.common.options.core_options import CoreOptions

        opts = CoreOptions(Options({
            "local-cache.enabled": "true",
            "local-cache.dir": "/custom/cache",
            "local-cache.max-size": "2gb",
            "local-cache.block-size": "4mb",
            "local-cache.whitelist": "meta,global-index,bucket-index",
        }))
        self.assertTrue(opts.local_cache_enabled())
        self.assertEqual("/custom/cache", opts.local_cache_dir())
        self.assertEqual(2 * 1024 * 1024 * 1024, opts.local_cache_max_size().get_bytes())
        self.assertEqual(4 * 1024 * 1024, opts.local_cache_block_size().get_bytes())
        self.assertEqual("meta,global-index,bucket-index", opts.local_cache_whitelist())


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