from __future__ import annotations

from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple

from aiel_sdk.errors import RuntimeDependencyError

try:
    from langchain_core.runnables import RunnableConfig
except Exception as e:  # pragma: no cover
    raise RuntimeDependencyError("langchain-core", "pip install 'aiel-sdk[langgraph]'") from e

try:
    from langgraph.checkpoint.base import (
        BaseCheckpointSaver,
        ChannelVersions,
        Checkpoint,
        CheckpointMetadata,
        CheckpointTuple,
        get_checkpoint_id,
    )
except Exception as e:  # pragma: no cover
    raise RuntimeDependencyError("langgraph", "pip install 'aiel-sdk[langgraph]'") from e

from ...memory.saver import Saver


class AIELLangGraphCheckpointer(BaseCheckpointSaver[str]):
    """
    LangGraph checkpointer backed by aiel_sdk.memory.Saver.

    Implementation note:
    - Saver exposes only "latest checkpoint" read semantics.
    - This adapter therefore persists the most recent LangGraph checkpoint per
      (thread_id, checkpoint_ns) and returns that on reads/lists.
    """

    _STATE_KEY = "_langgraph_checkpoint"
    _META_KEY = "_langgraph_metadata"
    _LG_ID_KEY = "_langgraph_checkpoint_id"
    _PARENT_ID_KEY = "_langgraph_parent_checkpoint_id"

    def __init__(self, saver: Saver) -> None:
        super().__init__()
        self._saver = saver

    @staticmethod
    def _configurable(config: RunnableConfig) -> Dict[str, Any]:
        return dict((config or {}).get("configurable") or {})

    @classmethod
    def _thread_ns(cls, config: RunnableConfig) -> Tuple[str, str]:
        configurable = cls._configurable(config)
        thread_id = str(configurable.get("thread_id") or "").strip()
        if not thread_id:
            raise ValueError("LangGraph config.configurable.thread_id is required")
        checkpoint_ns = str(configurable.get("checkpoint_ns") or "")
        return thread_id, checkpoint_ns

    @classmethod
    def _thread_key(cls, config: RunnableConfig) -> str:
        thread_id, checkpoint_ns = cls._thread_ns(config)
        return thread_id if not checkpoint_ns else f"{thread_id}::{checkpoint_ns}"

    @classmethod
    def _checkpoint_config(cls, *, thread_id: str, checkpoint_ns: str, checkpoint_id: Optional[str]) -> RunnableConfig:
        configurable: Dict[str, Any] = {
            "thread_id": thread_id,
            "checkpoint_ns": checkpoint_ns,
        }
        if checkpoint_id:
            configurable["checkpoint_id"] = checkpoint_id
        return {"configurable": configurable}

    @staticmethod
    def _matches_filter(metadata: Dict[str, Any], filter_values: Optional[Dict[str, Any]]) -> bool:
        if not filter_values:
            return True
        for key, expected in filter_values.items():
            if metadata.get(key) != expected:
                return False
        return True

    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
        thread_id, checkpoint_ns = self._thread_ns(config)
        thread_key = self._thread_key(config)
        requested_id = get_checkpoint_id(config)

        checkpoint = self._saver.checkpoint_read_latest(thread_key=thread_key)
        if checkpoint is None:
            return None

        raw_checkpoint = dict(checkpoint.state or {}).get(self._STATE_KEY)
        raw_metadata = dict(checkpoint.metadata or {}).get(self._META_KEY, {})
        langgraph_checkpoint_id = dict(checkpoint.metadata or {}).get(self._LG_ID_KEY)
        parent_id = dict(checkpoint.metadata or {}).get(self._PARENT_ID_KEY)

        if not isinstance(raw_checkpoint, dict):
            return None
        if not isinstance(raw_metadata, dict):
            raw_metadata = {}

        if requested_id and langgraph_checkpoint_id and requested_id != langgraph_checkpoint_id:
            return None

        parent_config = (
            self._checkpoint_config(thread_id=thread_id, checkpoint_ns=checkpoint_ns, checkpoint_id=str(parent_id))
            if parent_id
            else None
        )

        return CheckpointTuple(
            config=self._checkpoint_config(
                thread_id=thread_id,
                checkpoint_ns=checkpoint_ns,
                checkpoint_id=str(langgraph_checkpoint_id or raw_checkpoint.get("id") or ""),
            ),
            checkpoint=raw_checkpoint,
            metadata=raw_metadata,
            parent_config=parent_config,
            pending_writes=[],
        )

    def list(
        self,
        config: Optional[RunnableConfig],
        *,
        filter: Optional[Dict[str, Any]] = None,
        before: Optional[RunnableConfig] = None,
        limit: Optional[int] = None,
    ) -> Iterator[CheckpointTuple]:
        _ = before
        if limit is not None and limit <= 0:
            return
        if config is None:
            return
        out = self.get_tuple(config)
        if out is None:
            return
        if not self._matches_filter(dict(out.metadata or {}), filter):
            return
        yield out

    def put(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        _ = new_versions
        thread_id, checkpoint_ns = self._thread_ns(config)
        thread_key = self._thread_key(config)
        parent_id = get_checkpoint_id(config)
        checkpoint_id = str(checkpoint.get("id") or "")

        payload_state = {
            self._STATE_KEY: dict(checkpoint),
        }
        payload_metadata: Dict[str, Any] = {
            self._META_KEY: dict(metadata or {}),
            self._LG_ID_KEY: checkpoint_id,
        }
        if parent_id:
            payload_metadata[self._PARENT_ID_KEY] = str(parent_id)

        self._saver.checkpoint_write(
            thread_key=thread_key,
            state=payload_state,
            metadata=payload_metadata,
        )
        return self._checkpoint_config(thread_id=thread_id, checkpoint_ns=checkpoint_ns, checkpoint_id=checkpoint_id)

    def put_writes(
        self,
        config: RunnableConfig,
        writes: Sequence[Tuple[str, Any]],
        task_id: str,
        task_path: str = "",
    ) -> None:
        _ = config, writes, task_id, task_path
        # Saver does not expose pending-write persistence.
        return None

    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
        return self.get_tuple(config)

    async def alist(
        self,
        config: Optional[RunnableConfig],
        *,
        filter: Optional[Dict[str, Any]] = None,
        before: Optional[RunnableConfig] = None,
        limit: Optional[int] = None,
    ) -> AsyncIterator[CheckpointTuple]:
        for item in self.list(config, filter=filter, before=before, limit=limit):
            yield item

    async def aput(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        return self.put(config, checkpoint, metadata, new_versions)

    async def aput_writes(
        self,
        config: RunnableConfig,
        writes: Sequence[Tuple[str, Any]],
        task_id: str,
        task_path: str = "",
    ) -> None:
        self.put_writes(config, writes, task_id, task_path)

