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

from types import SimpleNamespace
import os
import threading

import numpy as np
import pandas as pd
import pyarrow as pa
import pytest

import tsfile.dataset.index as index_module
import tsfile.dataset.runtime as runtime_module
from tsfile import (
    ColumnCategory,
    ColumnSchema,
    TableSchema,
    TsFileDataFrame,
    TsFileTableWriter,
)
from tsfile.constants import TSDataType
from tsfile.dataset.index import (
    COLUMN_SCHEMA,
    DEVICE_FILE_SPAN,
    DIRECTORY,
    HEADER,
    LOGICAL_SERIES,
    MappedDatasetIndex,
    RECORDS,
    SERIES_FILE_SPAN,
    SERIES_LOCATOR,
    TSFILE_RECORD,
    build_sections_from_dataframe,
    crc32c,
    write_index_atomic,
)
from tsfile.dataset.metadata import MetadataCatalog, SeriesStats
from tsfile.dataset.runtime import RuntimeSeriesReader
from tsfile.dataset.runtime import DatasetRuntime
from tsfile.tsfile_reader import TsFileReaderPy


def _synthetic_dataframe(source_path):
    catalog = MetadataCatalog()
    table_id = catalog.add_table("root", (), (), ("s1",))
    device_id = catalog.add_device(table_id, (), 1, 10)
    catalog.series_stats_by_ref[(device_id, 0)] = SeriesStats(
        10,
        1,
        10,
        10,
        1,
        10,
        int(TSDataType.INT64),
        128,
        16,
        64,
        16,
        1,
        1,
        1,
        1,
    )
    reader = SimpleNamespace(file_path=source_path, catalog=catalog)
    index = SimpleNamespace(
        table_entries={"root": catalog.table_entries[0]},
        devices=[("root", ())],
        series=[(0, 0)],
        series_shards={(0, 0): [(reader, 0, 0)]},
    )
    return SimpleNamespace(
        _index=index,
        _readers={source_path: reader},
        _paths=[source_path],
    )


def test_binary_layout_matches_cpp_v1():
    assert HEADER.size == 64
    assert DIRECTORY.size == 32
    assert RECORDS[COLUMN_SCHEMA].size == 32
    assert RECORDS[LOGICAL_SERIES].size == 32
    assert RECORDS[TSFILE_RECORD].size == 32
    assert RECORDS[DEVICE_FILE_SPAN].size == 32
    assert RECORDS[SERIES_FILE_SPAN].size == 40
    assert RECORDS[SERIES_LOCATOR].size == 24
    assert index_module.SECTION_COUNT == 13


def test_crc32c_known_vector():
    assert crc32c(b"123456789") == 0xE3069283


def test_build_publish_map_and_lookup(tmp_path):
    source = tmp_path / "source.tsfile"
    source.write_bytes(b"T" * 4096)
    output = tmp_path / "dataset.tsidx"
    dataframe = _synthetic_dataframe(str(source))
    write_index_atomic(str(output), build_sections_from_dataframe(dataframe))

    with MappedDatasetIndex(str(output), verify_sections=True) as index:
        table_id = index.find_table_ids("root")[0]
        device_id = index.find_device_id(table_id, "root.")
        column_id = index.find_column_id(table_id, "s1")
        assert index.find_series_id(device_id, column_id) == 0
        assert index.count(LOGICAL_SERIES) == 1
        assert index.count(SERIES_LOCATOR) == 1
        file_record = index.record(TSFILE_RECORD, 0)
        assert index.string(file_record[0]) == str(source)


def test_child_lookup_checks_full_bytes_for_hash_collisions(monkeypatch):
    rows = [
        (0, 10, 42, 0, 0),
        (0, 11, 42, 1, 0),
        (0, 12, 42, 2, 0),
    ]
    names = [b"alpha", b"beta", b"gamma"]
    monkeypatch.setattr(index_module, "name_hash", lambda _value: 42)

    class _Index:
        @staticmethod
        def record(_section_type, record_id):
            return rows[record_id]

        @staticmethod
        def string_bytes(sid):
            return names[sid]

    assert MappedDatasetIndex._find_child(_Index(), 0, 0, "beta", 0, 3) == 11
    with pytest.raises(KeyError):
        MappedDatasetIndex._find_child(_Index(), 0, 0, "missing", 0, 3)


