import json
import time
from datetime import datetime, timedelta, timezone

import pytest

import ray
from ray.data._internal.datasource.kafka_datasink import KafkaDatasink
from ray.data._internal.datasource.kafka_datasource import (
    KafkaAuthConfig,
    _build_consumer_config_for_read,
    _datetime_to_ms,
)

pytest.importorskip("confluent_kafka")


def _wait_for_watermark(bootstrap_server, topic, expected_count, timeout=5):
    """Poll until the topic's high watermark reaches expected_count messages."""
    from confluent_kafka import Consumer, TopicPartition

    consumer = Consumer(
        {
            "bootstrap.servers": bootstrap_server,
            "group.id": "test-watermark-poller",
        }
    )
    try:
        deadline = time.monotonic() + timeout
        total = 0
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break

            metadata = consumer.list_topics(topic, timeout=remaining)
            topic_meta = metadata.topics.get(topic)
            if topic_meta and topic_meta.partitions:
                total = 0
                for pid in topic_meta.partitions:
                    remaining = deadline - time.monotonic()
                    if remaining <= 0:
                        break
                    _, high = consumer.get_watermark_offsets(
                        TopicPartition(topic, pid), timeout=remaining
                    )
                    total += high
                if total >= expected_count:
                    return
            time.sleep(0.1)
        raise TimeoutError(
            f"Timed out waiting for {expected_count} messages in topic {topic!r} "
            f"(got {total})"
        )
    finally:
        consumer.close()


@pytest.fixture(scope="session")
def kafka_container():
    from testcontainers.kafka import KafkaContainer

    print("\nStarting Kafka container (shared across all tests)...")
    with KafkaContainer() as kafka:
        bootstrap_server = kafka.get_bootstrap_server()
        print(f"Kafka container started at {bootstrap_server}")
        yield kafka
        print("\nShutting down Kafka container...")


@pytest.fixture(scope="session")
def bootstrap_server(kafka_container):
    return kafka_container.get_bootstrap_server()


def _json_value_serializer(obj, ctx):
    return json.dumps(obj).encode("utf-8")


def _str_key_serializer(obj, ctx):
    return obj.encode("utf-8") if obj else None


@pytest.fixture(scope="session")
def kafka_producer(bootstrap_server):
    from confluent_kafka.serializing_producer import SerializingProducer

    print(f"Creating shared Kafka producer for {bootstrap_server}")
    producer = SerializingProducer(
        {
            "bootstrap.servers": bootstrap_server,
            "value.serializer": _json_value_serializer,
            "key.serializer": _str_key_serializer,
        }
    )
    yield producer
    producer.flush()
    print("Closing shared Kafka producer")
    producer.close()


@pytest.fixture(scope="session")
def kafka_consumer(bootstrap_server):
    from confluent_kafka import Consumer

    print(f"Creating shared Kafka consumer for {bootstrap_server}")
    consumer = Consumer(
        {
            "bootstrap.servers": bootstrap_server,
            "group.id": "ray-test-consumer",
            "auto.offset.reset": "earliest",
            "enable.auto.commit": False,
        }
    )
    yield consumer
    print("Closing shared Kafka consumer")
    consumer.close()


def consume_messages(consumer, topic, expected_count, timeout=10):
    """Helper function to consume messages from a topic."""
    from confluent_kafka import KafkaError, TopicPartition

    # Discover partitions for the topic
    metadata = consumer.list_topics(timeout=10)
    topic_meta = metadata.topics.get(topic)
    if topic_meta is None or not topic_meta.partitions:
        time.sleep(2)  # Wait a bit more for topic to be created
        metadata = consumer.list_topics(timeout=10)
        topic_meta = metadata.topics.get(topic)

    if topic_meta and topic_meta.partitions:
        topic_partitions = [
            TopicPartition(topic, p, 0) for p in topic_meta.partitions.keys()
        ]
        consumer.assign(topic_partitions)

    messages = []
    start_time = time.time()

    while len(messages) < expected_count and (time.time() - start_time) < timeout:
        msg = consumer.poll(timeout=1.0)
        if msg is None:
            continue
        if msg.error():
            if msg.error().code() == KafkaError._PARTITION_EOF:
                continue
            break
        messages.append(msg)

    return messages


def test_build_consumer_config_for_read():
    """Test read config builder."""
    bootstrap_servers = ["localhost:9092"]

    # Test basic config
    config = _build_consumer_config_for_read(bootstrap_servers, None)
    assert config["bootstrap.servers"] == "localhost:9092"
    assert config["enable.auto.commit"] is False
    assert "group.id" in config

    # Test with authentication via consumer_config
    user_conf = {
        "security.protocol": "SASL_SSL",
        "sasl.mechanism": "PLAIN",
        "sasl.username": "user",
        "sasl.password": "pass",
    }
    config_with_auth = _build_consumer_config_for_read(
        bootstrap_servers, None, user_conf
    )
    assert config_with_auth["security.protocol"] == "SASL_SSL"
    assert config_with_auth["sasl.mechanism"] == "PLAIN"
    assert config_with_auth["sasl.username"] == "user"
    assert config_with_auth["sasl.password"] == "pass"


