"""
Copyright (c) 2023 Aiven Ltd
See LICENSE for details
"""

import datetime
import time
from dataclasses import replace
from pathlib import Path
from unittest import mock

import pytest
import xxhash

from karapace.backup.backends.reader import ProducerSend, RestoreTopic
from karapace.backup.backends.v3.backend import SchemaBackupV3Reader, SchemaBackupV3Writer, _PartitionStats
from karapace.backup.backends.v3.errors import (
    InconsistentOffset,
    InvalidChecksum,
    OffsetMismatch,
    TooFewRecords,
    TooManyRecords,
    UnknownChecksumAlgorithm,
)
from karapace.backup.backends.v3.readers import read_records
from karapace.backup.backends.v3.schema import ChecksumAlgorithm, DataFile
from karapace.core.kafka.types import Timestamp
from tests.utils import StubMessage


def test_writer_reader_roundtrip(tmp_path: Path) -> None:
    backup_path = tmp_path
    topic_name = "a-topic"
    partition_index = 123
    started_at = datetime.datetime.now(datetime.timezone.utc)
    finished_at = datetime.datetime.now(datetime.timezone.utc)

    records = (
        StubMessage(
            key=b"foo",
            value=b"bar",
            topic=topic_name,
            partition=partition_index,
            offset=10,
            timestamp=(Timestamp.CREATE_TIME, round(time.time())),
            headers=None,
        ),
        StubMessage(
            key=b"foo",
            value=b"bar",
            topic=topic_name,
            partition=partition_index,
            offset=14,
            timestamp=(Timestamp.CREATE_TIME, round(time.time())),
            headers=[("some-key", b"some-value")],
        ),
    )
    topic_configurations = {"max.message.bytes": "1024"}

    # Write backup to files.
    backup_writer = SchemaBackupV3Writer(
        checksum_algorithm=ChecksumAlgorithm.xxhash3_64_be,
        max_records_per_checkpoint=1_000,
        max_bytes_per_checkpoint=16_384,
    )
    file_path = backup_writer.start_partition(
        path=backup_path,
        topic_name=topic_name,
        index=partition_index,
    )
    with backup_writer.safe_writer(file_path, False) as buffer:
        for record in records:
            backup_writer.store_record(buffer, record)
    data_file = backup_writer.finalize_partition(
        index=partition_index,
        filename=file_path.name,
    )
    backup_writer.store_metadata(
        path=backup_path,
        topic_name=topic_name,
        topic_id=None,
        started_at=started_at,
        finished_at=finished_at,
        replication_factor=2,
        topic_configurations=topic_configurations,
        data_files=(data_file,),
        partition_count=1,
    )

    assert sorted(path.name for path in backup_path.iterdir()) == ["a-topic.metadata", "a-topic:123.data"]
    assert data_file == DataFile(
        filename="a-topic:123.data",
        partition=123,
        checksum=mock.ANY,
        record_count=2,
        start_offset=10,
        end_offset=14,
    )

    # Read backup into restore instructions.
    backup_reader = SchemaBackupV3Reader()
    metadata_path = backup_path / "a-topic.metadata"
    instructions = tuple(
        backup_reader.read(
            path=metadata_path,
            topic_name="a-topic",
        )
    )

    assert instructions == (
        RestoreTopic(
            topic_name=topic_name,
            partition_count=1,
            replication_factor=2,
            topic_configs={"max.message.bytes": "1024"},
        ),
        ProducerSend(
            topic_name=topic_name,
            partition_index=partition_index,
            key=records[0].key(),
            value=records[0].value(),
            headers=(),
            timestamp=records[0].timestamp()[1],
        ),
        ProducerSend(
            topic_name=topic_name,
            partition_index=partition_index,
            key=records[1].key(),
            value=records[1].value(),
            headers=((b"some-key", b"some-value"),),
            timestamp=records[0].timestamp()[1],
        ),
    )


def make_record(topic_name: str, partition_index: int, offset: int) -> StubMessage:
    return StubMessage(
        key=b"foo",
        value=b"bar",
        topic=topic_name,
        partition=partition_index,
        offset=offset,
        timestamp=(Timestamp.CREATE_TIME, round(time.time())),
        headers=[("some-key", b"some-value")],
    )


