"""
KafkaSaver — LangGraph checkpoint backend backed by Apache Kafka.

Architecture:
  • Write path: put() and put_writes() produce JSON messages to topics
    'checkpoints' and 'checkpoint_writes'. This is the hot path and is
    optimised for maximum throughput (async delivery, lz4 compression).

  • Read path: get_tuple() and list() are served from an in-memory dict
    that is seeded by consuming all existing messages on setup(). This
    mirrors event-sourced systems where Kafka is the durable log and a
    secondary projection handles reads.
"""

import json
import time
import logging
from contextlib import contextmanager
from typing import Any, Iterator, Optional, Sequence, Tuple, List

from confluent_kafka import Producer, Consumer, KafkaError
from confluent_kafka.admin import AdminClient, NewTopic

from langgraph.checkpoint.base import (
    BaseCheckpointSaver,
    Checkpoint,
    CheckpointMetadata,
    CheckpointTuple,
    ChannelVersions,
    PendingWrite,
    WRITES_IDX_MAP,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer

logger = logging.getLogger(__name__)

TOPIC_CHECKPOINTS = "checkpoints"
TOPIC_WRITES = "checkpoint_writes"
_SENTINEL = object()


class KafkaSaver(BaseCheckpointSaver[str]):
    """
    Kafka-backed LangGraph checkpointer.

    Parameters
    ----------
    bootstrap_servers : str
        Kafka broker address, e.g. ``"localhost:9092"``.
    num_partitions : int
        Number of partitions per topic (default 4).
    replication_factor : int
        Replication factor per topic (default 1, fine for single-broker dev).
    producer_config : dict, optional
        Extra ``confluent_kafka.Producer`` configuration to merge in.
    """

    def __init__(
        self,
        bootstrap_servers: str = "localhost:9092",
        *,
        num_partitions: int = 4,
        replication_factor: int = 1,
        producer_config: Optional[dict] = None,
        serde: Optional[Any] = None,
    ) -> None:
        super().__init__(serde=serde or JsonPlusSerializer())
        self.bootstrap_servers = bootstrap_servers
        self.num_partitions = num_partitions
        self.replication_factor = replication_factor

        # State projections (thread_id, checkpoint_ns, checkpoint_id) → data
        self._checkpoints: dict = {}   # (tid, ns, cid) → {checkpoint, metadata, parent, ts}
        # (tid, ns, cid) → list of (task_id, channel, type, value_blob, idx)
        self._writes: dict = {}

        base_producer_cfg = {
            "bootstrap.servers": bootstrap_servers,
            "compression.type": "lz4",
            "queue.buffering.max.ms": "5",        # batch window (ms)
            "queue.buffering.max.messages": 100_000,
            "linger.ms": 5,
            "acks": "1",                           # leader ack only → faster
        }
        if producer_config:
            base_producer_cfg.update(producer_config)
        self._producer = Producer(base_producer_cfg)

    # ------------------------------------------------------------------
    # Lifecycle helpers
    # ------------------------------------------------------------------

    @classmethod
    @contextmanager
    def from_bootstrap_servers(
        cls,
        bootstrap_servers: str = "localhost:9092",
        **kwargs: Any,
    ) -> Iterator["KafkaSaver"]:
        saver = cls(bootstrap_servers=bootstrap_servers, **kwargs)
        try:
            yield saver
        finally:
            saver.close()

    def setup(self) -> None:
        """Create topics if they don't exist, then replay existing messages."""
        admin = AdminClient({"bootstrap.servers": self.bootstrap_servers})

        existing = admin.list_topics(timeout=10).topics.keys()
        topics_to_create = []
        for name in (TOPIC_CHECKPOINTS, TOPIC_WRITES):
            if name not in existing:
                topics_to_create.append(
                    NewTopic(
                        name,
                        num_partitions=self.num_partitions,
                        replication_factor=self.replication_factor,
                        config={"retention.ms": str(7 * 24 * 3600 * 1000)},
                    )
                )
        if topics_to_create:
            futures = admin.create_topics(topics_to_create)
            for topic, future in futures.items():
                try:
                    future.result()
                    logger.info("Created Kafka topic: %s", topic)
                except Exception as exc:
                    logger.warning("Topic creation warning for %s: %s", topic, exc)

        # Small delay to let topic leadership settle
        time.sleep(0.5)

        # Replay all existing messages into local state
        self._replay_topic(TOPIC_CHECKPOINTS, self._handle_checkpoint_msg)
        self._replay_topic(TOPIC_WRITES, self._handle_write_msg)

    def close(self) -> None:
        """Flush and close the producer."""
        self._producer.flush(timeout=10)

    # ------------------------------------------------------------------
    # Internal: message handling
    # ------------------------------------------------------------------

    def _delivery_report(self, err, msg):
        if err:
            logger.error("Kafka delivery failed for %s: %s", msg.topic(), err)

    def _produce(self, topic: str, key: str, value: dict) -> None:
        self._producer.produce(
            topic,
            key=key.encode("utf-8"),
            value=json.dumps(value).encode("utf-8"),
            callback=self._delivery_report,
        )
        # Non-blocking: poll for delivery callbacks without waiting
        self._producer.poll(0)

    def _replay_topic(self, topic: str, handler) -> None:
        """Consume all messages in *topic* from the beginning."""
        consumer = Consumer({
            "bootstrap.servers": self.bootstrap_servers,
            "group.id": f"kafkasaver-replay-{topic}-{time.time_ns()}",
            "auto.offset.reset": "earliest",
            "enable.auto.commit": False,
        })
        consumer.subscribe([topic])

        idle_polls = 0
        max_idle = 5  # stop after 5 consecutive empty polls
        try:
            while idle_polls < max_idle:
                msg = consumer.poll(timeout=0.5)
                if msg is None:
                    idle_polls += 1
                    continue
                if msg.error():
                    if msg.error().code() == KafkaError._PARTITION_EOF:
                        idle_polls += 1
                    else:
                        logger.error("Replay error: %s", msg.error())
                    continue
                idle_polls = 0
                try:
                    data = json.loads(msg.value().decode("utf-8"))
                    handler(data)
                except Exception as exc:
                    logger.warning("Failed to replay message: %s", exc)
        finally:
            consumer.close()

    def _handle_checkpoint_msg(self, data: dict) -> None:
        key = (data["thread_id"], data.get("checkpoint_ns", ""), data["checkpoint_id"])
        self._checkpoints[key] = data

    def _handle_write_msg(self, data: dict) -> None:
        key = (data["thread_id"], data.get("checkpoint_ns", ""), data["checkpoint_id"])
        self._writes.setdefault(key, []).append(data)

    # ------------------------------------------------------------------
    # CheckpointSaver interface
    # ------------------------------------------------------------------

    def put(
        self,
        config: dict,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> dict:
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = checkpoint["id"]
        parent_checkpoint_id = config["configurable"].get("checkpoint_id")

        checkpoint_type, checkpoint_blob = self.serde.dumps_typed(checkpoint)
        metadata_blob = json.dumps(metadata)

        import base64
        msg = {
            "thread_id": thread_id,
            "checkpoint_ns": checkpoint_ns,
            "checkpoint_id": checkpoint_id,
            "parent_checkpoint_id": parent_checkpoint_id,
            "type": checkpoint_type,
            "checkpoint": base64.b64encode(checkpoint_blob).decode("utf-8"),
            "metadata": metadata_blob,
            "ts": time.time(),
        }
        key = (thread_id, checkpoint_ns, checkpoint_id)
        self._checkpoints[key] = msg
        self._produce(TOPIC_CHECKPOINTS, checkpoint_id, msg)

        return {
            "configurable": {
                "thread_id": thread_id,
                "checkpoint_ns": checkpoint_ns,
                "checkpoint_id": checkpoint_id,
            }
        }

    def put_writes(
        self,
        config: dict,
        writes: Sequence[Tuple[str, Any]],
        task_id: str,
        task_path: str = "",
    ) -> None:
        import base64
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = config["configurable"]["checkpoint_id"]

        key = (thread_id, checkpoint_ns, checkpoint_id)
        for idx, (channel, value) in enumerate(writes):
            mapped_idx = WRITES_IDX_MAP.get(channel, idx)
            type_, value_blob = self.serde.dumps_typed(value)
            msg = {
                "thread_id": thread_id,
                "checkpoint_ns": checkpoint_ns,
                "checkpoint_id": checkpoint_id,
                "task_id": task_id,
                "idx": mapped_idx,
                "channel": channel,
                "type": type_,
                "value": base64.b64encode(value_blob).decode("utf-8"),
                "ts": time.time(),
            }
            self._writes.setdefault(key, []).append(msg)
            self._produce(TOPIC_WRITES, f"{checkpoint_id}_{mapped_idx}", msg)

    def get_tuple(self, config: dict) -> Optional[CheckpointTuple]:
        import base64
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = config["configurable"].get("checkpoint_id")

        if checkpoint_id:
            key = (thread_id, checkpoint_ns, checkpoint_id)
            data = self._checkpoints.get(key)
        else:
            # Find the latest checkpoint for this thread
            candidates = [
                (k, v) for k, v in self._checkpoints.items()
                if k[0] == thread_id and k[1] == checkpoint_ns
            ]
            if not candidates:
                return None
            key, data = max(candidates, key=lambda x: x[1]["ts"])
            checkpoint_id = key[2]

        if not data:
            return None

        checkpoint = self.serde.loads_typed((
            data["type"],
            base64.b64decode(data["checkpoint"]),
        ))
        metadata = json.loads(data["metadata"])
        parent_config = None
        if data.get("parent_checkpoint_id"):
            parent_config = {
                "configurable": {
                    "thread_id": thread_id,
                    "checkpoint_ns": checkpoint_ns,
                    "checkpoint_id": data["parent_checkpoint_id"],
                }
            }

        pending_writes = self._get_writes_from_state(thread_id, checkpoint_ns, checkpoint_id)

        return CheckpointTuple(
            config={
                "configurable": {
                    "thread_id": thread_id,
                    "checkpoint_ns": checkpoint_ns,
                    "checkpoint_id": checkpoint_id,
                }
            },
            checkpoint=checkpoint,
            metadata=metadata,
            parent_config=parent_config,
            pending_writes=pending_writes,
        )

    def list(
        self,
        config: Optional[dict],
        *,
        filter: Optional[dict] = None,
        before: Optional[dict] = None,
        limit: Optional[int] = None,
    ) -> Iterator[CheckpointTuple]:
        import base64
        thread_id = config["configurable"]["thread_id"] if config else None
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "") if config else None

        candidates = [
            (k, v) for k, v in self._checkpoints.items()
            if (thread_id is None or k[0] == thread_id)
            and (checkpoint_ns is None or k[1] == checkpoint_ns)
        ]
        candidates.sort(key=lambda x: x[1]["ts"], reverse=True)

        if before and "checkpoint_id" in before["configurable"]:
            before_id = before["configurable"]["checkpoint_id"]
            candidates = [(k, v) for k, v in candidates if k[2] != before_id]

        if limit:
            candidates = candidates[:limit]

        for key, data in candidates:
            tid, ns, cid = key
            checkpoint = self.serde.loads_typed((
                data["type"],
                base64.b64decode(data["checkpoint"]),
            ))
            metadata = json.loads(data["metadata"])
            parent_config = None
            if data.get("parent_checkpoint_id"):
                parent_config = {
                    "configurable": {
                        "thread_id": tid,
                        "checkpoint_ns": ns,
                        "checkpoint_id": data["parent_checkpoint_id"],
                    }
                }
            pending_writes = self._get_writes_from_state(tid, ns, cid)
            yield CheckpointTuple(
                config={"configurable": {"thread_id": tid, "checkpoint_ns": ns, "checkpoint_id": cid}},
                checkpoint=checkpoint,
                metadata=metadata,
                parent_config=parent_config,
                pending_writes=pending_writes,
            )

    # ------------------------------------------------------------------
    # Internal: read-path helpers
    # ------------------------------------------------------------------

    def _get_writes_from_state(
        self, thread_id: str, checkpoint_ns: str, checkpoint_id: str
    ) -> List[PendingWrite]:
        import base64
        key = (thread_id, checkpoint_ns, checkpoint_id)
        rows = self._writes.get(key, [])
        rows_sorted = sorted(rows, key=lambda r: (r["ts"], r["idx"]))
        result = []
        for row in rows_sorted:
            value = self.serde.loads_typed((
                row["type"],
                base64.b64decode(row["value"]),
            ))
            result.append((row["task_id"], row["channel"], value))
        return result