def test_build_consumer_config_with_pass_through():
    """Test that extra consumer_config options pass through and cannot override bootstrap.servers."""
    bootstrap_servers = ["localhost:9092"]

    # Extra options should pass through
    extra = {
        "ssl.endpoint.identification.algorithm": "none",
        "group.id": "custom-group",
        "enable.auto.commit": True,
    }
    config = _build_consumer_config_for_read(bootstrap_servers, None, extra)
    assert config["bootstrap.servers"] == "localhost:9092"
    assert config["ssl.endpoint.identification.algorithm"] == "none"
    assert config["group.id"] == "custom-group"
    assert config["enable.auto.commit"] is True

    # Attempt to override bootstrap.servers should be ignored
    override = {"bootstrap.servers": "override:9092"}
    config2 = _build_consumer_config_for_read(bootstrap_servers, None, override)
    assert config2["bootstrap.servers"] == "localhost:9092"


def test_read_kafka_config_conflict_raises():
    """Specifying both kafka_auth_config and consumer_config should error."""
    with pytest.raises(
        ValueError, match="Provide only one of kafka_auth_config.* or consumer_config"
    ):
        ray.data.read_kafka(
            topics="t",
            bootstrap_servers="localhost:9092",
            kafka_auth_config=KafkaAuthConfig(security_protocol="SSL"),
            consumer_config={"security.protocol": "SSL"},
        )


def test_build_consumer_config_with_kafka_auth_config_deprecated():
    """Test kafka-python style KafkaAuthConfig mapping (deprecated path)."""
    bootstrap_servers = ["localhost:9092"]
    auth = KafkaAuthConfig(
        security_protocol="SASL_SSL",
        sasl_mechanism="SCRAM-SHA-256",
        sasl_plain_username="testuser",
        sasl_plain_password="testpass",
        sasl_kerberos_name="kafka/hostname@REALM",
        sasl_kerberos_service_name="kafka",
        # These are unsupported and should be ignored with warnings
        sasl_kerberos_domain_name="example.com",
        sasl_oauth_token_provider=object(),
        ssl_context=object(),
        # ssl_check_hostname False is unsafe to map; ensure not weakening
        ssl_check_hostname=False,
        # SSL files
        ssl_cafile="/path/to/ca.pem",
        ssl_certfile="/path/to/cert.pem",
        ssl_keyfile="/path/to/key.pem",
        ssl_password="keypassword",
        ssl_ciphers="HIGH:!aNULL",
        ssl_crlfile="/path/to/crl.pem",
    )
    config = _build_consumer_config_for_read(bootstrap_servers, auth, None)
    assert config["security.protocol"] == "SASL_SSL"
    assert config["sasl.mechanism"] == "SCRAM-SHA-256"
    assert config["sasl.username"] == "testuser"
    assert config["sasl.password"] == "testpass"
    assert config["sasl.kerberos.principal"] == "kafka/hostname@REALM"
    assert config["sasl.kerberos.service.name"] == "kafka"
    # No weakening of TLS verification
    assert "enable.ssl.certificate.verification" not in config
    assert config["ssl.ca.location"] == "/path/to/ca.pem"
    assert config["ssl.certificate.location"] == "/path/to/cert.pem"
    assert config["ssl.key.location"] == "/path/to/key.pem"
    assert config["ssl.key.password"] == "keypassword"
    assert config["ssl.cipher.suites"] == "HIGH:!aNULL"
    assert config["ssl.crl.location"] == "/path/to/crl.pem"


def test_datetime_to_ms_without_timezone():
    """Test that datetimes without timezone info are treated as UTC."""
    assert _datetime_to_ms(datetime(1970, 1, 1, 0, 0, 0)) == 0
    assert _datetime_to_ms(datetime(2025, 1, 1, 0, 0, 0)) == 1735689600000


def test_datetime_to_ms_with_timezone():
    """Test that timezone-aware datetimes are converted to UTC correctly."""
    from datetime import timezone

    assert _datetime_to_ms(datetime(1970, 1, 1, 0, 0, 0, tzinfo=timezone.utc)) == 0
    assert (
        _datetime_to_ms(datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc))
        == 1735689600000
    )


def test_read_kafka_datetime_validation():
    """Test that start_offset > end_offset with datetimes raises ValueError."""
    with pytest.raises(ValueError, match="start_offset must be less than end_offset"):
        ray.data.read_kafka(
            topics="test-topic",
            bootstrap_servers="localhost:9092",
            start_offset=datetime(2025, 6, 1),
            end_offset=datetime(2025, 1, 1),
        )


# Integration Tests (require Kafka container)


def test_read_kafka_basic(bootstrap_server, kafka_producer, ray_start_regular_shared):
    topic = "test-basic"

    for i in range(100):
        message = {"id": i, "value": f"message-{i}"}
        kafka_producer.produce(topic, value=message, key=f"key-{i}")
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 100)

    # Read from Kafka
    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
    )

    # Verify data
    records = ds.take_all()
    assert len(records) == 100

    # Check first record structure
    first_record = records[0]
    assert "offset" in first_record
    assert "key" in first_record
    assert "value" in first_record
    assert "topic" in first_record
    assert "partition" in first_record
    assert "timestamp" in first_record
    assert first_record["topic"] == topic

    # Verify data types: key is bytes, value is binary
    assert isinstance(first_record["key"], bytes)
    assert isinstance(first_record["value"], bytes)
    key_str = first_record["key"].decode("utf-8")
    assert key_str.startswith("key-")
    value_obj = json.loads(first_record["value"].decode("utf-8"))
    assert "id" in value_obj
    assert "value" in value_obj


