"""Checkpoint saver using existing PersistenceAdapter.

Implements LangGraph's checkpoint interface by delegating to
DataBridge's PersistenceAdapter (src/persistence/adapter.py).
This enables resume-from-failure for graph execution.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Dict, Optional, Sequence, Tuple

logger = logging.getLogger(__name__)

# Try to import LangGraph checkpoint base class
try:
    from langgraph.checkpoint.base import (
        BaseCheckpointSaver,
        Checkpoint,
        CheckpointMetadata,
        CheckpointTuple,
    )
    _HAS_CHECKPOINT_BASE = True
except ImportError:
    _HAS_CHECKPOINT_BASE = False

    # Stub base class for when langgraph is not installed
    class BaseCheckpointSaver:
        pass

    class CheckpointTuple:
        def __init__(self, config, checkpoint, metadata=None, parent_config=None):
            self.config = config
            self.checkpoint = checkpoint
            self.metadata = metadata
            self.parent_config = parent_config


class DataBridgeCheckpointSaver(BaseCheckpointSaver):
    """Checkpoint saver backed by DataBridge's PersistenceAdapter.

    Stores checkpoints as JSON files via FilePersistenceAdapter,
    enabling graph execution to resume from the last successful phase.
    """

    def __init__(self, base_dir: str = "data/langgraph_checkpoints"):
        if _HAS_CHECKPOINT_BASE:
            super().__init__()
        self._base_dir = base_dir

        from src.persistence.file_adapter import FilePersistenceAdapter
        self._adapter = FilePersistenceAdapter(base_dir=base_dir)

    def _thread_id(self, config: Dict[str, Any]) -> str:
        """Extract thread_id from config."""
        return config.get("configurable", {}).get("thread_id", "default")

    def _checkpoint_key(self, thread_id: str, checkpoint_id: str = "") -> str:
        """Build a storage key for a checkpoint."""
        if checkpoint_id:
            return f"ckpt_{thread_id}_{checkpoint_id}"
        return f"ckpt_{thread_id}_latest"

    # ------------------------------------------------------------------
    # BaseCheckpointSaver interface
    # ------------------------------------------------------------------

    def get_tuple(self, config: Dict[str, Any]) -> Optional[CheckpointTuple]:
        """Load the latest checkpoint for a thread."""
        thread_id = self._thread_id(config)
        checkpoint_id = config.get("configurable", {}).get("checkpoint_id", "")
        key = self._checkpoint_key(thread_id, checkpoint_id)

        data = self._adapter.load_run(key)
        if data is None:
            return None

        checkpoint = data.get("checkpoint", {})
        metadata = data.get("metadata", {})
        return CheckpointTuple(
            config=config,
            checkpoint=checkpoint,
            metadata=metadata,
        )

    def put(
        self,
        config: Dict[str, Any],
        checkpoint: Dict[str, Any],
        metadata: Dict[str, Any],
        new_versions: Optional[Dict] = None,
    ) -> Dict[str, Any]:
        """Save a checkpoint."""
        thread_id = self._thread_id(config)
        import uuid
        checkpoint_id = uuid.uuid4().hex[:8]

        key = self._checkpoint_key(thread_id, checkpoint_id)
        self._adapter.save_run(key, {
            "checkpoint": checkpoint,
            "metadata": metadata,
            "thread_id": thread_id,
            "checkpoint_id": checkpoint_id,
        })

        # Also save as 'latest'
        latest_key = self._checkpoint_key(thread_id)
        self._adapter.save_run(latest_key, {
            "checkpoint": checkpoint,
            "metadata": metadata,
            "thread_id": thread_id,
            "checkpoint_id": checkpoint_id,
        })

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

    def list(
        self,
        config: Optional[Dict[str, Any]] = None,
        *,
        filter: Optional[Dict[str, Any]] = None,
        before: Optional[Dict[str, Any]] = None,
        limit: Optional[int] = None,
    ):
        """List available checkpoints."""
        runs = self._adapter.list_runs(limit=limit or 50)
        results = []
        for run in runs:
            if run.get("thread_id"):
                results.append(CheckpointTuple(
                    config={"configurable": {
                        "thread_id": run["thread_id"],
                        "checkpoint_id": run.get("checkpoint_id", ""),
                    }},
                    checkpoint=run.get("checkpoint", {}),
                    metadata=run.get("metadata", {}),
                ))
        return results

    async def aget_tuple(self, config: Dict[str, Any]) -> Optional[CheckpointTuple]:
        """Async version of get_tuple."""
        return self.get_tuple(config)

    async def aput(
        self,
        config: Dict[str, Any],
        checkpoint: Dict[str, Any],
        metadata: Dict[str, Any],
        new_versions: Optional[Dict] = None,
    ) -> Dict[str, Any]:
        """Async version of put."""
        return self.put(config, checkpoint, metadata, new_versions)

    async def alist(
        self,
        config: Optional[Dict[str, Any]] = None,
        *,
        filter: Optional[Dict[str, Any]] = None,
        before: Optional[Dict[str, Any]] = None,
        limit: Optional[int] = None,
    ):
        """Async version of list."""
        for item in self.list(config, filter=filter, before=before, limit=limit):
            yield item

    # ------------------------------------------------------------------
    # Extra helpers
    # ------------------------------------------------------------------

    def clear(self, thread_id: str) -> bool:
        """Delete all checkpoints for a thread."""
        import os
        from pathlib import Path
        base = Path(self._base_dir)
        count = 0
        for f in base.glob(f"ckpt_{thread_id}_*.json"):
            f.unlink(missing_ok=True)
            count += 1
        return count > 0

    def get_latest_state(self, thread_id: str) -> Optional[Dict[str, Any]]:
        """Get the latest checkpoint state for a thread."""
        key = self._checkpoint_key(thread_id)
        data = self._adapter.load_run(key)
        if data:
            return data.get("checkpoint")
        return None
