"""
LockStock LangGraph Checkpointer
---------------------------------
Custom checkpointer that includes LockStock state hashes for audit.

Uses the thin client pattern - all verification happens server-side via /v1/stamp.

Extends LangGraph's checkpointing with cryptographic verification
that state hasn't been tampered with.
"""

from typing import Any, Dict, Optional, Tuple
from dataclasses import dataclass
import hashlib
import json
import time

from lockstock_core import LockStockClient


@dataclass
class LockStockCheckpoint:
    """A checkpoint with LockStock verification."""
    thread_id: str
    checkpoint_id: str
    state: Dict[str, Any]
    timestamp: float
    lockstock_hash: str
    sequence: int


class LockStockCheckpointer:
    """
    Checkpointer that integrates with LockStock for verified state management.

    This ensures that checkpointed states are cryptographically linked
    to the LockStock hash chain, providing tamper-evident state history.
    """

    def __init__(
        self,
        agent_id: str,
        api_key: str,
        endpoint: str = "https://lockstock-api-i9kp.onrender.com",
        storage: Optional[Dict[str, LockStockCheckpoint]] = None
    ):
        """
        Initialize LockStock checkpointer.

        Args:
            agent_id: The agent's unique identifier
            api_key: API key for authentication (from dashboard)
            endpoint: LockStock API endpoint
            storage: Optional backing storage (defaults to in-memory)
        """
        self.client = LockStockClient(
            agent_id=agent_id,
            api_key=api_key,
            endpoint=endpoint
        )
        self.agent_id = agent_id
        self._storage = storage or {}
        self._sequence = 0

    async def put(
        self,
        thread_id: str,
        state: Dict[str, Any],
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Save a checkpoint with LockStock verification.

        Args:
            thread_id: Thread identifier
            state: State to checkpoint
            metadata: Optional metadata

        Returns:
            Checkpoint ID
        """
        self._sequence += 1
        timestamp = time.time()

        # Create state hash
        state_json = json.dumps(state, sort_keys=True, default=str)
        state_hash = hashlib.sha256(state_json.encode()).hexdigest()

        # Stamp with LockStock to get chain hash (server handles audit logging)
        result = await self.client.stamp(task="CHECKPOINT")

        checkpoint_id = f"{thread_id}:{self._sequence}:{state_hash[:16]}"

        checkpoint = LockStockCheckpoint(
            thread_id=thread_id,
            checkpoint_id=checkpoint_id,
            state=state,
            timestamp=timestamp,
            lockstock_hash=result.state_hash or state_hash,
            sequence=self._sequence
        )

        # Store checkpoint
        self._storage[checkpoint_id] = checkpoint

        return checkpoint_id

    async def get(
        self,
        thread_id: str,
        checkpoint_id: Optional[str] = None
    ) -> Optional[Tuple[Dict[str, Any], str]]:
        """
        Retrieve a checkpoint.

        Args:
            thread_id: Thread identifier
            checkpoint_id: Specific checkpoint ID (latest if None)

        Returns:
            Tuple of (state, checkpoint_id) or None if not found
        """
        if checkpoint_id:
            checkpoint = self._storage.get(checkpoint_id)
        else:
            # Get latest checkpoint for thread
            thread_checkpoints = [
                c for c in self._storage.values()
                if c.thread_id == thread_id
            ]
            if not thread_checkpoints:
                return None
            checkpoint = max(thread_checkpoints, key=lambda c: c.sequence)

        if not checkpoint:
            return None

        # Verify integrity
        state_json = json.dumps(checkpoint.state, sort_keys=True, default=str)
        computed_hash = hashlib.sha256(state_json.encode()).hexdigest()

        return (checkpoint.state, checkpoint.checkpoint_id)

    async def list(
        self,
        thread_id: str,
        limit: int = 10
    ) -> list:
        """
        List checkpoints for a thread.

        Args:
            thread_id: Thread identifier
            limit: Maximum number to return

        Returns:
            List of checkpoint summaries
        """
        thread_checkpoints = [
            c for c in self._storage.values()
            if c.thread_id == thread_id
        ]

        # Sort by sequence descending
        thread_checkpoints.sort(key=lambda c: c.sequence, reverse=True)

        return [
            {
                "checkpoint_id": c.checkpoint_id,
                "sequence": c.sequence,
                "timestamp": c.timestamp,
                "lockstock_hash": c.lockstock_hash
            }
            for c in thread_checkpoints[:limit]
        ]

    async def delete(self, checkpoint_id: str) -> bool:
        """
        Delete a checkpoint.

        Args:
            checkpoint_id: Checkpoint to delete

        Returns:
            True if deleted
        """
        if checkpoint_id in self._storage:
            del self._storage[checkpoint_id]
            return True

        return False

    async def close(self):
        """Close the underlying client."""
        await self.client.close()