def test_rejects_damaged_header_checksum(tmp_path):
    source = tmp_path / "source.tsfile"
    source.write_bytes(b"T" * 4096)
    output = tmp_path / "dataset.tsidx"
    write_index_atomic(
        str(output), build_sections_from_dataframe(_synthetic_dataframe(str(source)))
    )
    with output.open("r+b") as stream:
        stream.seek(32)
        stream.write((1).to_bytes(8, "little"))
    with pytest.raises(ValueError, match="header shape|checksum"):
        MappedDatasetIndex(str(output))


def test_reader_exposes_exact_aligned_locator_ranges(tmp_path):
    source = tmp_path / "aligned.tsfile"
    schema = TableSchema(
        "weather",
        [
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
        ],
    )
    with TsFileTableWriter(str(source), schema) as writer:
        writer.write_dataframe(
            pd.DataFrame(
                {
                    "time": [0, 1, 2],
                    "device": ["d0", "d0", "d0"],
                    "value": [1.0, 2.0, 3.0],
                }
            )
        )
    reader = TsFileReaderPy(str(source))
    try:
        groups = reader.get_timeseries_metadata()
        metadata = [item for group in groups.values() for item in group.timeseries]
        assert metadata
        assert all(item.value_metadata_length > 0 for item in metadata)
        assert all(
            item.value_metadata_offset + item.value_metadata_length
            <= os.path.getsize(source)
            for item in metadata
        )
        aligned = [item for item in metadata if item.layout == 1]
        assert aligned
        assert all(item.time_metadata_length > 0 for item in aligned)
        assert all(
            item.time_chunk_meta_count == item.chunk_meta_count for item in aligned
        )
    finally:
        reader.close()


def _write_runtime_file(path, start):
    schema = TableSchema(
        "weather",
        [
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
        ],
    )
    with TsFileTableWriter(str(path), schema) as writer:
        writer.write_dataframe(
            pd.DataFrame(
                {
                    "time": [start, start + 1],
                    "device": ["d0", "d0"],
                    "value": [float(start), float(start + 1)],
                }
            )
        )


def _write_runtime_devices_file(path):
    schema = TableSchema(
        "weather",
        [
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
        ],
    )
    with TsFileTableWriter(str(path), schema) as writer:
        writer.write_dataframe(
            pd.DataFrame(
                {
                    "time": [0, 1, 0, 1, 0, 1],
                    "device": ["d0", "d0", "d1", "d1", "d2", "d2"],
                    "value": [0.0, 1.0, 10.0, 11.0, 20.0, 21.0],
                }
            )
        )


def test_hot_construction_maps_index_without_opening_readers(tmp_path, monkeypatch):
    source = tmp_path / "part.tsfile"
    _write_runtime_file(source, 0)
    with TsFileDataFrame(str(source), show_progress=False, use_index=True) as first:
        assert len(first) == 1

    def fail_legacy_scan(*_args, **_kwargs):
        raise AssertionError("hot construction must not build a legacy catalog")

    monkeypatch.setattr("tsfile.dataset.reader.TsFileSeriesReader", fail_legacy_scan)
    with TsFileDataFrame(str(source), show_progress=False, use_index=True) as second:
        assert len(second) == 1
        assert second._runtime.readers.open_count == 0
        series = second[0]
        assert second._runtime.readers.open_count == 0
        assert series[0] == 0.0
        assert second._runtime.readers.open_count == 1
        assert second._runtime.prepared.size == 1
        assert series[1] == 1.0
        assert second._runtime.prepared.size == 1
        series.close()


def test_dataframe_does_not_use_or_create_index_by_default(tmp_path):
    source = tmp_path / "part.tsfile"
    _write_runtime_file(source, 0)
    index_path = index_module.index_path_for([str(source)])

    with TsFileDataFrame(str(source), show_progress=False) as dataframe:
        assert dataframe._runtime is None
        assert len(dataframe) == 1
        np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0]))
        aligned = dataframe.loc[0:1, [0]]
        np.testing.assert_array_equal(aligned.timestamps, np.array([0, 1]))
        np.testing.assert_array_equal(aligned.values, np.array([[0.0], [1.0]]))

    assert not os.path.exists(index_path)