@pytest.mark.parametrize(
    "total_messages,start_offset,end_offset,expected_count,test_id",
    [
        (100, 20, 80, 60, "both-set"),
        (100, 50, None, 50, "start-offset-only"),
        (100, None, 50, 50, "end-offset-only"),
        (100, None, None, 100, "both-none"),
        (100, "earliest", 30, 30, "earliest-start-offset-number-end-offset"),
        (100, 50, "latest", 50, "number-start-offset-number-end-offset"),
    ],
)
def test_read_kafka_with_offsets(
    bootstrap_server,
    kafka_producer,
    ray_start_regular_shared,
    total_messages,
    start_offset,
    end_offset,
    expected_count,
    test_id,
):
    topic = f"test-{test_id}"

    for i in range(total_messages):
        message = {"id": i, "value": f"message-{i}"}
        kafka_producer.produce(topic, value=message)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, total_messages)

    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset=start_offset,
        end_offset=end_offset,
    )

    records = ds.take_all()
    assert len(records) == expected_count


def test_read_kafka_multiple_partitions(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    from confluent_kafka.admin import AdminClient, NewTopic

    topic = "test-multi-partition"

    # Create topic with 3 partitions
    admin_client = AdminClient({"bootstrap.servers": bootstrap_server})
    topic_config = NewTopic(topic, num_partitions=3, replication_factor=1)
    admin_client.create_topics([topic_config])

    time.sleep(2)  # Wait for topic creation

    # Send messages to different partitions
    for i in range(150):
        message = {"id": i, "value": f"message-{i}"}
        partition = i % 3
        kafka_producer.produce(topic, value=message, partition=partition)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 150)

    # Read from all partitions
    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
    )

    records = ds.take_all()
    assert len(records) == 150


def test_read_kafka_multiple_topics(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    topic1 = "test-multi-topic-1"
    topic2 = "test-multi-topic-2"

    # Send messages to topic1
    for i in range(50):
        message = {"id": i, "value": f"topic1-message-{i}"}
        kafka_producer.produce(topic1, value=message)

    # Send messages to topic2
    for i in range(30):
        message = {"id": i, "value": f"topic2-message-{i}"}
        kafka_producer.produce(topic2, value=message)

    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic1, 50)
    _wait_for_watermark(bootstrap_server, topic2, 30)

    # Read from both topics
    ds = ray.data.read_kafka(
        topics=[topic1, topic2],
        bootstrap_servers=[bootstrap_server],
    )

    records = ds.take_all()
    assert len(records) == 80


def test_read_kafka_with_message_headers(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    topic = "test-headers"

    for i in range(10):
        message = {"id": i, "value": f"message-{i}"}
        headers = [
            ("header1", b"value1"),
            ("header2", f"value-{i}".encode("utf-8")),
        ]
        kafka_producer.produce(topic, value=message, headers=headers)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 10)

    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
    )

    records = ds.take_all()
    assert len(records) == 10

    first_record = records[0]
    assert "headers" in first_record
    assert isinstance(first_record["headers"]["header1"], bytes)
    assert first_record["headers"]["header1"].decode("utf-8") == "value1"


@pytest.mark.parametrize(
    "start_offset,end_offset,expected_count,test_id",
    [
        (0, 150, 100, "end-offset-exceeds-available-messages"),
        (
            "earliest",
            150,
            100,
            "earliest-start-offset-end-offset-exceeds-available-messages",
        ),
    ],
)
def test_read_kafka_offset_exceeds_available_messages(
    bootstrap_server,
    kafka_producer,
    ray_start_regular_shared,
    start_offset,
    end_offset,
    expected_count,
    test_id,
):
    topic = f"test-offset-timeout-{test_id}"

    for i in range(100):
        message = {"id": i, "value": f"message-{i}"}
        kafka_producer.produce(topic, value=message)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 100)

    # end_offset exceeds available messages, but it gets clamped to the high
    # watermark during offset resolution, so the read completes without
    # hanging.
    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset=start_offset,
        end_offset=end_offset,
    )

    records = ds.take_all()

    assert len(records) == expected_count


def test_read_kafka_default_no_timeout(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    topic = "test-default-no-timeout"

    for i in range(50):
        message = {"id": i, "value": f"message-{i}"}
        kafka_producer.produce(topic, value=message)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 50)

    # timeout_ms is intentionally omitted (defaults to None).
    # Verifies the no-timeout code path works correctly.
    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset=0,
        end_offset=50,
    )

    records = ds.take_all()

    assert len(records) == 50


def test_read_kafka_invalid_topic(bootstrap_server, ray_start_regular_shared):
    with pytest.raises(ValueError, match="has no partitions or doesn't exist"):
        ds = ray.data.read_kafka(
            topics=["non-existent-topic"],
            bootstrap_servers=[bootstrap_server],
        )
        ds.take_all()


