agents/storage_agent.py

"""
Storage Agent
Abstract storage interface with Supabase implementation stub.
"""

from typing import Dict, Any, Optional
from abc import abstractmethod

from core.agent_base import Agent
from core.errors import StorageError, ProofNotFoundError
from models.proof import Proof
from config.settings import settings


class StorageAgent(Agent):
    """
    Abstract storage interface for proof persistence.
    """
    
    @abstractmethod
    def save_proof(self, proof: Proof) -> Dict[str, Any]:
        """
        Save proof to storage.
        
        Args:
            proof: Proof object to save
            
        Returns:
            Storage response with proof_id and status
        """
        pass
    
    @abstractmethod
    def get_proof(self, proof_id: str) -> Optional[Proof]:
        """
        Retrieve proof from storage.
        
        Args:
            proof_id: Unique proof identifier
            
        Returns:
            Proof object if found, None otherwise
        """
        pass
    
    def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute storage operation based on action.
        
        Expected input_data:
            {
                "action": "save" | "get",
                "proof": Proof (for save),
                "proof_id": str (for get)
            }
        """
        action = input_data.get("action")
        
        if action == "save":
            proof = input_data.get("proof")
            if not proof:
                raise StorageError("Missing 'proof' for save action")
            return self.save_proof(proof)
        
        elif action == "get":
            proof_id = input_data.get("proof_id")
            if not proof_id:
                raise StorageError("Missing 'proof_id' for get action")
            proof = self.get_proof(proof_id)
            if not proof:
                raise ProofNotFoundError(f"Proof not found: {proof_id}")
            return {"proof": proof}
        
        else:
            raise StorageError(f"Invalid action: {action}")


class SupabaseStorageAgent(StorageAgent):
    """
    Supabase-backed storage implementation.
    Uses environment variables for credentials - never hardcoded.
    """
    
    def __init__(self):
        super().__init__()
        self._client = None
        self._init_client()
    
    def _init_client(self):
        """
        Initialize Supabase client using environment variables.
        This is a stub - actual implementation would use supabase-py.
        """
        if not settings.validate():
            raise StorageError(
                "Supabase credentials not configured. "
                "Set SUPABASE_URL and SUPABASE_KEY environment variables."
            )
        
        # TODO: Initialize actual Supabase client when library is added
        # from supabase import create_client
        # self._client = create_client(settings.SUPABASE_URL, settings.SUPABASE_KEY)
        
        # For now, using in-memory storage for testing
        self._memory_store = {}
    
    def save_proof(self, proof: Proof) -> Dict[str, Any]:
        """
        Save proof to Supabase table.
        Currently using in-memory stub.
        """
        try:
            proof_data = proof.to_dict()
            
            # TODO: Replace with actual Supabase insert
            # response = self._client.table(settings.SUPABASE_TABLE).insert(proof_data).execute()
            
            # In-memory stub
            self._memory_store[proof.proof_id] = proof_data
            
            return {
                "success": True,
                "proof_id": proof.proof_id,
                "message": "Proof saved successfully"
            }
            
        except Exception as e:
            raise StorageError(f"Failed to save proof: {str(e)}") from e
    
    def get_proof(self, proof_id: str) -> Optional[Proof]:
        """
        Retrieve proof from Supabase table.
        Currently using in-memory stub.
        """
        try:
            # TODO: Replace with actual Supabase query
            # response = self._client.table(settings.SUPABASE_TABLE).select("*").eq("proof_id", proof_id).execute()
            
            # In-memory stub
            proof_data = self._memory_store.get(proof_id)
            
            if not proof_data:
                return None
            
            return Proof.from_dict(proof_data)
            
        except Exception as e:
            raise StorageError(f"Failed to retrieve proof: {str(e)}") from e