def test_persistent_index_path_is_scoped_to_expanded_file_set(tmp_path, monkeypatch):
    first = tmp_path / "part1.tsfile"
    second = tmp_path / "part2.tsfile"
    third = tmp_path / "part3.tsfile"
    _write_runtime_file(first, 0)
    _write_runtime_file(second, 2)
    _write_runtime_file(third, 10)

    first_set = [str(first), str(second)]
    second_set = [str(first), str(third)]
    first_index = index_module.index_path_for(first_set)
    second_index = index_module.index_path_for(second_set)
    assert first_index != second_index

    with TsFileDataFrame(first_set, show_progress=False, use_index=True) as dataframe:
        assert len(dataframe) == 1
    with TsFileDataFrame(second_set, show_progress=False, use_index=True) as dataframe:
        assert len(dataframe) == 1
    assert os.path.exists(first_index)
    assert os.path.exists(second_index)

    def fail_legacy_scan(*_args, **_kwargs):
        raise AssertionError("matching file-set index should be reused")

    monkeypatch.setattr("tsfile.dataset.reader.TsFileSeriesReader", fail_legacy_scan)
    with TsFileDataFrame(first_set, show_progress=False, use_index=True) as dataframe:
        np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0, 2.0, 3.0]))
    with TsFileDataFrame(second_set, show_progress=False, use_index=True) as dataframe:
        np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0, 10.0, 11.0]))


def test_named_selection_reuses_bounded_runtime_descriptor(tmp_path, monkeypatch):
    source = tmp_path / "part.tsfile"
    _write_runtime_file(source, 0)

    with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe:
        name = str(dataframe.list_timeseries()[0])
        find_device_calls = 0
        original_find_device = dataframe._runtime.index.find_device_id

        def count_find_device(*args, **kwargs):
            nonlocal find_device_calls
            find_device_calls += 1
            return original_find_device(*args, **kwargs)

        def fail_series_info(*_args, **_kwargs):
            raise AssertionError(
                "descriptor-backed selection must not rebuild series info"
            )

        monkeypatch.setattr(
            dataframe._runtime.index, "find_device_id", count_find_device
        )
        monkeypatch.setattr(
            RuntimeSeriesReader, "get_series_info_by_ref", fail_series_info
        )

        first = dataframe[name]
        second = dataframe[name]
        assert first.stats == {"start_time": 0, "end_time": 1, "count": 2}
        np.testing.assert_array_equal(first[:], np.array([0.0, 1.0]))
        assert second.stats == first.stats
        assert find_device_calls == 1
        assert len(dataframe._index._descriptor_cache) == 1
        first.close()
        second.close()


def test_listed_series_path_resolves_directly_by_snapshot_series_id(
    tmp_path, monkeypatch
):
    source = tmp_path / "part.tsfile"
    _write_runtime_file(source, 0)

    with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe:
        path = dataframe.list_timeseries()[0]
        assert isinstance(path, str)
        assert path.series_id == 0
        assert path._index_identity == dataframe._runtime.index.identity

        def fail_name_lookup(*_args, **_kwargs):
            raise AssertionError("listed SeriesPath must bypass device name lookup")

        def fail_name_rebuild(*_args, **_kwargs):
            raise AssertionError("listed SeriesPath must not be rebuilt from the index")

        def fail_span_lookup(*_args, **_kwargs):
            raise AssertionError("descriptor locator must bypass series span lookup")

        monkeypatch.setattr(
            dataframe._runtime.index, "find_device_id", fail_name_lookup
        )
        monkeypatch.setattr(dataframe, "_build_series_name", fail_name_rebuild)
        monkeypatch.setattr(RuntimeSeriesReader, "_span", fail_span_lookup)
        series = dataframe[path]
        assert series.name is path
        np.testing.assert_array_equal(series[:], np.array([0.0, 1.0]))
        series.close()

        # Converting to a plain str deliberately drops the snapshot-local id.
        assert not hasattr(str(path), "series_id")