@pytest.mark.parametrize(
    "start_offset,end_offset,expected_error,topic",
    [
        (0, "earliest", "end_offset cannot be 'earliest'", "test-invalid-offsets-0"),
        ("latest", 1000, "start_offset cannot be 'latest'", "test-invalid-offsets-1"),
        (80, 20, "start_offset must be less than end_offset", "test-invalid-offsets-2"),
        (
            150,
            "latest",
            r"start_offset \(150\) > end_offset \(latest \(resolved to 100\)\) for partition 0 in topic test-invalid-offsets-3",
            "test-invalid-offsets-3",
        ),
        (
            150,
            200,
            r"start_offset \(150\) > end_offset \(200 \(resolved to 100\)\) for partition 0 in topic test-invalid-offsets-4",
            "test-invalid-offsets-4",
        ),
    ],
)
def test_read_kafka_invalid_offsets(
    bootstrap_server,
    kafka_producer,
    ray_start_regular_shared,
    start_offset,
    end_offset,
    expected_error,
    topic,
):
    for i in range(100):
        message = {"id": i, "value": f"message-{i}"}
        kafka_producer.produce(topic, value=message)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 100)

    with pytest.raises(ValueError, match=expected_error):
        ds = ray.data.read_kafka(
            topics=[topic],
            bootstrap_servers=[bootstrap_server],
            start_offset=start_offset,
            end_offset=end_offset,
        )
        ds.take_all()


def test_read_kafka_with_datetime_offsets(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Test reading Kafka messages using datetime-based start and end offsets."""
    topic = "test-datetime-offsets"

    now = datetime.now(timezone.utc)
    time_before = now - timedelta(hours=1)
    time_after = now + timedelta(hours=1)
    msg_ts = _datetime_to_ms(now)

    for i in range(3):
        kafka_producer.produce(topic, value={"id": i}, timestamp=msg_ts)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 3)

    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset=time_before,
        end_offset=time_after,
    )
    records = ds.take_all()
    assert len(records) == 3


def test_read_kafka_datetime_partial_range(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Test that only messages within the datetime range are returned."""
    topic = "test-datetime-partial-range"

    # First batch at Jan 1 2025, second batch at Feb 1 2025
    batch1_ts = _datetime_to_ms(datetime(2025, 1, 1))
    batch2_ts = _datetime_to_ms(datetime(2025, 2, 1))
    boundary_time = datetime(2025, 1, 15)  # Between the two batches

    kafka_producer.produce(
        topic,
        value={"batch": 1, "id": 0},
        timestamp=batch1_ts,
    )
    kafka_producer.produce(
        topic,
        value={"batch": 1, "id": 1},
        timestamp=batch1_ts,
    )
    kafka_producer.produce(
        topic,
        value={"batch": 2, "id": 0},
        timestamp=batch2_ts,
    )
    kafka_producer.produce(
        topic,
        value={"batch": 2, "id": 1},
        timestamp=batch2_ts,
    )
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 4)

    # Read only the second batch using boundary_time as start
    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset=boundary_time,
        end_offset="latest",
    )
    records = ds.take_all()
    assert len(records) == 2

    # Read only the first batch using boundary_time as end
    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset="earliest",
        end_offset=boundary_time,
    )
    records = ds.take_all()
    assert len(records) == 2


def test_read_kafka_datetime_after_all_messages(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Test datetime start_offset after all messages returns 0 rows."""
    topic = "test-datetime-after-all"

    kafka_producer.produce(topic, value={"id": 0})
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 1)

    future_time = datetime(2099, 1, 1)

    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset=future_time,
        end_offset="latest",
    )
    records = ds.take_all()
    assert len(records) == 0


def test_read_kafka_datetime_before_all_messages(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Test datetime end_offset before all messages returns 0 rows."""
    topic = "test-datetime-before-all"

    kafka_producer.produce(topic, value={"id": 0})
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 1)

    past_time = datetime(1970, 1, 2)

    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset="earliest",
        end_offset=past_time,
    )
    records = ds.take_all()
    # All messages have timestamps after 1970, so offsets_for_times will
    # return the first offset. This means end_offset resolves to the
    # beginning, yielding 0 rows.
    assert len(records) == 0


# Kafka Datasink Unit Tests


def test_kafka_datasink_initialization():
    """Test KafkaDatasink initialization and validation."""
    # Valid initialization
    sink = KafkaDatasink(
        topic="test-topic",
        bootstrap_servers="localhost:9092",
        key_field="id",
        key_serializer="string",
        value_serializer="json",
    )
    assert sink.topic == "test-topic"
    assert sink.bootstrap_servers == "localhost:9092"
    assert sink.key_field == "id"
    assert sink.key_serializer == "string"
    assert sink.value_serializer == "json"

    # Invalid key serializer
    with pytest.raises(ValueError, match="key_serializer must be one of"):
        KafkaDatasink(
            topic="test-topic",
            bootstrap_servers="localhost:9092",
            key_serializer="invalid",
        )

    # Invalid value serializer
    with pytest.raises(ValueError, match="value_serializer must be one of"):
        KafkaDatasink(
            topic="test-topic",
            bootstrap_servers="localhost:9092",
            value_serializer="invalid",
        )