def test_reader_raises_unknown_checksum_algorithm() -> None:
    topic_name = "a-topic"
    # This backup file was generated by temporarily adding an algorithm to the
    # ChecksumAlgorithm enum, and writing the backup using it.
    backup_path = Path(__file__).parent / "test_data" / "backup_v3_future_algorithm" / f"{topic_name}.metadata"

    # Read backup into restore instructions.
    backup_reader = SchemaBackupV3Reader()
    instructions = backup_reader.read(
        path=backup_path,
        topic_name=topic_name,
    )

    with pytest.raises(
        UnknownChecksumAlgorithm,
        match=r"^Tried restoring from a backup with an unknown checksum algorithm\.$",
    ):
        # Note: important that we use next() here, so that we assert that no
        # instructions are yielded before reaching the exception.
        next(instructions)


def test_reader_raises_invalid_checksum(tmp_path: Path) -> None:
    backup_path = tmp_path
    topic_name = "a-topic"
    partition_index = 123
    started_at = datetime.datetime.now(datetime.timezone.utc)
    finished_at = datetime.datetime.now(datetime.timezone.utc)

    backup_writer = SchemaBackupV3Writer(
        checksum_algorithm=ChecksumAlgorithm.xxhash3_64_be,
        # Write checkpoints on all records (except first).
        max_bytes_per_checkpoint=0,
        max_records_per_checkpoint=0,
    )
    file_path = backup_writer.start_partition(
        path=backup_path,
        topic_name=topic_name,
        index=partition_index,
    )

    # Write 2 records, then insert an invalid checksum on the third, and write
    # another record.
    with backup_writer.safe_writer(file_path, False) as buffer:
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 0))
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 1))
        with mock.patch.object(
            backup_writer._partition_stats[partition_index],
            "get_checkpoint",
            return_value=b"not what you expected!",
            autospec=True,
        ):
            backup_writer.store_record(buffer, make_record(topic_name, partition_index, 2))
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 3))

    # Finalize backup.
    data_file = backup_writer.finalize_partition(partition_index, file_path.name)
    backup_writer.store_metadata(
        path=backup_path,
        topic_name=topic_name,
        topic_id=None,
        started_at=started_at,
        finished_at=finished_at,
        replication_factor=1,
        topic_configurations={},
        data_files=(data_file,),
        partition_count=1,
    )

    # Read backup into restore instructions.
    backup_reader = SchemaBackupV3Reader()
    metadata_path = backup_path / "a-topic.metadata"

    with pytest.raises(
        InvalidChecksum,
        match=r"^Found invalid checksum at record number 2 \(counting from 0\), file byte offset 170\.$",
    ):
        tuple(
            backup_reader.read(
                path=metadata_path,
                topic_name="a-topic",
            )
        )


# This is tested separately, because if the last record in a file does not have
# checkpoint, we still want to make sure we discover if it's corrupted.
def test_reader_raises_invalid_checksum_for_corruption_in_last_record(tmp_path: Path) -> None:
    backup_path = tmp_path
    topic_name = "a-topic"
    partition_index = 123
    started_at = datetime.datetime.now(datetime.timezone.utc)
    finished_at = datetime.datetime.now(datetime.timezone.utc)

    backup_writer = SchemaBackupV3Writer(
        checksum_algorithm=ChecksumAlgorithm.xxhash3_64_be,
        # Make sure we have no checkpoints in the file.
        max_bytes_per_checkpoint=1_000_000,
        max_records_per_checkpoint=1_000_000,
    )
    file_path = backup_writer.start_partition(
        path=backup_path,
        topic_name=topic_name,
        index=partition_index,
    )

    with backup_writer.safe_writer(file_path, False) as buffer:
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 123))
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 124))

    # Flip last bit in byte before last.
    pos = -2
    file_bytes = file_path.read_bytes()
    file_path.write_bytes(file_bytes[:pos] + (file_bytes[pos] ^ 0b00000001).to_bytes(1, "big") + file_bytes[pos + 1 :])

    # Finalize backup.
    data_file = backup_writer.finalize_partition(partition_index, file_path.name)
    backup_writer.store_metadata(
        path=backup_path,
        topic_name=topic_name,
        topic_id=None,
        started_at=started_at,
        finished_at=finished_at,
        replication_factor=1,
        topic_configurations={},
        data_files=(data_file,),
        partition_count=1,
    )

    # Read backup into restore instructions.
    backup_reader = SchemaBackupV3Reader()
    metadata_path = backup_path / "a-topic.metadata"

    with pytest.raises(
        InvalidChecksum,
        match=r"^Found checksum mismatch after reading full data file.$",
    ):
        tuple(
            backup_reader.read(
                path=metadata_path,
                topic_name="a-topic",
            )
        )


