# -*- coding: UTF-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.

"""
Persistence checkpointer implementation using BaseKVStore interface.

This module provides a checkpointer implementation that uses BaseKVStore
for persistent storage, supporting any KV store implementation (shelve, database, etc.).
"""

import base64
from abc import ABC, abstractmethod
from pathlib import Path
from typing import (
    Any,
    Optional,
    Tuple,
)

from sqlalchemy import event
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    create_async_engine,
)

from openjiuwen.core.common.constants.constant import INTERACTIVE_INPUT
from openjiuwen.core.common.exception.codes import StatusCode
from openjiuwen.core.common.exception.errors import build_error
from openjiuwen.core.common.logging import session_logger, LogEventType
from openjiuwen.core.foundation.store import (
    BaseKVStore,
    DbBasedKVStore,
)
from openjiuwen.core.foundation.store.kv import ShelveStore
from openjiuwen.core.graph.pregel import TASK_STATUS_INTERRUPT
from openjiuwen.core.graph.store import (
    create_serializer,
    GraphState,
    Serializer,
    Store,
)
from openjiuwen.core.session import (
    BaseSession,
    Checkpointer,
    InteractiveInput,
    NodeSession,
)
from openjiuwen.core.session.checkpointer import (
    build_key,
    build_key_with_namespace,
    CheckpointerFactory,
    CheckpointerProvider,
    SESSION_NAMESPACE_AGENT,
    SESSION_NAMESPACE_AGENT_TEAM,
    SESSION_NAMESPACE_WORKFLOW,
    Storage,
    WORKFLOW_NAMESPACE_GRAPH,
)
from openjiuwen.core.session.constants import FORCE_DEL_WORKFLOW_STATE_KEY


class BaseStorage(Storage, ABC):
    """
    Base class for persistence-based storage implementations with common functionality.
    
    This class uses BaseKVStore interface and does not depend on specific implementations.
    """

    def __init__(self, kv_store: BaseKVStore):
        """
        Initialize BaseStorage with a BaseKVStore instance.

        Args:
            kv_store (BaseKVStore): The BaseKVStore instance for all storage operations.
            ttl (Optional[dict[str, Any]]): Optional TTL configuration.
        """
        self._kv_store = kv_store
        self._serde: Serializer = create_serializer("pickle")

    def _serialize_state(self, state: Any) -> Optional[Tuple[str, bytes]]:
        """Serialize state and return (dump_type, blob) tuple."""
        return self._serde.dumps_typed(state)

    def _decode_dump_type(self, dump_type: Any) -> str:
        """Decode dump_type from bytes to string if needed."""
        if isinstance(dump_type, bytes):
            return dump_type.decode("utf-8")
        return dump_type if dump_type is not None else ""

    def _deserialize_state(self, dump_type: Any, blob: Any) -> Any:
        """Deserialize state from (dump_type, blob) tuple."""
        if dump_type is None or blob is None:
            return None
        # Decode dump_type if needed
        dump_type_str = self._decode_dump_type(dump_type)
        try:
            # Convert blob string back to bytes if needed
            if isinstance(blob, str):
                # Assume base64 encoding for bytes stored as string
                blob_bytes = base64.b64decode(blob)
            else:
                blob_bytes = blob
            return self._serde.loads_typed((dump_type_str, blob_bytes))
        except Exception as e:
            session_logger.error(
                "Failed to deserialize state",
                event_type=LogEventType.CHECKPOINT_ERROR,
                error_message=str(e),
                metadata={"operation": "deserialize"}
            )
            return None


