"""
Cubesquare Session Helper

Simplifies agent state persistence in the agent_sandbox.
"""
import json
import logging
from pathlib import Path
from typing import Dict, Any, Optional
from datetime import datetime

class SessionHelper:
    """
    Manages session state for an agent in the sandbox.
    
    Usage:
        session = SessionHelper("agent_01")
        session.save({"status": "active", "task": "observing"})
        data = session.load()
    """
    
    def __init__(self, agent_id: str):
        self.agent_id = agent_id
        # Use relative path from execution root or package root
        self.sandbox_root = Path("agent_sandbox")
        self.agent_dir = self.sandbox_root / agent_id
        self.session_file = self.agent_dir / "session.json"
        
        # Ensure directory exists
        if not self.agent_dir.exists():
            try:
                self.agent_dir.mkdir(parents=True, exist_ok=True)
            except Exception as e:
                logging.warning(f"Could not create session directory {self.agent_dir}: {e}")

    def save(self, data: Dict[str, Any]) -> bool:
        """Save session data atomically."""
        try:
            # Add metadata if not present
            if "_meta" not in data:
                data["_meta"] = {
                    "agent_id": self.agent_id,
                    "updated_at": datetime.now().isoformat()
                }
            else:
                data["_meta"]["updated_at"] = datetime.now().isoformat()
                
            # atomic write
            temp_file = self.session_file.with_suffix('.tmp')
            with open(temp_file, 'w') as f:
                json.dump(data, f, indent=2)
            temp_file.replace(self.session_file)
            return True
        except Exception as e:
            logging.error(f"Failed to save session for {self.agent_id}: {e}")
            return False

    def load(self) -> Dict[str, Any]:
        """Load session data or return empty dict if not found."""
        if not self.session_file.exists():
            return {}
        
        try:
            with open(self.session_file, 'r') as f:
                return json.load(f)
        except Exception as e:
            logging.error(f"Failed to load session for {self.agent_id}: {e}")
            return {}

    def clear(self) -> bool:
        """Clear session data."""
        try:
            if self.session_file.exists():
                self.session_file.unlink()
            return True
        except Exception as e:
            logging.error(f"Failed to clear session for {self.agent_id}: {e}")
            return False
