"""
Tests for KafkaSaver.

Requirements: a running Kafka broker at localhost:9092 (see docker-compose.yml).

Run with:
    pytest tests/test_kafka.py -v
"""

import pytest
import uuid
from langgraph.checkpoint.base import CheckpointMetadata, empty_checkpoint
from langgraph_checkpoint_timeseries import KafkaSaver

KAFKA_BOOTSTRAP = "localhost:9092"


# ──────────────────────────────────────────────────────────────
# Fixtures
# ──────────────────────────────────────────────────────────────

@pytest.fixture(scope="module")
def saver():
    """Module-scoped KafkaSaver — topics are created once, state accumulates."""
    s = KafkaSaver(bootstrap_servers=KAFKA_BOOTSTRAP)
    s.setup()
    yield s
    s.close()


def make_config(thread_id: str | None = None) -> dict:
    return {"configurable": {"thread_id": thread_id or f"test_{uuid.uuid4().hex[:8]}", "checkpoint_ns": ""}}


def make_metadata() -> CheckpointMetadata:
    return CheckpointMetadata(source="input", step=0, run_id="test_run", parents={})


# ──────────────────────────────────────────────────────────────
# Basic put / get_tuple
# ──────────────────────────────────────────────────────────────

def test_put_and_get_tuple(saver: KafkaSaver):
    config = make_config()
    checkpoint = empty_checkpoint()
    metadata = make_metadata()

    returned_config = saver.put(config, checkpoint, metadata, {})

    # get_tuple using the returned config (has checkpoint_id)
    result = saver.get_tuple(returned_config)

    assert result is not None
    assert result.checkpoint["id"] == checkpoint["id"]
    assert result.metadata["source"] == "input"


def test_get_tuple_returns_latest_without_checkpoint_id(saver: KafkaSaver):
    """get_tuple with no checkpoint_id should return the most recent."""
    tid = f"latest_{uuid.uuid4().hex[:8]}"
    config = make_config(tid)
    metadata = make_metadata()

    cp1 = empty_checkpoint()
    cp2 = empty_checkpoint()

    saver.put(config, cp1, metadata, {})
    saver.put(config, cp2, metadata, {})

    # Query without checkpoint_id
    bare_config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
    result = saver.get_tuple(bare_config)

    assert result is not None
    assert result.checkpoint["id"] == cp2["id"]


def test_get_tuple_nonexistent_returns_none(saver: KafkaSaver):
    config = {"configurable": {"thread_id": "does_not_exist_xyz", "checkpoint_ns": "", "checkpoint_id": "fake"}}
    assert saver.get_tuple(config) is None


# ──────────────────────────────────────────────────────────────
# put_writes / pending_writes
# ──────────────────────────────────────────────────────────────

def test_put_writes_and_retrieve(saver: KafkaSaver):
    config = make_config()
    checkpoint = empty_checkpoint()
    metadata = make_metadata()

    returned_config = saver.put(config, checkpoint, metadata, {})
    saver.put_writes(returned_config, [("my_channel", "hello"), ("other_channel", 42)], "task_001")

    result = saver.get_tuple(returned_config)
    assert result is not None
    assert result.pending_writes is not None
    assert len(result.pending_writes) == 2

    channels = {w[1] for w in result.pending_writes}
    assert "my_channel" in channels
    assert "other_channel" in channels


def test_put_writes_values_are_deserialized(saver: KafkaSaver):
    config = make_config()
    checkpoint = empty_checkpoint()
    metadata = make_metadata()

    returned_config = saver.put(config, checkpoint, metadata, {})
    payload = {"key": "value", "nested": [1, 2, 3]}
    saver.put_writes(returned_config, [("data_channel", payload)], "task_002")

    result = saver.get_tuple(returned_config)
    write = next(w for w in result.pending_writes if w[1] == "data_channel")
    assert write[2] == payload


# ──────────────────────────────────────────────────────────────
# list()
# ──────────────────────────────────────────────────────────────

def test_list_returns_all_checkpoints_for_thread(saver: KafkaSaver):
    tid = f"list_test_{uuid.uuid4().hex[:8]}"
    config = make_config(tid)
    metadata = make_metadata()

    n = 5
    for _ in range(n):
        cp = empty_checkpoint()
        saver.put(config, cp, metadata, {})

    listed = list(saver.list(config))
    # May have more from previous test runs on same topic; at minimum n
    assert len(listed) >= n


def test_list_limit(saver: KafkaSaver):
    tid = f"limit_test_{uuid.uuid4().hex[:8]}"
    config = make_config(tid)
    metadata = make_metadata()

    for _ in range(10):
        cp = empty_checkpoint()
        saver.put(config, cp, metadata, {})

    listed = list(saver.list(config, limit=3))
    assert len(listed) == 3


def test_list_ordered_desc_by_timestamp(saver: KafkaSaver):
    tid = f"order_test_{uuid.uuid4().hex[:8]}"
    config = make_config(tid)
    metadata = make_metadata()

    ids = []
    for _ in range(4):
        cp = empty_checkpoint()
        saver.put(config, cp, metadata, {})
        ids.append(cp["id"])

    listed = list(saver.list(config))
    listed_ids = [t.checkpoint["id"] for t in listed]
    # Most recent first — last inserted should be first in list
    assert listed_ids[0] == ids[-1]


# ──────────────────────────────────────────────────────────────
# Parent checkpoint linkage
# ──────────────────────────────────────────────────────────────

def test_parent_checkpoint_id_is_linked(saver: KafkaSaver):
    tid = f"parent_test_{uuid.uuid4().hex[:8]}"
    config = make_config(tid)
    metadata = make_metadata()

    cp1 = empty_checkpoint()
    returned1 = saver.put(config, cp1, metadata, {})

    # Second checkpoint treats first as parent
    cp2 = empty_checkpoint()
    saver.put(returned1, cp2, metadata, {})

    result = saver.get_tuple({"configurable": {"thread_id": tid, "checkpoint_ns": "", "checkpoint_id": cp2["id"]}})
    assert result is not None
    assert result.parent_config is not None
    assert result.parent_config["configurable"]["checkpoint_id"] == cp1["id"]


# ──────────────────────────────────────────────────────────────
# context manager helper
# ──────────────────────────────────────────────────────────────

def test_from_bootstrap_servers_context_manager():
    with KafkaSaver.from_bootstrap_servers(KAFKA_BOOTSTRAP) as s:
        s.setup()
        config = make_config()
        cp = empty_checkpoint()
        metadata = make_metadata()
        returned = s.put(config, cp, metadata, {})
        result = s.get_tuple(returned)
        assert result is not None
