# 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 gc
import os
import tempfile
import threading
import time
import unittest
import weakref
from concurrent.futures import Future, ThreadPoolExecutor
from unittest.mock import patch

import pyarrow as pa
import pyarrow.fs as pafs
import pyarrow.parquet as pq
from fsspec.implementations.local import LocalFileSystem as FsspecLocalFileSystem

from pypaimon.common.options import Options
from pypaimon.common.options.config import CatalogOptions
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.filesystem.local_file_io import LocalFileIO
from pypaimon.read.reader import format_pyarrow_reader as reader_module
from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader
from pypaimon.schema.data_types import AtomicType, DataField


DEFAULT_CACHE_SIZE = 50 * 1024 * 1024


class _CountingInputFile:
    def __init__(self, wrapped, file_system):
        self._wrapped = wrapped
        self._file_system = file_system

    def read(self, size=-1):
        offset = self._wrapped.tell()
        data = self._wrapped.read(size)
        self._file_system.reads.append((offset, len(data)))
        return data

    def readinto(self, buffer):
        offset = self._wrapped.tell()
        size = self._wrapped.readinto(buffer)
        self._file_system.reads.append((offset, size))
        return size

    def __getattr__(self, name):
        return getattr(self._wrapped, name)


class _CountingLocalFileSystem(FsspecLocalFileSystem):
    def __init__(self):
        super().__init__()
        self.opens = 0
        self.reads = []

    def _open(self, path, mode="rb", **kwargs):
        wrapped = super()._open(path, mode=mode, **kwargs)
        if "r" not in mode:
            return wrapped
        self.opens += 1
        return _CountingInputFile(wrapped, self)

    def reset_counts(self):
        self.opens = 0
        self.reads = []