def test_kafka_datasink_row_to_dict():
    """Test _row_to_dict conversion for different row types."""
    sink = KafkaDatasink(topic="test", bootstrap_servers="localhost:9092")

    # Test with dict
    dict_row = {"a": 1, "b": 2}
    assert sink._row_to_dict(dict_row) == dict_row

    # Test with object that has as_pydict
    class MockArrowRow:
        def as_pydict(self):
            return {"x": 10, "y": 20}

    arrow_row = MockArrowRow()
    assert sink._row_to_dict(arrow_row) == {"x": 10, "y": 20}

    # Test with NamedTuple
    from collections import namedtuple

    Point = namedtuple("Point", ["x", "y"])
    point = Point(x=5, y=10)
    result = sink._row_to_dict(point)
    assert result == {"x": 5, "y": 10}

    # Test with primitive
    assert sink._row_to_dict("string") == "string"
    assert sink._row_to_dict(123) == 123


def test_kafka_datasink_serialize_value():
    """Test value serialization for different formats."""
    # JSON serializer
    sink_json = KafkaDatasink(
        topic="test", bootstrap_servers="localhost:9092", value_serializer="json"
    )
    result = sink_json._serialize_value({"key": "value"})
    assert result == b'{"key": "value"}'

    # String serializer
    sink_string = KafkaDatasink(
        topic="test", bootstrap_servers="localhost:9092", value_serializer="string"
    )
    result = sink_string._serialize_value({"key": "value"})
    assert result == b"{'key': 'value'}"

    # Bytes serializer
    sink_bytes = KafkaDatasink(
        topic="test", bootstrap_servers="localhost:9092", value_serializer="bytes"
    )
    result = sink_bytes._serialize_value(b"raw bytes")
    assert result == b"raw bytes"


def test_kafka_datasink_serialize_key():
    """Test key serialization for different formats."""
    # JSON serializer
    sink_json = KafkaDatasink(
        topic="test", bootstrap_servers="localhost:9092", key_serializer="json"
    )
    result = sink_json._serialize_key({"id": 123})
    assert result == b'{"id": 123}'

    # String serializer (default)
    sink_string = KafkaDatasink(
        topic="test", bootstrap_servers="localhost:9092", key_serializer="string"
    )
    result = sink_string._serialize_key(456)
    assert result == b"456"

    # Bytes serializer
    sink_bytes = KafkaDatasink(
        topic="test", bootstrap_servers="localhost:9092", key_serializer="bytes"
    )
    result = sink_bytes._serialize_key(b"key-bytes")
    assert result == b"key-bytes"


def test_kafka_datasink_extract_key():
    """Test key extraction from rows."""
    sink = KafkaDatasink(
        topic="test",
        bootstrap_servers="localhost:9092",
        key_field="user_id",
        key_serializer="string",
    )

    # Test with dict containing key_field
    row = {"user_id": 123, "name": "Alice"}
    key = sink._extract_key(row)
    assert key == b"123"

    # Test with dict without key_field
    row_no_key = {"name": "Bob"}
    key = sink._extract_key(row_no_key)
    assert key is None

    # Test with no key_field configured
    sink_no_key = KafkaDatasink(topic="test", bootstrap_servers="localhost:9092")
    key = sink_no_key._extract_key({"user_id": 456})
    assert key is None


def test_kafka_datasink_extract_key_uses_serializer():
    """Test that _extract_key properly uses the configured key_serializer."""
    # JSON serializer should produce valid JSON
    sink_json = KafkaDatasink(
        topic="test",
        bootstrap_servers="localhost:9092",
        key_field="key",
        key_serializer="json",
    )
    # Test with dict key value - JSON uses double quotes, str() uses single quotes
    row_dict = {"key": {"nested": "value"}, "data": "test"}
    key = sink_json._extract_key(row_dict)
    assert key == b'{"nested": "value"}'  # JSON format with double quotes

    # Test with string that needs JSON escaping
    row_str = {"key": 'hello "world"', "data": "test"}
    key = sink_json._extract_key(row_str)
    assert key == b'"hello \\"world\\""'  # Properly JSON-escaped

    # Bytes serializer should preserve bytes
    sink_bytes = KafkaDatasink(
        topic="test",
        bootstrap_servers="localhost:9092",
        key_field="key",
        key_serializer="bytes",
    )
    row_bytes = {"key": b"raw-bytes", "data": "test"}
    key = sink_bytes._extract_key(row_bytes)
    assert key == b"raw-bytes"


@pytest.fixture
def mock_write_env():
    from types import SimpleNamespace
    from unittest.mock import MagicMock, patch

    sink = KafkaDatasink(topic="test-topic", bootstrap_servers="localhost:9092")
    mock_producer = MagicMock()
    mock_producer.flush.return_value = 0

    mock_block = MagicMock()
    mock_accessor = MagicMock()
    mock_accessor.iter_rows.return_value = [{"id": 1}]

    with (
        patch("confluent_kafka.Producer", return_value=mock_producer),
        patch(
            "ray.data._internal.datasource.kafka_datasink.BlockAccessor.for_block",
            return_value=mock_accessor,
        ),
    ):
        yield SimpleNamespace(
            sink=sink,
            producer=mock_producer,
            block=mock_block,
            accessor=mock_accessor,
            ctx=MagicMock(),
        )