class BaseSingleStateStorage(BaseStorage, ABC):
    _KEY_NUMS = 2

    @property
    @abstractmethod
    def _namespace(self) -> str:
        ...

    @property
    @abstractmethod
    def _entity_label(self) -> str:
        ...

    @property
    @abstractmethod
    def _state_blobs_key(self) -> str:
        ...

    @property
    @abstractmethod
    def _state_dump_type_key(self) -> str:
        ...

    @abstractmethod
    def _get_entity_id(self, session: BaseSession) -> str:
        ...

    @abstractmethod
    def _get_state_to_save(self, session: BaseSession) -> Any:
        ...

    @abstractmethod
    def _restore_state(self, session: BaseSession, state: Any) -> None:
        ...

    def _build_state_keys(self, session_id: str, entity_id: str) -> tuple[str, str]:
        dump_type_key = build_key_with_namespace(
            session_id, self._namespace, entity_id, self._state_dump_type_key
        )
        blob_key = build_key_with_namespace(
            session_id, self._namespace, entity_id, self._state_blobs_key
        )
        return dump_type_key, blob_key

    def _entity_log_extra(self, entity_id: str) -> dict[str, str]:
        if self._entity_label == "agent":
            return {"agent_id": entity_id}
        return {"workflow_id": entity_id}

    def _log_kwargs(self, session_id: str, entity_id: str, operation: str) -> dict[str, Any]:
        log_kwargs = {
            "event_type": LogEventType.CHECKPOINT_ERROR,
            "session_id": session_id,
            "metadata": {"operation": operation, "storage_type": "persistence"},
        }
        log_kwargs.update(self._entity_log_extra(entity_id))
        return log_kwargs

    async def save(self, session: BaseSession):
        """Save state to KV store."""
        state = self._get_state_to_save(session)
        session_id = session.session_id()
        entity_id = self._get_entity_id(session)

        state_blob = self._serialize_state(state)
        if not state_blob:
            session_logger.warning(
                f"Failed to serialize {self._entity_label} state",
                **self._log_kwargs(session_id, entity_id, "serialize"),
            )
            return

        try:
            dump_type, blob = state_blob
            pipeline = self._kv_store.pipeline()
            dump_type_key, blob_key = self._build_state_keys(session_id, entity_id)
            await pipeline.set(dump_type_key, dump_type)
            await pipeline.set(blob_key, blob)
            await pipeline.execute()
            session_logger.debug(
                f"{self._entity_label.capitalize()} state saved successfully",
                event_type=LogEventType.CHECKPOINT_SAVE,
                session_id=session_id,
                metadata={"storage_type": "persistence"},
                **self._entity_log_extra(entity_id),
            )
        except Exception as e:
            session_logger.error(
                f"Failed to save {self._entity_label} state",
                error_message=str(e),
                **self._log_kwargs(session_id, entity_id, "save"),
            )
            raise

    async def recover(self, session: BaseSession, inputs: InteractiveInput = None):
        """Recover state from KV store."""
        session_id = session.session_id()
        entity_id = self._get_entity_id(session)

        pipeline = self._kv_store.pipeline()
        dump_type_key, blob_key = self._build_state_keys(session_id, entity_id)
        await pipeline.get(dump_type_key)
        await pipeline.get(blob_key)
        results = await pipeline.execute()

        if len(results) != self._KEY_NUMS:
            session_logger.debug(
                f"Unexpected key count during {self._entity_label} state recovery",
                event_type=LogEventType.CHECKPOINT_RESTORE,
                session_id=session_id,
                metadata={"expected_keys": self._KEY_NUMS, "actual_keys": len(results)},
                **self._entity_log_extra(entity_id),
            )
            return

        dump_type, blob = results[0], results[1]
        state = self._deserialize_state(dump_type, blob)
        if state is None:
            session_logger.debug(
                f"No {self._entity_label} state found",
                event_type=LogEventType.CHECKPOINT_RESTORE,
                session_id=session_id,
                metadata={"storage_type": "persistence"},
                **self._entity_log_extra(entity_id),
            )
            return

        try:
            self._restore_state(session, state)
            session_logger.debug(
                f"{self._entity_label.capitalize()} state recovered successfully",
                event_type=LogEventType.CHECKPOINT_RESTORE,
                session_id=session_id,
                metadata={"storage_type": "persistence"},
                **self._entity_log_extra(entity_id),
            )
        except Exception as e:
            session_logger.error(
                f"Failed to set {self._entity_label} state",
                error_message=str(e),
                **self._log_kwargs(session_id, entity_id, "set_state"),
            )
            raise

    async def clear(self, entity_id: str, session_id: str):
        """Clear state from KV store."""
        dump_type_key, blob_key = self._build_state_keys(session_id, entity_id)
        deleted = await self._kv_store.batch_delete([dump_type_key, blob_key])
        session_logger.debug(
            f"{self._entity_label.capitalize()} checkpoint cleared",
            event_type=LogEventType.CHECKPOINT_CLEAR,
            session_id=session_id,
            **self._entity_log_extra(entity_id),
            metadata={"deleted_keys": deleted, "storage_type": "persistence"}
        )

    async def exists(self, session: BaseSession) -> bool:
        """Check if state exists in KV store."""
        session_id = session.session_id()
        entity_id = self._get_entity_id(session)

        pipeline = self._kv_store.pipeline()
        dump_type_key, blob_key = self._build_state_keys(session_id, entity_id)
        await pipeline.exists(dump_type_key)
        await pipeline.exists(blob_key)
        results = await pipeline.execute()

        if len(results) != self._KEY_NUMS:
            return False

        # Both keys must exist for the state to be considered existing
        return results[0] is True and results[1] is True