def test_reader_raises_too_many_records(tmp_path: Path) -> None:
    backup_path = tmp_path
    topic_name = "a-topic"
    partition_index = 123
    started_at = datetime.datetime.now(datetime.timezone.utc)
    finished_at = datetime.datetime.now(datetime.timezone.utc)

    backup_writer = SchemaBackupV3Writer()
    file_path = backup_writer.start_partition(
        path=backup_path,
        topic_name=topic_name,
        index=partition_index,
    )

    with backup_writer.safe_writer(file_path, False) as buffer:
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 0))
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 1))

    # Finalize partition.
    data_file = backup_writer.finalize_partition(partition_index, file_path.name)

    # Alter the record count.
    data_file = replace(
        data_file,
        record_count=data_file.record_count - 1,
    )

    # Finalize backup
    backup_writer.store_metadata(
        path=backup_path,
        topic_name=topic_name,
        topic_id=None,
        started_at=started_at,
        finished_at=finished_at,
        replication_factor=1,
        topic_configurations={},
        data_files=(data_file,),
        partition_count=1,
    )

    # Read backup into restore instructions.
    backup_reader = SchemaBackupV3Reader()
    metadata_path = backup_path / "a-topic.metadata"

    with pytest.raises(
        TooManyRecords,
        match=r"^Data file contains data beyond last expected record\.$",
    ):
        tuple(
            backup_reader.read(
                path=metadata_path,
                topic_name="a-topic",
            )
        )


def test_reader_raises_too_few_records(tmp_path: Path) -> None:
    backup_path = tmp_path
    topic_name = "a-topic"
    partition_index = 123
    started_at = datetime.datetime.now(datetime.timezone.utc)
    finished_at = datetime.datetime.now(datetime.timezone.utc)

    backup_writer = SchemaBackupV3Writer()
    file_path = backup_writer.start_partition(
        path=backup_path,
        topic_name=topic_name,
        index=partition_index,
    )

    with backup_writer.safe_writer(file_path, False) as buffer:
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 0))

    # Finalize partition.
    data_file = backup_writer.finalize_partition(partition_index, file_path.name)

    # Alter the record count.
    data_file = replace(
        data_file,
        record_count=data_file.record_count + 1,
        end_offset=data_file.end_offset + 1,
    )

    # Finalize backup
    backup_writer.store_metadata(
        path=backup_path,
        topic_name=topic_name,
        topic_id=None,
        started_at=started_at,
        finished_at=finished_at,
        replication_factor=1,
        topic_configurations={},
        data_files=(data_file,),
        partition_count=1,
    )

    # Read backup into restore instructions.
    backup_reader = SchemaBackupV3Reader()
    metadata_path = backup_path / "a-topic.metadata"

    with pytest.raises(
        TooFewRecords,
        match=r"^Data file contains fewer records than expected\.$",
    ):
        tuple(
            backup_reader.read(
                path=metadata_path,
                topic_name="a-topic",
            )
        )


def test_reader_raises_offset_mismatch_for_first_record(tmp_path: Path) -> None:
    backup_path = tmp_path
    topic_name = "a-topic"
    partition_index = 123
    started_at = datetime.datetime.now(datetime.timezone.utc)
    finished_at = datetime.datetime.now(datetime.timezone.utc)

    backup_writer = SchemaBackupV3Writer()
    file_path = backup_writer.start_partition(
        path=backup_path,
        topic_name=topic_name,
        index=partition_index,
    )

    with backup_writer.safe_writer(file_path, False) as buffer:
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 123))
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 124))

    # Finalize partition.
    data_file = backup_writer.finalize_partition(partition_index, file_path.name)

    # Insert bad offset.
    data_file = replace(data_file, start_offset=122)

    # Finalize backup.
    backup_writer.store_metadata(
        path=backup_path,
        topic_name=topic_name,
        topic_id=None,
        started_at=started_at,
        finished_at=finished_at,
        data_files=(data_file,),
        partition_count=1,
        topic_configurations={},
        replication_factor=1,
    )

    # Read backup into restore instructions.
    backup_reader = SchemaBackupV3Reader()
    metadata_path = backup_path / "a-topic.metadata"

    with pytest.raises(
        OffsetMismatch,
        match=r"^First record in data file does not match expected start_offset \(expected: 122, actual: 123\)\.$",
    ):
        tuple(
            backup_reader.read(
                path=metadata_path,
                topic_name="a-topic",
            )
        )