def test_write_delivery_failure_raises(mock_write_env):
    """Test that any delivery failure raises RuntimeError."""
    from unittest.mock import MagicMock

    from confluent_kafka import KafkaError

    env = mock_write_env
    mock_kafka_error = MagicMock(spec=KafkaError)
    mock_kafka_error.str.return_value = "MSG_SIZE_TOO_LARGE"

    def fake_produce(topic, value, key, on_delivery):
        on_delivery(mock_kafka_error, None)

    env.producer.produce.side_effect = fake_produce

    with pytest.raises(RuntimeError, match="Failed to write 1 out of 1"):
        env.sink.write([env.block], ctx=env.ctx)


def test_write_in_flight_after_flush_raises(mock_write_env):
    """Test that messages stuck in-flight after flush raises RuntimeError."""
    env = mock_write_env
    env.producer.produce.return_value = None
    env.producer.flush.return_value = 1  # 1 message still in-flight

    with pytest.raises(
        RuntimeError,
        match="1 out of 1 messages were still in-flight after flush timeout",
    ):
        env.sink.write([env.block], ctx=env.ctx)


def test_write_buffer_error_retry(mock_write_env):
    """Test that BufferError triggers poll and retry, then succeeds."""
    env = mock_write_env
    # First produce() raises BufferError, retry succeeds
    env.producer.produce.side_effect = [BufferError("queue full"), None]

    result = env.sink.write([env.block], ctx=env.ctx)

    assert result["total_records"] == 1
    assert result["failed_messages"] == 0


def test_write_buffer_error_persistent_raises(mock_write_env):
    """Test that persistent BufferError (retry also fails) raises RuntimeError."""
    env = mock_write_env
    env.producer.produce.side_effect = BufferError("queue full")

    with pytest.raises(RuntimeError, match="producer queue is still full"):
        env.sink.write([env.block], ctx=env.ctx)


def test_write_success_returns_stats(mock_write_env):
    """Test successful write returns correct stats."""
    env = mock_write_env
    env.producer.produce.return_value = None
    env.accessor.iter_rows.return_value = [{"id": i} for i in range(3)]

    result = env.sink.write([env.block], ctx=env.ctx)

    assert result == {"total_records": 3, "failed_messages": 0}


# Kafka Datasink Integration Tests (require Kafka container)


def test_write_kafka_basic(bootstrap_server, kafka_consumer, ray_start_regular_shared):
    """Test basic write to Kafka."""
    topic = "test-write-basic"

    # Create dataset
    ds = ray.data.range(100)

    # Write to Kafka
    ds.write_kafka(topic=topic, bootstrap_servers=bootstrap_server)

    # Consume and verify
    messages = consume_messages(kafka_consumer, topic, expected_count=100)
    assert len(messages) == 100

    # Verify message structure
    first_msg = messages[0]
    assert first_msg.topic() == topic
    assert first_msg.key() is None  # No key field specified

    # Verify value is JSON encoded
    value = json.loads(first_msg.value().decode("utf-8"))
    assert "id" in value
    assert isinstance(value["id"], int)


def test_write_kafka_with_keys(
    bootstrap_server, kafka_consumer, ray_start_regular_shared
):
    """Test writing to Kafka with keys."""
    topic = "test-write-with-keys"

    # Create dataset with id field
    data = [{"id": i, "name": f"user-{i}", "value": i * 10} for i in range(50)]
    ds = ray.data.from_items(data)

    # Write to Kafka with id as key
    ds.write_kafka(
        topic=topic,
        bootstrap_servers=bootstrap_server,
        key_field="id",
        key_serializer="string",
    )

    # Consume and verify
    messages = consume_messages(kafka_consumer, topic, expected_count=50)
    assert len(messages) == 50

    # Verify keys
    for msg in messages:
        assert msg.key() is not None
        key_str = msg.key().decode("utf-8")
        assert key_str.isdigit()

        # Verify value
        value = json.loads(msg.value().decode("utf-8"))
        assert value["id"] == int(key_str)


@pytest.mark.parametrize(
    "key_serializer,value_serializer",
    [
        ("string", "json"),
        ("json", "string"),
        ("string", "string"),
        ("json", "json"),
    ],
)
def test_write_kafka_serializers(
    bootstrap_server,
    kafka_consumer,
    ray_start_regular_shared,
    key_serializer,
    value_serializer,
):
    """Test different serializer combinations."""
    topic = f"test-serializers-{key_serializer}-{value_serializer}"

    # Create dataset
    data = [{"id": i, "message": f"msg-{i}"} for i in range(20)]
    ds = ray.data.from_items(data)

    # Write with specified serializers
    ds.write_kafka(
        topic=topic,
        bootstrap_servers=bootstrap_server,
        key_field="id",
        key_serializer=key_serializer,
        value_serializer=value_serializer,
    )

    # Consume and verify
    messages = consume_messages(kafka_consumer, topic, expected_count=20)
    assert len(messages) == 20

    # Verify first message can be deserialized correctly
    first_msg = messages[0]

    if key_serializer == "json":
        key_data = json.loads(first_msg.key().decode("utf-8"))
        assert isinstance(key_data, int)
    else:  # string
        key_str = first_msg.key().decode("utf-8")
        assert key_str.isdigit()

    if value_serializer == "json":
        value_data = json.loads(first_msg.value().decode("utf-8"))
        assert "id" in value_data
        assert "message" in value_data
    else:  # string
        value_str = first_msg.value().decode("utf-8")
        assert "id" in value_str
        assert "message" in value_str