class FileFormatMetadataCacheTest(unittest.TestCase):
    def setUp(self):
        reader_module._reset_file_format_dataset_cache()
        self.temp_dir = tempfile.TemporaryDirectory()
        self.file_io = LocalFileIO(self.temp_dir.name, Options({}))
        self.paths = []
        for index in range(3):
            path = os.path.join(self.temp_dir.name, "data-{}.parquet".format(index))
            pq.write_table(
                pa.table({"value": list(range(index * 10, index * 10 + 10))}),
                path,
                row_group_size=2,
            )
            self.paths.append(path)

    def tearDown(self):
        reader_module._reset_file_format_dataset_cache()
        self.temp_dir.cleanup()

    def _file_io(self, max_size="50 mb"):
        return LocalFileIO(self.temp_dir.name, Options({
            "file-format.metadata-cache.max-size": max_size,
        }))

    def _read(self, path, file_io=None, options=None):
        reader = FormatPyArrowReader(
            file_io or self.file_io,
            "parquet",
            path,
            [DataField(0, "value", AtomicType("BIGINT"))],
            None,
            options=options,
        )
        values = []
        try:
            while True:
                batch = reader.read_arrow_batch()
                if batch is None:
                    return values
                values.extend(batch.column(0).to_pylist())
        finally:
            reader.close()

    def test_enabled_by_default(self):
        self.assertEqual(
            DEFAULT_CACHE_SIZE,
            self.file_io.properties.get(
                CatalogOptions.FILE_FORMAT_METADATA_CACHE_MAX_SIZE).get_bytes())

        original = reader_module.ds.dataset
        with patch.object(reader_module.ds, "dataset", wraps=original) as dataset:
            self._read(self.paths[0])
            self._read(self.paths[0])
        self.assertEqual(1, dataset.call_count)

    def test_zero_size_bypasses_and_removes_entry(self):
        enabled = self._file_io()
        disabled = self._file_io("0 b")
        original = reader_module.ds.dataset
        with patch.object(reader_module.ds, "dataset", wraps=original) as dataset:
            self._read(self.paths[0], enabled)
            self._read(self.paths[0], disabled)
            self._read(self.paths[0], enabled)
        self.assertEqual(3, dataset.call_count)

    def test_zero_size_clears_other_entries(self):
        enabled = self._file_io()
        disabled = self._file_io("0 b")
        original = reader_module.ds.dataset
        with patch.object(reader_module.ds, "dataset", wraps=original) as dataset:
            self._read(self.paths[0], enabled)
            self._read(self.paths[1], enabled)
            self._read(self.paths[0], disabled)
            self._read(self.paths[1], enabled)
        self.assertEqual(4, dataset.call_count)

    def test_reuses_dataset(self):
        file_io = self._file_io()
        original = reader_module.ds.dataset
        with patch.object(reader_module.ds, "dataset", wraps=original) as dataset:
            first = self._read(self.paths[0], file_io)
            second = self._read(self.paths[0], file_io)

        self.assertEqual(list(range(10)), first)
        self.assertEqual(first, second)
        self.assertEqual(1, dataset.call_count)

    def test_repeated_scan_skips_footer_io(self):
        path = os.path.join(self.temp_dir.name, "footer-io.parquet")
        pq.write_table(
            pa.table({
                "value": list(range(10000)),
                "payload": ["x" * 100] * 10000,
            }),
            path,
            row_group_size=100,
            compression="none",
        )
        counting = _CountingLocalFileSystem()
        file_io = self._file_io("0 b")
        file_io.filesystem = pafs.PyFileSystem(pafs.FSSpecHandler(counting))

        uncached = self._read(path, file_io)
        uncached_opens = counting.opens
        uncached_reads = len(counting.reads)

        counting.reset_counts()
        reader_module._reset_file_format_dataset_cache()
        file_io.properties.set(
            CatalogOptions.FILE_FORMAT_METADATA_CACHE_MAX_SIZE, "50 mb")
        self._read(path, file_io)
        counting.reset_counts()
        cached = self._read(path, file_io)

        self.assertEqual(uncached, cached)
        self.assertLess(counting.opens, uncached_opens)
        self.assertLess(len(counting.reads), uncached_reads)

    def test_evicts_least_recently_used_entry_by_estimated_size(self):
        cache = reader_module._FileFormatDatasetCache(10)
        first_key = (None, "parquet", "first")
        second_key = (None, "parquet", "second")
        third_key = (None, "parquet", "third")

        cache.get_or_load(first_key, lambda: "first", lambda _: 4)
        cache.get_or_load(second_key, lambda: "second", lambda _: 4)
        cache.get_or_load(first_key, lambda: "unused", lambda _: 4)
        cache.get_or_load(third_key, lambda: "third", lambda _: 4)

        self.assertEqual([first_key, third_key], list(cache._entries.keys()))
        self.assertEqual(8, cache.estimated_size)

    def test_evicts_least_recently_used_entry_by_entry_count(self):
        cache = reader_module._FileFormatDatasetCache(
            1024, max_entries=2)
        first_key = (None, "parquet", "first")
        second_key = (None, "parquet", "second")
        third_key = (None, "parquet", "third")

        cache.get_or_load(first_key, lambda: "first", lambda _: 1)
        cache.get_or_load(second_key, lambda: "second", lambda _: 1)
        cache.get_or_load(third_key, lambda: "third", lambda _: 1)

        self.assertEqual(
            [second_key, third_key], list(cache._entries.keys()))
        self.assertEqual(2, cache.estimated_size)

    def test_does_not_retain_entry_larger_than_size_limit(self):
        cache = reader_module._FileFormatDatasetCache(5)
        loads = []
        key = (None, "parquet", "large")

        def load():
            loads.append(True)
            return "large"

        self.assertEqual(
            "large", cache.get_or_load(key, load, lambda _: 6))
        self.assertEqual(
            "large", cache.get_or_load(key, load, lambda _: 6))
        self.assertEqual(2, len(loads))
        self.assertEqual(0, len(cache._entries))
        self.assertEqual(0, cache.estimated_size)

    def test_does_not_retain_entry_without_size_estimate(self):
        cache = reader_module._FileFormatDatasetCache(10)
        key = (None, "unknown", "data")
        loads = []

        def load():
            loads.append(True)
            return "unknown"

        self.assertEqual(
            "unknown", cache.get_or_load(key, load, lambda _: None))
        self.assertEqual(
            "unknown", cache.get_or_load(key, load, lambda _: None))
        self.assertEqual(2, len(loads))
        self.assertEqual(0, len(cache._entries))

    def test_coalesces_load_while_uncached_result_completes(self):
        cache = reader_module._FileFormatDatasetCache(10)
        key = (None, "unknown", "data")
        setting_result = threading.Event()
        waiter_started = threading.Event()
        release_result = threading.Event()
        loads = []
        future_count = []

        def new_future():
            future = Future()
            future_count.append(future)
            if len(future_count) == 1:
                original_set_result = future.set_result
                original_result = future.result

                def delayed_set_result(result):
                    setting_result.set()
                    release_result.wait()
                    original_set_result(result)

                def observed_result(*args, **kwargs):
                    waiter_started.set()
                    return original_result(*args, **kwargs)

                future.set_result = delayed_set_result
                future.result = observed_result
            return future

        def load():
            loads.append(True)
            return "unknown"

        with patch.object(reader_module, "Future", side_effect=new_future):
            with ThreadPoolExecutor(max_workers=2) as executor:
                first = executor.submit(
                    cache.get_or_load, key, load, lambda _: None)
                self.assertTrue(setting_result.wait(1))
                second = executor.submit(
                    cache.get_or_load, key, load, lambda _: None)
                waited = waiter_started.wait(1)
                release_result.set()

                self.assertTrue(waited)
                self.assertEqual("unknown", first.result())
                self.assertEqual("unknown", second.result())
        self.assertEqual(1, len(loads))

    def test_estimates_serialized_parquet_footer_size(self):
        dataset = reader_module.ds.dataset(self.paths[0], format="parquet")
        expected = sum(
            fragment.metadata.serialized_size
            for fragment in dataset.get_fragments()
        )
        self.assertGreater(expected, 0)
        self.assertEqual(
            expected,
            reader_module._estimate_file_format_dataset_size(
                dataset, "parquet"))

    def test_cache_entry_estimate_has_conservative_floor(self):
        dataset = reader_module.ds.dataset(
            self.paths[0], format="parquet")
        key = (
            reader_module._FilesystemIdentity(self.file_io.filesystem),
            "parquet",
            self.paths[0],
        )

        estimated = reader_module._estimate_file_format_cache_entry_size(
            key, dataset, "parquet")

        self.assertGreaterEqual(
            estimated,
            reader_module._FILE_FORMAT_METADATA_CACHE_MIN_ENTRY_SIZE)

    def test_process_cache_can_shrink_requested_capacity(self):
        cache = reader_module._file_format_dataset_cache(10)
        cache.get_or_load(("first", "parquet", "first"),
                          lambda: "first", lambda _: 4)
        cache.get_or_load(("second", "parquet", "second"),
                          lambda: "second", lambda _: 4)
        same_cache = reader_module._file_format_dataset_cache(5)

        self.assertIs(cache, same_cache)
        self.assertEqual(5, cache.max_size)
        self.assertEqual(
            [("second", "parquet", "second")],
            list(cache._entries.keys()))
        self.assertEqual(4, cache.estimated_size)

    def test_table_option_does_not_configure_process_cache(self):
        table_options = CoreOptions(Options({
            "file-format.metadata-cache.max-size": "0 b",
        }))
        original = reader_module.ds.dataset
        with patch.object(reader_module.ds, "dataset", wraps=original) as dataset:
            self._read(self.paths[0], options=table_options)
            self._read(self.paths[0], options=table_options)
        self.assertEqual(1, dataset.call_count)

    def test_shares_cache_across_file_io_with_same_filesystem(self):
        other_file_io = LocalFileIO(self.temp_dir.name, Options({}))
        other_file_io.filesystem = self.file_io.filesystem

        original = reader_module.ds.dataset
        with patch.object(reader_module.ds, "dataset", wraps=original) as dataset:
            reader_module._file_format_dataset(
                self.file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE)
            reader_module._file_format_dataset(
                other_file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE)

        self.assertEqual(1, dataset.call_count)

    def test_does_not_share_across_filesystems(self):
        other_file_io = LocalFileIO(self.temp_dir.name, Options({}))
        original = reader_module.ds.dataset
        with patch.object(reader_module.ds, "dataset", wraps=original) as dataset:
            reader_module._file_format_dataset(
                self.file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE)
            reader_module._file_format_dataset(
                other_file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE)
        self.assertEqual(2, dataset.call_count)

    def test_does_not_share_across_file_formats(self):
        parquet_dataset = object()
        orc_dataset = object()
        with patch.object(
                reader_module.ds, "dataset",
                side_effect=[parquet_dataset, orc_dataset]) as dataset:
            with patch.object(
                    reader_module, "_estimate_file_format_dataset_size",
                    return_value=1):
                first = reader_module._file_format_dataset(
                    self.file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE)
                second = reader_module._file_format_dataset(
                    self.file_io, "orc", self.paths[0], DEFAULT_CACHE_SIZE)

        self.assertIs(parquet_dataset, first)
        self.assertIs(orc_dataset, second)
        self.assertEqual(2, dataset.call_count)

    def test_cache_key_retains_filesystem_wrapper(self):
        root = pafs.LocalFileSystem()
        filesystem = pafs.SubTreeFileSystem(self.temp_dir.name, root)
        filesystem_ref = weakref.ref(filesystem)
        file_io = LocalFileIO(self.temp_dir.name, Options({}))
        file_io.filesystem = filesystem

        reader_module._file_format_dataset(
            file_io, "parquet", os.path.basename(self.paths[0]),
            DEFAULT_CACHE_SIZE)
        file_io.filesystem = root
        del filesystem
        gc.collect()

        self.assertIsNotNone(filesystem_ref())

    def test_filesystem_hash_collision_does_not_share_dataset(self):
        first_dir = tempfile.TemporaryDirectory()
        second_dir = tempfile.TemporaryDirectory()
        try:
            file_name = "same.parquet"
            pq.write_table(
                pa.table({"value": [1]}), os.path.join(first_dir.name, file_name))
            pq.write_table(
                pa.table({"value": [2]}), os.path.join(second_dir.name, file_name))
            file_io = LocalFileIO(first_dir.name, Options({}))
            first_filesystem = pafs.SubTreeFileSystem(
                first_dir.name, pafs.LocalFileSystem())
            second_filesystem = pafs.SubTreeFileSystem(
                second_dir.name, pafs.LocalFileSystem())

            with patch.object(
                    reader_module._FilesystemIdentity, "__hash__", return_value=1):
                file_io.filesystem = first_filesystem
                first = reader_module._file_format_dataset(
                    file_io, "parquet", file_name,
                    DEFAULT_CACHE_SIZE).to_table()
                file_io.filesystem = second_filesystem
                second = reader_module._file_format_dataset(
                    file_io, "parquet", file_name,
                    DEFAULT_CACHE_SIZE).to_table()

            self.assertEqual([1], first.column("value").to_pylist())
            self.assertEqual([2], second.column("value").to_pylist())
        finally:
            first_dir.cleanup()
            second_dir.cleanup()

    def test_resets_after_process_change(self):
        parent_cache = reader_module._file_format_dataset_cache(DEFAULT_CACHE_SIZE)
        with patch.object(reader_module.os, "getpid", return_value=os.getpid() + 1):
            child_cache = reader_module._file_format_dataset_cache(DEFAULT_CACHE_SIZE)
        self.assertIsNot(parent_cache, child_cache)

    def test_coalesces_concurrent_loads(self):
        original = reader_module.ds.dataset

        def delayed_dataset(*args, **kwargs):
            time.sleep(0.05)
            return original(*args, **kwargs)

        with patch.object(
                reader_module.ds, "dataset", side_effect=delayed_dataset) as dataset:
            with ThreadPoolExecutor(max_workers=8) as executor:
                results = list(executor.map(
                    lambda _: self._read(self.paths[0]),
                    range(8),
                ))

        self.assertEqual(1, dataset.call_count)
        self.assertTrue(all(value == list(range(10)) for value in results))


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