def test_series_path_from_another_index_falls_back_to_its_name(tmp_path, monkeypatch):
    first_dir = tmp_path / "first"
    second_dir = tmp_path / "second"
    first_dir.mkdir()
    second_dir.mkdir()
    first_source = first_dir / "part.tsfile"
    second_source = second_dir / "part.tsfile"
    _write_runtime_file(first_source, 0)
    _write_runtime_file(second_source, 10)

    with TsFileDataFrame(
        str(first_source), show_progress=False, use_index=True
    ) as first:
        foreign_path = first.list_timeseries()[0]
        with TsFileDataFrame(
            str(second_source), show_progress=False, use_index=True
        ) as second:
            assert foreign_path._index_identity != second._runtime.index.identity
            find_device_calls = 0
            original_find_device = second._runtime.index.find_device_id

            def count_find_device(*args, **kwargs):
                nonlocal find_device_calls
                find_device_calls += 1
                return original_find_device(*args, **kwargs)

            monkeypatch.setattr(
                second._runtime.index, "find_device_id", count_find_device
            )
            series = second[foreign_path]
            np.testing.assert_array_equal(series[:], np.array([10.0, 11.0]))
            assert find_device_calls == 1
            series.close()


def test_runtime_descriptor_cache_evicts_least_recent_name(tmp_path, monkeypatch):
    source = tmp_path / "devices.tsfile"
    _write_runtime_devices_file(source)
    monkeypatch.setattr(runtime_module, "_SERIES_DESCRIPTOR_CACHE_SIZE", 2)

    with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe:
        names = [str(name) for name in dataframe.list_timeseries()]
        find_device_calls = 0
        original_find_device = dataframe._runtime.index.find_device_id

        def count_find_device(*args, **kwargs):
            nonlocal find_device_calls
            find_device_calls += 1
            return original_find_device(*args, **kwargs)

        monkeypatch.setattr(
            dataframe._runtime.index, "find_device_id", count_find_device
        )
        for name in names:
            dataframe[name].close()

        assert find_device_calls == 3
        assert len(dataframe._index._descriptor_cache) == 2
        assert len(dataframe._index.series_shards._cache) == 2

        # d0 was the least recently used name and must be resolved again.
        dataframe[names[0]].close()
        assert find_device_calls == 4


def test_reader_pool_enforces_open_file_cap(tmp_path, monkeypatch):
    first = tmp_path / "part1.tsfile"
    second = tmp_path / "part2.tsfile"
    _write_runtime_file(first, 0)
    _write_runtime_file(second, 2)
    monkeypatch.setenv("TSFILE_DATAFRAME_MAX_OPEN_FILES", "1")
    with TsFileDataFrame(
        [str(first), str(second)], show_progress=False, use_index=True
    ) as dataframe:
        series = dataframe[0]
        assert list(series[:]) == [0.0, 1.0, 2.0, 3.0]
        assert dataframe._runtime.readers.open_count == 1
        assert dataframe._runtime.prepared.size == 2
        series.close()


def test_runtime_consume_concatenates_arrow_batches_without_scalar_reads():
    class _ArrowResult:
        def __init__(self):
            self._batches = iter(
                [
                    pa.table(
                        {
                            "time": pa.array([], type=pa.int64()),
                            "value": pa.array([], type=pa.float64()),
                        }
                    ),
                    pa.table(
                        {
                            "time": pa.array([1, 2], type=pa.int64()),
                            "value": pa.array([10.0, None], type=pa.float64()),
                        }
                    ),
                    pa.table(
                        {
                            "time": pa.array([3], type=pa.int64()),
                            "value": pa.array([30.0], type=pa.float64()),
                        }
                    ),
                ]
            )
            self.closed = False

        def __enter__(self):
            return self

        def __exit__(self, *_):
            self.closed = True

        def read_arrow_batch(self):
            return next(self._batches, None)

        def next(self):
            raise AssertionError("Runtime must not consume prepared rows one by one")

    result = _ArrowResult()
    timestamps, values = RuntimeSeriesReader._consume(result)

    np.testing.assert_array_equal(timestamps, np.array([1, 2, 3], dtype=np.int64))
    np.testing.assert_allclose(values, np.array([10.0, np.nan, 30.0]), equal_nan=True)
    assert result.closed