def test_write_kafka_multiple_blocks(
    bootstrap_server, kafka_consumer, ray_start_regular_shared
):
    """Test writing dataset with multiple blocks."""
    topic = "test-write-multiple-blocks"

    # Create dataset and repartition to ensure multiple blocks
    ds = ray.data.range(200).repartition(5)

    # Write to Kafka
    ds.write_kafka(topic=topic, bootstrap_servers=bootstrap_server)

    # Consume and verify
    messages = consume_messages(kafka_consumer, topic, expected_count=200)
    assert len(messages) == 200

    # Verify all ids are present
    ids = set()
    for msg in messages:
        value = json.loads(msg.value().decode("utf-8"))
        ids.add(value["id"])

    assert len(ids) == 200
    assert ids == set(range(200))


def test_write_kafka_empty_dataset(
    bootstrap_server, kafka_consumer, ray_start_regular_shared
):
    """Test writing an empty dataset."""
    topic = "test-write-empty"

    # Create empty dataset
    ds = ray.data.from_items([])

    # Write to Kafka (should succeed without errors)
    ds.write_kafka(topic=topic, bootstrap_servers=bootstrap_server)

    # Try to consume (should get no messages)
    messages = consume_messages(kafka_consumer, topic, expected_count=0, timeout=3)
    assert len(messages) == 0


def test_write_kafka_with_producer_config(
    bootstrap_server, kafka_consumer, ray_start_regular_shared
):
    """Test writing with custom producer configuration."""
    topic = "test-write-producer-config"

    # Create dataset
    ds = ray.data.range(30)

    # Write with custom producer config (confluent-kafka/librdkafka format)
    ds.write_kafka(
        topic=topic,
        bootstrap_servers=bootstrap_server,
        producer_config={
            "acks": "all",
            "retries": 3,
            "max.in.flight.requests.per.connection": 1,
        },
    )

    # Consume and verify
    messages = consume_messages(kafka_consumer, topic, expected_count=30)
    assert len(messages) == 30


def test_write_kafka_with_complex_data(
    bootstrap_server, kafka_consumer, ray_start_regular_shared
):
    """Test writing complex nested data structures."""
    topic = "test-write-complex"

    # Create dataset with nested structures
    data = [
        {
            "id": i,
            "user": {"name": f"user-{i}", "email": f"user{i}@example.com"},
            "tags": [f"tag{j}" for j in range(3)],
            "metadata": {"created": "2024-01-01", "score": i * 1.5},
        }
        for i in range(15)
    ]
    ds = ray.data.from_items(data)

    # Write to Kafka
    ds.write_kafka(topic=topic, bootstrap_servers=bootstrap_server)

    # Consume and verify
    messages = consume_messages(kafka_consumer, topic, expected_count=15)
    assert len(messages) == 15

    # Verify nested structure is preserved
    first_value = json.loads(messages[0].value().decode("utf-8"))
    assert "user" in first_value
    assert "name" in first_value["user"]
    assert "tags" in first_value
    assert isinstance(first_value["tags"], list)
    assert "metadata" in first_value


def test_write_kafka_invalid_bootstrap_server(ray_start_regular_shared):
    """Test error handling with invalid bootstrap server."""
    topic = "test-invalid-server"

    ds = ray.data.range(10)

    # Use a short message.timeout.ms so librdkafka gives up quickly
    # instead of waiting the full 30s flush timeout.
    with pytest.raises(Exception):
        ds.write_kafka(
            topic=topic,
            bootstrap_servers="invalid-server:9999",
            producer_config={"message.timeout.ms": 3000},
        )


def test_write_kafka_dataset_with_nulls(
    bootstrap_server, kafka_consumer, ray_start_regular_shared
):
    """Test writing dataset with null/None values."""
    topic = "test-write-nulls"

    # Create dataset with None values
    data = [{"id": i, "value": f"val-{i}" if i % 2 == 0 else None} for i in range(20)]
    ds = ray.data.from_items(data)

    # Write to Kafka
    ds.write_kafka(topic=topic, bootstrap_servers=bootstrap_server)

    # Consume and verify
    messages = consume_messages(kafka_consumer, topic, expected_count=20)
    assert len(messages) == 20

    # Verify None values are serialized
    for msg in messages:
        value = json.loads(msg.value().decode("utf-8"))
        assert "id" in value
        # value["value"] should be either a string or null