class AgentStorage(BaseSingleStateStorage):
    """Agent state storage using BaseKVStore."""

    _namespace = SESSION_NAMESPACE_AGENT
    _entity_label = "agent"
    _state_blobs_key = "agent_state_blobs"
    _state_dump_type_key = "agent_state_blobs_dump_type"

    def _get_entity_id(self, session: BaseSession) -> str:
        return session.agent_id()

    def _get_state_to_save(self, session: BaseSession) -> Any:
        return session.state().get_state()

    def _restore_state(self, session: BaseSession, state: Any) -> None:
        session.state().set_state(state)


class AgentTeamStorage(BaseSingleStateStorage):
    """Agent team global state storage using BaseKVStore."""

    _namespace = SESSION_NAMESPACE_AGENT_TEAM
    _entity_label = "agent_team"
    _state_blobs_key = "agent_team_state_blobs"
    _state_dump_type_key = "agent_team_state_blobs_dump_type"

    def _get_entity_id(self, session: BaseSession) -> str:
        return session.team_id()

    def _get_state_to_save(self, session: BaseSession) -> Any:
        return session.state().get_global(None)

    def _restore_state(self, session: BaseSession, state: Any) -> None:
        session.state().global_state.set_state(state)


class WorkflowStorage(BaseStorage):
    """Workflow state storage using BaseKVStore."""

    _STATE_BLOBS = "workflow_state_blobs"
    _STATE_BLOBS_DUMP_TYPE = "workflow_state_blobs_dump_type"
    _UPDATE_BLOBS = "workflow_update_blobs"
    _UPDATE_BLOBS_DUMP_TYPE = "workflow_update_blobs_dump_type"
    _KEY_NUMS = 4

    def _process_interactive_inputs(self, session: BaseSession, inputs: InteractiveInput) -> None:
        """Process interactive inputs and update workflow state."""
        if inputs.raw_inputs is not None:
            session.state().update_and_commit_workflow_state({INTERACTIVE_INPUT: inputs.raw_inputs})
            return

        if not (hasattr(inputs, 'user_inputs') and inputs.user_inputs):
            return

        for node_id, value in inputs.user_inputs.items():
            node_session = NodeSession(session, node_id)
            interactive_input = node_session.state().get(INTERACTIVE_INPUT)
            if isinstance(interactive_input, list):
                interactive_input.append(value)
                node_session.state().update({INTERACTIVE_INPUT: interactive_input})
            else:
                node_session.state().update({INTERACTIVE_INPUT: [value]})
        session.state().commit()

    async def save(self, session: BaseSession):
        """Save workflow state to KV store."""
        state = session.state().get_state()
        workflow_id = session.workflow_id()
        session_id = session.session_id()

        pipeline = self._kv_store.pipeline()
        has_operations = False

        state_blob = self._serialize_state(state)
        if state_blob:
            dump_type, blob = state_blob
            dump_type_key = build_key_with_namespace(
                session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._STATE_BLOBS_DUMP_TYPE
            )
            blob_key = build_key_with_namespace(
                session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._STATE_BLOBS
            )
            await pipeline.set(dump_type_key, dump_type)
            await pipeline.set(blob_key, blob)
            has_operations = True
        else:
            session_logger.warning(
                "Failed to serialize workflow state",
                event_type=LogEventType.CHECKPOINT_ERROR,
                session_id=session_id,
                workflow_id=workflow_id,
                metadata={"operation": "serialize"}
            )

        updates = session.state().get_updates()
        updates_blob = self._serialize_state(updates)
        if updates_blob:
            dump_type, blob = updates_blob
            dump_type_key = build_key_with_namespace(
                session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._UPDATE_BLOBS_DUMP_TYPE
            )
            blob_key = build_key_with_namespace(
                session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._UPDATE_BLOBS
            )
            await pipeline.set(dump_type_key, dump_type)
            await pipeline.set(blob_key, blob)
            has_operations = True

        if has_operations:
            try:
                await pipeline.execute()
                session_logger.debug(
                    "Workflow state saved successfully",
                    event_type=LogEventType.CHECKPOINT_SAVE,
                    session_id=session_id,
                    workflow_id=workflow_id,
                    metadata={"storage_type": "persistence"}
                )
            except Exception as e:
                session_logger.error(
                    "Failed to save workflow state",
                    event_type=LogEventType.CHECKPOINT_ERROR,
                    session_id=session_id,
                    workflow_id=workflow_id,
                    error_message=str(e),
                    metadata={"operation": "save", "storage_type": "persistence"}
                )
                raise

    async def recover(self, session: BaseSession, inputs: InteractiveInput = None):
        """Recover workflow state from KV store."""
        workflow_id = session.workflow_id()
        session_id = session.session_id()

        pipeline = self._kv_store.pipeline()
        state_dump_type_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._STATE_BLOBS_DUMP_TYPE
        )
        state_blob_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._STATE_BLOBS
        )
        await pipeline.get(state_dump_type_key)
        await pipeline.get(state_blob_key)
        updates_dump_type_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._UPDATE_BLOBS_DUMP_TYPE
        )
        updates_blob_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._UPDATE_BLOBS
        )
        await pipeline.get(updates_dump_type_key)
        await pipeline.get(updates_blob_key)
        results = await pipeline.execute()

        if len(results) != self._KEY_NUMS:
            session_logger.warning(
                "Unexpected key count during workflow state recovery",
                event_type=LogEventType.CHECKPOINT_RESTORE,
                session_id=session_id,
                workflow_id=workflow_id,
                metadata={"expected_keys": self._KEY_NUMS, "actual_keys": len(results)}
            )
            return

        # Recover state
        state_dump_type, state_blob = results[0], results[1]
        state_dump_type_str = self._decode_dump_type(state_dump_type)

        if state_blob and state_dump_type_str and state_dump_type_str != "empty":
            try:
                state = self._deserialize_state(state_dump_type_str, state_blob)
                if state is not None:
                    session.state().set_state(state)
            except Exception as e:
                session_logger.error(
                    "Failed to deserialize workflow state",
                    event_type=LogEventType.CHECKPOINT_ERROR,
                    session_id=session_id,
                    workflow_id=workflow_id,
                    error_message=str(e),
                    metadata={"operation": "deserialize_state"}
                )

        # Process interactive inputs
        if inputs is not None:
            self._process_interactive_inputs(session, inputs)

        # Recover updates
        updates_dump_type, updates_blob = results[2], results[3]
        updates_dump_type_str = self._decode_dump_type(updates_dump_type)

        if updates_blob and updates_dump_type_str and updates_dump_type_str != "empty":
            try:
                state_updates = self._deserialize_state(updates_dump_type_str, updates_blob)
                if state_updates is not None:
                    session.state().set_updates(state_updates)
            except Exception as e:
                session_logger.error(
                    "Failed to deserialize workflow updates",
                    event_type=LogEventType.CHECKPOINT_ERROR,
                    session_id=session_id,
                    workflow_id=workflow_id,
                    error_message=str(e),
                    metadata={"operation": "deserialize_updates"}
                )

    async def clear(self, workflow_id: str, session_id: str):
        """Clear workflow state from KV store."""
        state_dump_type_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._STATE_BLOBS_DUMP_TYPE
        )
        state_blob_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._STATE_BLOBS
        )
        state_updates_dump_type_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._UPDATE_BLOBS_DUMP_TYPE
        )
        state_updates_blob_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._UPDATE_BLOBS
        )
        # Use batch_delete for multiple keys
        deleted = await self._kv_store.batch_delete([
            state_dump_type_key, state_blob_key,
            state_updates_dump_type_key, state_updates_blob_key
        ])
        session_logger.debug(
            "Workflow checkpoint cleared",
            event_type=LogEventType.CHECKPOINT_CLEAR,
            session_id=session_id,
            workflow_id=workflow_id,
            metadata={"deleted_keys": deleted, "storage_type": "persistence"}
        )

    async def exists(self, session: BaseSession) -> bool:
        """Check if workflow state exists in KV store."""
        workflow_id = session.workflow_id()
        session_id = session.session_id()

        pipeline = self._kv_store.pipeline()
        # Check state keys
        state_dump_type_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._STATE_BLOBS_DUMP_TYPE
        )
        state_blob_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._STATE_BLOBS
        )
        await pipeline.exists(state_dump_type_key)
        await pipeline.exists(state_blob_key)
        # Check updates keys (optional)
        state_updates_dump_type_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._UPDATE_BLOBS_DUMP_TYPE
        )
        state_updates_blob_key = build_key_with_namespace(
            session_id, SESSION_NAMESPACE_WORKFLOW, workflow_id, self._UPDATE_BLOBS
        )
        await pipeline.exists(state_updates_dump_type_key)
        await pipeline.exists(state_updates_blob_key)
        results = await pipeline.execute()

        if len(results) != self._KEY_NUMS:
            return False

        # At least state keys must exist for the workflow state to be considered existing
        # Updates are optional, so we only require state keys to exist
        return results[0] is True and results[1] is True