def test_prepared_query_reads_nullable_offset_window_in_arrow_batches(tmp_path):
    source = tmp_path / "nullable.tsfile"
    schema = TableSchema(
        "weather",
        [
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
        ],
    )
    expected = np.arange(10, dtype=np.float64)
    expected[2] = np.nan
    expected[6] = np.nan
    with TsFileTableWriter(str(source), schema) as writer:
        writer.write_dataframe(
            pd.DataFrame(
                {
                    "time": np.arange(10, dtype=np.int64),
                    "device": ["d0"] * 10,
                    "value": expected,
                }
            )
        )

    with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe:
        runtime = dataframe._runtime
        series = runtime.index.record(LOGICAL_SERIES, 0)
        span = runtime.index.record(SERIES_FILE_SPAN, series[2])
        with runtime.readers.acquire(0) as reader:
            prepared = runtime.prepared.get(0, span[2], reader)
            with reader.query_prepared(prepared, offset=1, limit=7) as result:
                batches = []
                while True:
                    batch = result.read_arrow_batch()
                    if batch is None:
                        break
                    batches.append(batch)
            with reader.query_prepared(
                prepared, start_time=100, end_time=200
            ) as empty_result:
                assert empty_result.read_arrow_batch() is None

    assert batches
    table = pa.concat_tables(batches)
    np.testing.assert_array_equal(
        table.column("time").to_numpy(), np.arange(1, 8, dtype=np.int64)
    )
    np.testing.assert_allclose(
        table.column("value").to_numpy(zero_copy_only=False),
        expected[1:8],
        equal_nan=True,
    )


def test_prepared_locator_rejects_stale_generation_and_bad_range(tmp_path):
    source = tmp_path / "part.tsfile"
    _write_runtime_file(source, 0)
    with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe:
        runtime = dataframe._runtime
        series = runtime.index.record(LOGICAL_SERIES, 0)
        span = runtime.index.record(SERIES_FILE_SPAN, series[2])
        locator = list(runtime.prepared._locator_tuple(0, span[2]))
        with runtime.readers.acquire(0) as reader:
            stale = list(locator)
            stale[3] ^= 1
            with pytest.raises(Exception, match="prepare Dataset Index locator"):
                reader.prepare_series(stale)

            out_of_range = list(locator)
            out_of_range[7] = os.path.getsize(source) + 1
            with pytest.raises(Exception, match="prepare Dataset Index locator"):
                reader.prepare_series(out_of_range)


def test_reader_session_revalidates_generation_when_reused(tmp_path):
    source = tmp_path / "part.tsfile"
    _write_runtime_file(source, 0)
    with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe:
        pool = dataframe._runtime.readers
        with pool.acquire(0):
            pass
        stat = os.stat(source)
        os.utime(source, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1_000_000))
        with pytest.raises(RuntimeError, match="generation changed"):
            with pool.acquire(0):
                pass


def test_runtime_lease_close_waits_for_query_lease(tmp_path):
    source = tmp_path / "part.tsfile"
    _write_runtime_file(source, 0)
    with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe:
        runtime = DatasetRuntime(str(dataframe._runtime.index.path), query_workers=1)
        lease = runtime.lease()
        entered = threading.Event()
        release = threading.Event()
        query_done = threading.Event()

        def run_query():
            with lease.query_lease():
                entered.set()
                assert release.wait(timeout=2)
            query_done.set()

        query_thread = threading.Thread(target=run_query)
        query_thread.start()
        assert entered.wait(timeout=2)

        close_done = threading.Event()

        def close_lease():
            lease.close()
            close_done.set()

        close_thread = threading.Thread(target=close_lease)
        close_thread.start()
        assert not close_done.wait(timeout=0.05)

        release.set()
        query_thread.join(timeout=2)
        close_thread.join(timeout=2)
        assert query_done.is_set()
        assert close_done.is_set()