@pytest.mark.parametrize(
    "start_offset,expected_error",
    [
        (
            {"my-topic": {0: "latest"}},
            r"start_offset\['my-topic'\]\[0\] cannot be 'latest'",
        ),
        (
            {"my-topic": {"not-an-int": 100}},
            r"start_offset\['my-topic'\] keys must be integers",
        ),
        (
            {"my-topic": "not-a-dict"},
            r"start_offset\['my-topic'\] must be a dict",
        ),
    ],
)
def test_per_partition_start_offset_invalid_values(start_offset, expected_error):
    """Per-partition start_offset with disallowed values raises ValueError at init."""
    with pytest.raises(ValueError, match=expected_error):
        ray.data.read_kafka(
            topics="my-topic",
            bootstrap_servers="localhost:9092",
            start_offset=start_offset,
        )


@pytest.mark.parametrize(
    "end_offset,expected_error",
    [
        (
            {"my-topic": {0: "earliest"}},
            r"end_offset\['my-topic'\]\[0\] cannot be 'earliest'",
        ),
        (
            {"my-topic": {"not-an-int": 100}},
            r"end_offset\['my-topic'\] keys must be integers",
        ),
        (
            {"my-topic": "not-a-dict"},
            r"end_offset\['my-topic'\] must be a dict",
        ),
    ],
)
def test_per_partition_end_offset_invalid_values(end_offset, expected_error):
    """Per-partition end_offset with disallowed values raises ValueError at init."""
    with pytest.raises(ValueError, match=expected_error):
        ray.data.read_kafka(
            topics="my-topic",
            bootstrap_servers="localhost:9092",
            end_offset=end_offset,
        )


def test_per_partition_start_offset_non_existent_partition(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Per-partition dict referencing a non-existent partition raises ValueError."""
    topic = "test-per-partition-bad-partition"

    for i in range(10):
        message = {"id": i, "value": f"message-{i}"}
        kafka_producer.produce(topic, value=message)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 10)

    with pytest.raises(
        ValueError,
        match=r"start_offset references partition 99 in topic",
    ):
        ds = ray.data.read_kafka(
            topics=[topic],
            bootstrap_servers=[bootstrap_server],
            start_offset={topic: {99: 0}},
        )
        ds.take_all()


def test_per_partition_start_offset_specific_offsets(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Per-partition start_offset reads correct slice from each partition."""
    from confluent_kafka.admin import AdminClient, NewTopic

    topic = "test-per-partition-start-offset"

    admin_client = AdminClient({"bootstrap.servers": bootstrap_server})
    topic_config = NewTopic(topic, num_partitions=2, replication_factor=1)
    admin_client.create_topics([topic_config])
    time.sleep(2)  # Wait for topic creation

    # Send 50 messages to partition 0 and 50 to partition 1
    for i in range(50):
        kafka_producer.produce(topic, value={"id": i}, partition=0)
        kafka_producer.produce(topic, value={"id": i}, partition=1)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 100)

    # Read partition 0 from offset 20, partition 1 from offset 40
    # Expected: 30 messages from partition 0 + 10 messages from partition 1 = 40
    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset={topic: {0: 20, 1: 40}},
        end_offset="latest",
    )

    records = ds.take_all()
    assert len(records) == 40


def test_per_partition_end_offset_specific_offsets(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Per-partition end_offset reads correct slice from each partition."""
    from confluent_kafka.admin import AdminClient, NewTopic

    topic = "test-per-partition-end-offset"

    admin_client = AdminClient({"bootstrap.servers": bootstrap_server})
    topic_config = NewTopic(topic, num_partitions=2, replication_factor=1)
    admin_client.create_topics([topic_config])
    time.sleep(2)  # Wait for topic creation

    # Send 50 messages to partition 0 and 50 to partition 1
    for i in range(50):
        kafka_producer.produce(topic, value={"id": i}, partition=0)
        kafka_producer.produce(topic, value={"id": i}, partition=1)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 100)

    # End partition 0 at offset 30, partition 1 at offset 20
    # Expected: 30 messages from partition 0 + 20 messages from partition 1 = 50
    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset="earliest",
        end_offset={topic: {0: 30, 1: 20}},
    )

    records = ds.take_all()
    assert len(records) == 50


def test_per_partition_start_offset_fallback_to_earliest(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Partitions not listed in per-partition dict fall back to 'earliest'."""
    from confluent_kafka.admin import AdminClient, NewTopic

    topic = "test-per-partition-fallback"

    admin_client = AdminClient({"bootstrap.servers": bootstrap_server})
    topic_config = NewTopic(topic, num_partitions=2, replication_factor=1)
    admin_client.create_topics([topic_config])
    time.sleep(2)  # Wait for topic creation

    # Send 50 messages to each partition
    for i in range(50):
        kafka_producer.produce(topic, value={"id": i}, partition=0)
        kafka_producer.produce(topic, value={"id": i}, partition=1)
    kafka_producer.flush()
    _wait_for_watermark(bootstrap_server, topic, 100)

    # Only specify offset for partition 0; partition 1 should fall back to earliest (0)
    # Expected: 30 messages from partition 0 + 50 messages from partition 1 = 80
    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset={topic: {0: 20}},
        end_offset="latest",
    )

    records = ds.take_all()
    assert len(records) == 80


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))