class GraphStore(Store):
    """
    Graph state store implementation using BaseKVStore.
    
    Graph state keys are structured as: session:workflow-graph:workflow_id:suffix
    This separates graph state from workflow's own state.
    """

    _DATA_TYPE = "checkpoint_data_type"
    _DATA_VALUE = "checkpoint_data_value"
    _KEY_NUMS = 2

    def __init__(self, kv_store: BaseKVStore) -> None:
        """
        Initialize GraphStore with a BaseKVStore instance.

        Args:
            kv_store (BaseKVStore): The BaseKVStore instance for all storage operations.
        """
        self._kv_store = kv_store
        self._serde: Serializer = create_serializer("pickle")

    def _serialize_graph_state(self, graph_state: GraphState) -> Optional[Tuple[str, bytes]]:
        """Serialize graph state and return (dump_type, blob) tuple."""
        return self._serde.dumps_typed(graph_state)

    def _decode_dump_type(self, dump_type: Any) -> str:
        """Decode dump_type from bytes to string if needed."""
        if isinstance(dump_type, bytes):
            return dump_type.decode("utf-8")
        return dump_type if dump_type is not None else ""

    def _deserialize_graph_state(self, dump_type: Any, blob: Any) -> Optional[GraphState]:
        """Deserialize graph state from (dump_type, blob) tuple."""
        if not dump_type or blob is None:
            return None
        dump_type_str = self._decode_dump_type(dump_type)
        try:
            # Convert blob string back to bytes if needed
            if isinstance(blob, str):
                blob_bytes = base64.b64decode(blob)
            else:
                blob_bytes = blob
            return self._serde.loads_typed((dump_type_str, blob_bytes))
        except Exception as e:
            session_logger.error(
                "Failed to deserialize graph state",
                event_type=LogEventType.CHECKPOINT_ERROR,
                error_message=str(e),
                metadata={"operation": "deserialize"}
            )
            return None

    async def get(self, session_id: str, ns: str) -> Optional[GraphState]:
        """Get graph state from KV store."""
        pipeline = self._kv_store.pipeline()
        key_type = build_key_with_namespace(
            session_id, WORKFLOW_NAMESPACE_GRAPH, ns, self._DATA_TYPE
        )
        await pipeline.get(key_type)
        key_value = build_key_with_namespace(
            session_id, WORKFLOW_NAMESPACE_GRAPH, ns, self._DATA_VALUE
        )
        await pipeline.get(key_value)
        results = await pipeline.execute()

        if len(results) != self._KEY_NUMS:
            session_logger.error(
                "Unexpected key count during graph state retrieval",
                event_type=LogEventType.CHECKPOINT_ERROR,
                session_id=session_id,
                metadata={"expected_keys": self._KEY_NUMS, "actual_keys": len(results), "namespace": ns}
            )
            return None

        _type, _value = results[0], results[1]
        if not _type or not _value:
            session_logger.debug(
                "Graph state not found in KV store",
                event_type=LogEventType.CHECKPOINT_RESTORE,
                session_id=session_id,
                metadata={"namespace": ns, "has_type": bool(_type), "has_value": bool(_value)}
            )
            return None

        graph_state = self._deserialize_graph_state(_type, _value)
        if graph_state is None:
            session_logger.debug(
                "Failed to deserialize graph state",
                event_type=LogEventType.CHECKPOINT_ERROR,
                session_id=session_id,
                metadata={"namespace": ns}
            )
            return None
        return graph_state

    async def save(self, session_id: str, ns: str, state: GraphState) -> None:
        """Save graph state to KV store."""
        serialized = self._serialize_graph_state(state)
        if not serialized:
            session_logger.warning(
                "Failed to serialize graph state",
                event_type=LogEventType.CHECKPOINT_ERROR,
                session_id=session_id,
                metadata={"namespace": ns, "operation": "serialize"}
            )
            return

        dump_type, blob = serialized
        dump_type_str = dump_type if isinstance(dump_type, str) else dump_type.decode("utf-8")

        try:
            key_type = build_key_with_namespace(
                session_id, WORKFLOW_NAMESPACE_GRAPH, ns, self._DATA_TYPE
            )
            pipeline = self._kv_store.pipeline()
            await pipeline.set(key_type, dump_type_str)
            key_value = build_key_with_namespace(
                session_id, WORKFLOW_NAMESPACE_GRAPH, ns, self._DATA_VALUE
            )
            await pipeline.set(key_value, blob)
            await pipeline.execute()
            session_logger.debug(
                "Graph state saved successfully",
                event_type=LogEventType.CHECKPOINT_SAVE,
                session_id=session_id,
                metadata={"namespace": ns, "storage_type": "graph"}
            )
        except Exception as e:
            session_logger.error(
                "Failed to save graph state",
                event_type=LogEventType.CHECKPOINT_ERROR,
                session_id=session_id,
                error_message=str(e),
                metadata={"namespace": ns, "operation": "save", "storage_type": "graph"}
            )
            raise

    async def delete(self, session_id: str, ns: Optional[str] = None) -> None:
        """
        Delete graph state keys for the given session_id and namespace.
        
        Args:
            session_id: Session identifier.
            ns: Namespace identifier. If None or empty, deletes all graph state data
                for the session_id (all namespaces under this session).
        """
        if not ns:
            # Delete all graph state data for this session_id
            prefix = build_key(session_id, WORKFLOW_NAMESPACE_GRAPH)
            await self._kv_store.delete_by_prefix(prefix)
            session_logger.debug(
                "Graph checkpoint cleared for all namespaces",
                event_type=LogEventType.CHECKPOINT_CLEAR,
                session_id=session_id,
                metadata={"storage_type": "graph"}
            )
        else:
            # Delete specific namespace
            prefix = build_key_with_namespace(
                session_id, WORKFLOW_NAMESPACE_GRAPH, ns
            )
            await self._kv_store.delete_by_prefix(prefix)
            session_logger.debug(
                "Graph checkpoint cleared",
                event_type=LogEventType.CHECKPOINT_CLEAR,
                session_id=session_id,
                metadata={"namespace": ns, "storage_type": "graph"}
            )