def test_reader_raises_offset_mismatch_for_last_record(tmp_path: Path) -> None:
    backup_path = tmp_path
    topic_name = "a-topic"
    partition_index = 123
    started_at = datetime.datetime.now(datetime.timezone.utc)
    finished_at = datetime.datetime.now(datetime.timezone.utc)

    backup_writer = SchemaBackupV3Writer()
    file_path = backup_writer.start_partition(
        path=backup_path,
        topic_name=topic_name,
        index=partition_index,
    )

    with backup_writer.safe_writer(file_path, False) as buffer:
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 123))
        backup_writer.store_record(buffer, make_record(topic_name, partition_index, 124))

    # Finalize partition.
    data_file = backup_writer.finalize_partition(partition_index, file_path.name)

    # Insert bad offset.
    data_file = replace(data_file, end_offset=125)

    # Finalize backup.
    backup_writer.store_metadata(
        path=backup_path,
        topic_name=topic_name,
        topic_id=None,
        started_at=started_at,
        finished_at=finished_at,
        data_files=(data_file,),
        partition_count=1,
        topic_configurations={},
        replication_factor=1,
    )

    # Read backup into restore instructions.
    backup_reader = SchemaBackupV3Reader()
    metadata_path = backup_path / "a-topic.metadata"

    with pytest.raises(
        OffsetMismatch,
        match=r"^Last record in data file does not match expected end_offset \(expected: 125, actual: 124\)\.$",
    ):
        tuple(
            backup_reader.read(
                path=metadata_path,
                topic_name="a-topic",
            )
        )


def test_writer_respects_max_records_per_checkpoint(tmp_path: Path) -> None:
    backup_path = tmp_path
    topic_name = "a-topic"
    partition_index = 123
    records_per_checkpoint = 3

    records = (
        make_record(topic_name, partition_index, 0),
        make_record(topic_name, partition_index, 1),
        make_record(topic_name, partition_index, 2),
        make_record(topic_name, partition_index, 3),
        make_record(topic_name, partition_index, 4),
        make_record(topic_name, partition_index, 5),
        make_record(topic_name, partition_index, 6),
        make_record(topic_name, partition_index, 7),
        make_record(topic_name, partition_index, 8),
        make_record(topic_name, partition_index, 9),
        make_record(topic_name, partition_index, 10),
    )

    backup_writer = SchemaBackupV3Writer(
        checksum_algorithm=ChecksumAlgorithm.xxhash3_64_be,
        max_bytes_per_checkpoint=16_000_000,
        max_records_per_checkpoint=records_per_checkpoint,
    )
    file_path = backup_writer.start_partition(
        path=backup_path,
        topic_name=topic_name,
        index=partition_index,
    )
    with backup_writer.safe_writer(file_path, False) as buffer:
        for record in records:
            backup_writer.store_record(buffer, record)
    data_file = backup_writer.finalize_partition(partition_index, file_path.name)
    data_file_path = backup_path / data_file.filename

    checksum = xxhash.xxh64()
    with data_file_path.open("rb") as buffer:
        checkpoints = tuple(
            isinstance(record.checksum_checkpoint, bytes)
            for record in read_records(
                buffer=buffer,
                num_records=11,
                running_checksum=checksum,
            )
        )

    # Expect checkpoint for every third entry.
    assert checkpoints == (
        False,
        False,
        False,
        True,
        False,
        False,
        True,
        False,
        False,
        True,
        False,
    )

    # Check running checksum matches full file.
    full_file_checksum = xxhash.xxh64(data_file_path.read_bytes())
    assert checksum.digest() == full_file_checksum.digest()


def test_writer_respects_max_bytes_per_checkpoint(tmp_path: Path) -> None:
    backup_path = tmp_path
    topic_name = "a-topic"
    partition_index = 123
    # Each record + length is 43 bytes.
    bytes_per_checkpoint = 100

    records = (
        make_record(topic_name, partition_index, 0),
        make_record(topic_name, partition_index, 1),
        make_record(topic_name, partition_index, 2),
        make_record(topic_name, partition_index, 3),
        make_record(topic_name, partition_index, 4),
        make_record(topic_name, partition_index, 5),
        make_record(topic_name, partition_index, 6),
        make_record(topic_name, partition_index, 7),
        make_record(topic_name, partition_index, 8),
        make_record(topic_name, partition_index, 9),
        make_record(topic_name, partition_index, 10),
    )

    backup_writer = SchemaBackupV3Writer(
        checksum_algorithm=ChecksumAlgorithm.xxhash3_64_be,
        max_bytes_per_checkpoint=bytes_per_checkpoint,
        max_records_per_checkpoint=1_000,
    )
    file_path = backup_writer.start_partition(
        path=backup_path,
        topic_name=topic_name,
        index=partition_index,
    )
    with backup_writer.safe_writer(file_path, False) as buffer:
        for record in records:
            backup_writer.store_record(buffer, record)
    data_file = backup_writer.finalize_partition(partition_index, file_path.name)
    data_file_path = backup_path / data_file.filename

    checksum = xxhash.xxh64()
    with data_file_path.open("rb") as buffer:
        checkpoints = tuple(
            isinstance(record.checksum_checkpoint, bytes)
            for record in read_records(
                buffer=buffer,
                num_records=11,
                running_checksum=checksum,
            )
        )

    # Expect checkpoint for about every third entry.
    assert checkpoints == (
        False,
        False,
        False,
        True,
        False,
        False,
        True,
        False,
        False,
        True,
        False,
    )

    # Check running checksum matches full file.
    full_file_checksum = xxhash.xxh64(data_file_path.read_bytes())
    assert checksum.digest() == full_file_checksum.digest()


