# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

"""Unit tests for UDTF checkpoint key helpers and store operations."""

import pyarrow as pa

from geneva.checkpoint import InMemoryCheckpointStore
from geneva.checkpoint_utils import (
    UDTF_PARTITION_DONE_MARKER,
    format_udtf_batch_key,
    format_udtf_checkpoint_prefix,
    format_udtf_partition_prefix,
)

# ---------------------------------------------------------------------------
# Key helper tests
# ---------------------------------------------------------------------------


class TestUDTFCheckpointKeyHelpers:
    def test_format_checkpoint_prefix(self) -> None:
        prefix = format_udtf_checkpoint_prefix(
            udtf_name="dedup",
            udtf_version="abc123",
            source_version=42,
        )
        assert prefix == "udtf_dedup_abc123_src-42"

    def test_format_partition_prefix_unpartitioned(self) -> None:
        top = "udtf_dedup_abc123_src-42"
        result = format_udtf_partition_prefix(
            top, partition_col=None, partition_value=None
        )
        assert result == "udtf_dedup_abc123_src-42___all__"

    def test_format_partition_prefix_partitioned(self) -> None:
        top = "udtf_dedup_abc123_src-42"
        result = format_udtf_partition_prefix(
            top, partition_col="group", partition_value="a"
        )
        assert result == "udtf_dedup_abc123_src-42_group=a"

    def test_format_partition_prefix_sanitizes_underscores(self) -> None:
        top = "udtf_dedup_abc123_src-42"
        result = format_udtf_partition_prefix(
            top, partition_col="category", partition_value="foo_bar"
        )
        assert result == "udtf_dedup_abc123_src-42_category=foo-bar"

    def test_format_batch_key_zero_padded(self) -> None:
        part = "udtf_dedup_abc123_src-42_group=a"
        key = format_udtf_batch_key(part, 5)
        assert key == "udtf_dedup_abc123_src-42_group=a_batch-0005"

    def test_format_batch_key_zero(self) -> None:
        part = "udtf_dedup_abc123_src-42___all__"
        key = format_udtf_batch_key(part, 0)
        assert key == "udtf_dedup_abc123_src-42___all___batch-0000"

    def test_sorted_batch_keys_numeric_order(self) -> None:
        part = "udtf_x_v1_src-1_group=a"
        keys = [format_udtf_batch_key(part, i) for i in [9, 0, 100, 3]]
        assert sorted(keys) == [
            format_udtf_batch_key(part, 0),
            format_udtf_batch_key(part, 3),
            format_udtf_batch_key(part, 9),
            format_udtf_batch_key(part, 100),
        ]


# ---------------------------------------------------------------------------
# Checkpoint store write/read tests
# ---------------------------------------------------------------------------


class TestUDTFCheckpointWriteRead:
    def test_write_read_roundtrip(self) -> None:
        store = InMemoryCheckpointStore()
        batch = pa.RecordBatch.from_pydict({"id": [1, 2], "val": [10.0, 20.0]})
        key = format_udtf_batch_key("udtf_x_v1_src-1___all__", 0)

        store[key] = batch
        assert key in store
        result = store[key]
        assert result.equals(batch)

    def test_list_partition_batches(self) -> None:
        store = InMemoryCheckpointStore()
        prefix_a = "udtf_x_v1_src-1_group=a"
        prefix_b = "udtf_x_v1_src-1_group=b"

        for i in range(3):
            batch = pa.RecordBatch.from_pydict({"id": [i]})
            store[format_udtf_batch_key(prefix_a, i)] = batch

        store[format_udtf_batch_key(prefix_b, 0)] = pa.RecordBatch.from_pydict(
            {"id": [99]}
        )

        a_keys = sorted(store.list_keys(prefix_a))
        assert len(a_keys) == 3
        # Ensure no prefix_b keys snuck in
        for k in a_keys:
            assert k.startswith(prefix_a)

    def test_done_marker_signals_completion(self) -> None:
        store = InMemoryCheckpointStore()
        prefix = "udtf_x_v1_src-1___all__"
        done_key = f"{prefix}_{UDTF_PARTITION_DONE_MARKER}"

        assert done_key not in store

        store[done_key] = pa.RecordBatch.from_pydict({"_done": [True]})
        assert done_key in store