class PersistenceCheckpointer(Checkpointer):
    """
    Persistence-based checkpointer implementation using BaseKVStore.
    
    This checkpointer uses BaseKVStore interface for persistent storage,
    supporting any KV store implementation (shelve, database, etc.).
    """

    def __init__(self, kv_store: BaseKVStore):
        """
        Initialize PersistenceCheckpointer with a BaseKVStore instance.

        Args:
            kv_store (BaseKVStore): The BaseKVStore instance for all storage operations.
        """
        self._kv_store = kv_store
        self._agent_storage = AgentStorage(kv_store)
        self._agent_team_storage = AgentTeamStorage(kv_store)
        self._workflow_storage = WorkflowStorage(kv_store)
        self._graph_state = GraphStore(kv_store)

    async def pre_agent_execute(self, session: BaseSession, inputs):
        """Prepare agent execution by recovering agent state."""
        session_logger.info(
            "Agent checkpoint restore initiated",
            event_type=LogEventType.CHECKPOINT_RESTORE,
            session_id=session.session_id(),
            agent_id=session.agent_id(),
            metadata={"operation": "pre_execute", "storage_type": "persistence"}
        )
        await self._agent_storage.recover(session)
        if inputs is not None:
            session.state().update({INTERACTIVE_INPUT: [inputs]})

    async def pre_agent_team_execute(self, session: BaseSession, inputs):
        session_logger.info(
            "Agent team checkpoint restore initiated",
            event_type=LogEventType.CHECKPOINT_RESTORE,
            session_id=session.session_id(),
            workflow_id=session.team_id(),
            metadata={"operation": "pre_execute", "storage_type": "persistence"}
        )
        await self._agent_team_storage.recover(session)
        if inputs is not None:
            session.state().update_global({INTERACTIVE_INPUT: [inputs]})

    async def interrupt_agent_execute(self, session: BaseSession):
        """Save agent state when interaction is required."""
        session_logger.info(
            "Agent checkpoint save on interrupt",
            event_type=LogEventType.CHECKPOINT_SAVE,
            session_id=session.session_id(),
            agent_id=session.agent_id(),
            metadata={"reason": "interaction_required", "storage_type": "persistence"}
        )
        await self._agent_storage.save(session)

    async def post_agent_execute(self, session: BaseSession):
        """Save agent state after execution completes."""
        session_logger.info(
            "Agent checkpoint save on completion",
            event_type=LogEventType.CHECKPOINT_SAVE,
            session_id=session.session_id(),
            agent_id=session.agent_id(),
            metadata={"reason": "agent_finished", "storage_type": "persistence"}
        )
        await self._agent_storage.save(session)

    async def post_agent_team_execute(self, session: BaseSession):
        session_logger.info(
            "Agent team checkpoint save on completion",
            event_type=LogEventType.CHECKPOINT_SAVE,
            session_id=session.session_id(),
            workflow_id=session.team_id(),
            metadata={"reason": "team_finished", "storage_type": "persistence"}
        )
        await self._agent_team_storage.save(session)

    async def pre_workflow_execute(self, session: BaseSession, inputs: InteractiveInput):
        """
        Prepare workflow execution by recovering or clearing workflow state.
        
        If inputs is an InteractiveInput, recover the workflow state.
        If inputs is not an InteractiveInput and workflow state exists:
            - If FORCE_DEL_WORKFLOW_STATE_KEY is True, delete graph state and workflow state
            - Otherwise, raise WORKFLOW_STATE_INVALID exception
        
        Args:
            session (BaseSession): The session for the workflow.
            inputs (InteractiveInput): The input for the workflow execution.
        """
        workflow_id = session.workflow_id()
        session_logger.info(
            "Workflow checkpoint restore initiated",
            event_type=LogEventType.CHECKPOINT_RESTORE,
            session_id=session.session_id(),
            workflow_id=workflow_id,
            metadata={"operation": "pre_execute", "storage_type": "persistence"}
        )
        if isinstance(inputs, InteractiveInput):
            await self._workflow_storage.recover(session, inputs)
        else:
            # Check if workflow state exists
            if not await self._workflow_storage.exists(session):
                return

            # If FORCE_DEL_WORKFLOW_STATE_KEY is enabled, delete the state
            if session.config().get_env(FORCE_DEL_WORKFLOW_STATE_KEY, False):
                workflow_id = session.workflow_id()
                if workflow_id is None:
                    session_logger.warning(
                        "Workflow ID is None during state cleanup",
                        event_type=LogEventType.CHECKPOINT_ERROR,
                        session_id=session.session_id(),
                        metadata={"operation": "force_delete"}
                    )
                    return
                session_id = session.session_id()
                await self._graph_state.delete(session_id, workflow_id)
                await self._workflow_storage.clear(workflow_id, session_id)
                session_logger.info(
                    "Workflow state force deleted",
                    event_type=LogEventType.CHECKPOINT_CLEAR,
                    session_id=session_id,
                    workflow_id=workflow_id,
                    metadata={"reason": "force_delete", "storage_type": "persistence"}
                )
            else:
                # Raise exception if state exists but cleanup is disabled
                raise build_error(
                    StatusCode.CHECKPOINTER_PRE_WORKFLOW_EXECUTION_ERROR,
                    workflow=workflow_id,
                    reason="workflow state exists but non-interactive input and cleanup is disabled"
                )

    async def post_workflow_execute(self, session: BaseSession, result, exception):
        """Handle workflow execution completion."""
        workflow_id = session.workflow_id()
        session_id = session.session_id()

        if exception is not None:
            session_logger.info(
                "Workflow checkpoint save on exception",
                event_type=LogEventType.CHECKPOINT_SAVE,
                session_id=session_id,
                workflow_id=workflow_id,
                metadata={"reason": "exception", "storage_type": "persistence"}
            )
            await self._workflow_storage.save(session)
            raise exception

        if result.get(TASK_STATUS_INTERRUPT) is None:
            session_logger.info(
                "Workflow checkpoint cleared on completion",
                event_type=LogEventType.CHECKPOINT_CLEAR,
                session_id=session_id,
                workflow_id=workflow_id,
                metadata={"reason": "workflow_completed", "storage_type": "persistence"}
            )
            await self._graph_state.delete(session_id, workflow_id)
            await self._workflow_storage.clear(workflow_id, session_id)
        else:
            session_logger.info(
                "Workflow checkpoint save on interrupt",
                event_type=LogEventType.CHECKPOINT_SAVE,
                session_id=session_id,
                workflow_id=workflow_id,
                metadata={"reason": "interaction_required", "storage_type": "persistence"}
            )
            await self._workflow_storage.save(session)

    async def session_exists(self, session_id: str) -> bool:
        """
        Check if a session exists in KV store.
        
        Args:
            session_id (str): The session ID to check.
        
        Returns:
            bool: True if the session exists (has associated keys), False otherwise.
        """
        if self._kv_store is None:
            return False

        # Check if any keys exist with the session_id prefix
        prefix = f"{session_id}:"
        keys = await self._kv_store.get_by_prefix(prefix)
        return len(keys) > 0

    async def release(self, session_id: str, agent_id: Optional[str] = None):
        """
        Release resources for a session, optionally for a specific agent.
        
        Args:
            session_id (str): The session ID to release resources for.
            agent_id (str, optional): If provided, only release resources for this specific agent.
        """
        if self._kv_store is None:
            session_logger.warning(
                "Cannot release resources: KV store is None",
                event_type=LogEventType.CHECKPOINT_ERROR,
                session_id=session_id,
                metadata={"operation": "release"}
            )
            return

        if agent_id is not None:
            session_logger.info(
                "Agent checkpoint cleared",
                event_type=LogEventType.CHECKPOINT_CLEAR,
                session_id=session_id,
                agent_id=agent_id,
                metadata={"operation": "release", "storage_type": "persistence"}
            )
            await self._agent_storage.clear(agent_id, session_id)
        else:
            session_logger.info(
                "Session cleared",
                event_type=LogEventType.CHECKPOINT_CLEAR,
                session_id=session_id,
                metadata={"operation": "release_all", "storage_type": "persistence"}
            )
            # Delete all keys matching the session prefix
            prefix = f"{session_id}:"
            await self._kv_store.delete_by_prefix(prefix)
            session_logger.debug(
                "All session resources released",
                event_type=LogEventType.CHECKPOINT_CLEAR,
                session_id=session_id,
                metadata={"storage_type": "persistence"}
            )

    def graph_store(self) -> Store:
        """Return the graph store instance."""
        return self._graph_state