class TestPartitionStats:
    def test_get_checkpoint_returns_bytes_when_byte_threshold_exceeded(self) -> None:
        stats = _PartitionStats(
            running_checksum=xxhash.xxh64(),
            records_written=100,
            bytes_written_since_checkpoint=1023,
            records_written_since_checkpoint=23,
        )
        checkpoint = stats.get_checkpoint(
            bytes_threshold=1024,
            records_threshold=24,
        )
        assert checkpoint is None
        assert stats.records_written == 100
        assert stats.bytes_written_since_checkpoint == 1023
        assert stats.records_written_since_checkpoint == 23

        checkpoint = stats.get_checkpoint(
            bytes_threshold=1023,
            records_threshold=24,
        )
        assert isinstance(checkpoint, bytes)
        assert stats.records_written == 100
        assert stats.bytes_written_since_checkpoint == 0
        assert stats.records_written_since_checkpoint == 0

    def test_get_checkpoint_returns_bytes_when_record_threshold_exceeded(self) -> None:
        stats = _PartitionStats(
            running_checksum=xxhash.xxh64(),
            records_written=100,
            bytes_written_since_checkpoint=1023,
            records_written_since_checkpoint=23,
        )
        checkpoint = stats.get_checkpoint(
            bytes_threshold=1024,
            records_threshold=24,
        )
        assert checkpoint is None
        assert stats.records_written == 100
        assert stats.bytes_written_since_checkpoint == 1023
        assert stats.records_written_since_checkpoint == 23

        checkpoint = stats.get_checkpoint(
            bytes_threshold=1024,
            records_threshold=23,
        )
        assert isinstance(checkpoint, bytes)
        assert stats.records_written == 100
        assert stats.bytes_written_since_checkpoint == 0
        assert stats.records_written_since_checkpoint == 0

    def test_can_increase_stats(self) -> None:
        stats = _PartitionStats(running_checksum=xxhash.xxh64())
        assert stats == _PartitionStats(
            running_checksum=stats.running_checksum,
            records_written=0,
            bytes_written_since_checkpoint=0,
            records_written_since_checkpoint=0,
            min_offset=None,
            max_offset=None,
        )
        stats.update(
            bytes_offset=1,
            record_offset=1,
        )
        assert stats == _PartitionStats(
            running_checksum=stats.running_checksum,
            records_written=1,
            bytes_written_since_checkpoint=1,
            records_written_since_checkpoint=1,
            min_offset=1,
            max_offset=1,
        )
        stats.update(
            bytes_offset=1023,
            record_offset=23,
        )
        assert stats == _PartitionStats(
            running_checksum=stats.running_checksum,
            records_written=2,
            bytes_written_since_checkpoint=1024,
            records_written_since_checkpoint=2,
            min_offset=1,
            max_offset=23,
        )

    def test_update_raises_inconsistent_offset_when_not_increasing(self) -> None:
        stats = _PartitionStats(
            running_checksum=xxhash.xxh64(),
            min_offset=0,
            max_offset=0,
        )

        with pytest.raises(
            InconsistentOffset,
            match=(
                r"^Read record offset that's less than or equal to an already seen "
                r"record. Expected record offset 0 > 0\.$"
            ),
        ):
            stats.update(bytes_offset=1, record_offset=0)

    def test_update_raises_inconsistent_offset_when_decreasing(self) -> None:
        stats = _PartitionStats(
            running_checksum=xxhash.xxh64(),
            min_offset=10,
            max_offset=20,
        )

        with pytest.raises(
            InconsistentOffset,
            match=(
                r"^Read record offset that's less than or equal to an already seen "
                r"record. Expected record offset 19 > 20\.$"
            ),
        ):
            stats.update(bytes_offset=1, record_offset=19)
