
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, ChannelVersions, BaseCheckpointSaver, \
    CheckpointTuple
from datetime import datetime, timedelta
from typing import Optional, Dict, Sequence, Tuple, Any
from collections.abc import Iterator
import threading

from langgraph.checkpoint.memory import MemorySaver


class MemorySaverWithExpiry(MemorySaver):
    """
    支持自动过期的内存检查点
    """

    def __init__(self, expire_time: int = 3600, clean_interval: int = 300):
        """
        自动清理过期 检查点
        """
        super().__init__()
        self._checkpoints: Dict[str, Checkpoint] = {}
        self._thread_ts: Dict[str, datetime] = {}
        self._metadata: Dict[str, CheckpointMetadata] = {}
        self.expire_time = timedelta(seconds=expire_time)
        self.clean_interval = clean_interval
        self._lock = threading.RLock()

        # 后台清除线程
        self.clean_thread = threading.Thread(target=self._background_clean, daemon=True)
        self.clean_thread.start()

    @staticmethod
    def get_thread_id(config: RunnableConfig) -> str:
        """安全获取线程ID"""
        try:
            return str(config["configurable"]["thread_id"])
        except KeyError:
            raise ValueError("RunnableConfig must contain 'configurable.thread_id'")

    def _background_clean(self):
        """
        后台清除线程
        """
        while True:
            # 等待
            threading.Event().wait(timeout=self.clean_interval)
            with self._lock:
                # 找出过期的 thread_id
                expired = [
                    thread_id
                    for thread_id, ts in self._thread_ts.items()
                        if (datetime.now() - ts) > self.expire_time
                ]
                for thread_id in expired:
                    self._delete(thread_id)

    def _is_expired(self, thread_id: str):
        """
        判断 thread_id 是否过期
        """
        ts = self._thread_ts.get(thread_id)
        return ts is None or (datetime.now() - ts) > self.expire_time

    def _delete(self, thread_id: str):
        """
        删除检查点
        """
        if thread_id in self._thread_ts:
            del self._thread_ts[thread_id]
        if thread_id in self._checkpoints:
            del self._checkpoints[thread_id]
        if thread_id in self._metadata:
            del self._metadata[thread_id]

    def put(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        """
        保存 检查点
        """
        thread_id = self.get_thread_id(config)
        with self._lock:
            self._thread_ts[thread_id] = datetime.now()
            self._checkpoints[thread_id] = checkpoint
            self._metadata[thread_id] = metadata

        return config

    def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
        """
        获取 检查点
        """
        thread_id = self.get_thread_id(config=config)
        with self._lock:
            if self._is_expired(thread_id):
                self._delete(thread_id)
                return None
        return self._checkpoints.get(thread_id)

    def list(
        self,
        config: Optional[RunnableConfig],
        *,
        filter: Optional[Dict[str, any]] = None,
        before: Optional[RunnableConfig] = None,
        limit: Optional[int] = None,
    ) -> Iterator[CheckpointTuple]:
        return super().list(config=config)

    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
        """获取检查点元组"""
        thread_id = self.get_thread_id(config)
        with self._lock:
            checkpoint = self._checkpoints.get(thread_id)
            if not checkpoint:
                return None

        # 从单独存储的元数据字典中获取 metadata
        metadata = self._metadata.get(thread_id, CheckpointMetadata())  # 添加此行
        ts = self._thread_ts.get(thread_id)
        if not ts:
            return None
        return CheckpointTuple(
            config=config,
            checkpoint=checkpoint,
            metadata=CheckpointMetadata(
                source="memory",
                timestamp=ts,
                step=metadata.get("step", 0)  # 添加 step 字段
            )
        )