def _enable_sqlite_wal(engine: AsyncEngine) -> None:
    """
    Enable WAL mode on SQLite to reduce 'database is locked' under concurrency.

    WAL allows one writer and concurrent readers. Listens for sync_engine connect
    and runs PRAGMA journal_mode=WAL on each new connection.
    """
    @event.listens_for(engine.sync_engine, "connect")
    def _set_sqlite_pragma(dbapi_conn, connection_record):
        cursor = dbapi_conn.cursor()
        cursor.execute("PRAGMA journal_mode=WAL")
        cursor.close()


@CheckpointerFactory.register("persistence")
class PersistenceCheckpointerProvider(CheckpointerProvider):
    """
    Provider for creating persistence-based checkpointers.
    
    Configuration format:
        {
            "db_type": "sqlite" | "shelve",  # Optional: Type of storage backend (default: "sqlite")
            "db_path": "checkpointer.db",     # Optional: Path to database file
            "db_client": AsyncEngine,         # Optional: Pre-configured database engine
            "db_timeout": 30,                 # Optional: SQLite lock wait seconds (default: 30)
            "db_enable_wal": True,            # Optional: Enable SQLite WAL mode (default: True)
        }
    """

    async def create(self, conf: dict) -> Checkpointer:
        """
        Create a PersistenceCheckpointer instance.
        
        Args:
            conf (dict): Configuration dictionary with optional keys:
                - 'db_type': Type of storage backend ("sqlite" or "shelve", default: "sqlite")
                - 'db_path': Path to database file (default: "checkpointer")
                - 'db_client': Pre-configured AsyncEngine instance (optional)
                - 'db_timeout': Seconds to wait when SQLite DB is locked (default: 30)
                - 'db_enable_wal': Enable SQLite WAL mode for better concurrency (default: True)
        
        Returns:
            Checkpointer: A PersistenceCheckpointer instance.
        """
        db_type = conf.get("db_type", "sqlite")
        db_path = conf.get("db_path", "checkpointer")

        if db_type == "sqlite":
            db_client = conf.get("db_client")
            if db_client is not None and isinstance(db_client, AsyncEngine):
                kv_store = DbBasedKVStore(db_client)
            else:
                if not db_path.endswith(".db"):
                    db_path = f"{db_path}.db"
                # Ensure parent directory exists; SQLite cannot create it.
                if db_path and not db_path.strip().startswith(":memory:"):
                    Path(db_path).parent.mkdir(parents=True, exist_ok=True)
                # Timeout (seconds) to wait when DB is locked; reduces OperationalError.
                db_timeout = conf.get("db_timeout", 30)
                engine = create_async_engine(
                    f"sqlite+aiosqlite:///{db_path}",
                    connect_args={"timeout": db_timeout},
                )
                # Enable WAL mode if configured (default: True)
                if conf.get("db_enable_wal", True):
                    _enable_sqlite_wal(engine)
                kv_store = DbBasedKVStore(engine)
        elif db_type == "shelve":
            if db_path.endswith(".db"):
                db_path = db_path.removesuffix(".db")
            kv_store = ShelveStore(db_path)
        else:
            raise build_error(StatusCode.CHECKPOINTER_CONFIG_ERROR, reason=f"db type[{db_type}] is not supported")
        return PersistenceCheckpointer(kv_store)
