rpa_vision_v3/core/coaching/session_persistence.py

"""
COACHING Session Persistence Module

Provides persistence layer for COACHING sessions to enable:
- Save session state for recovery after interruption
- Resume sessions from last known state
- Track session history and statistics
"""

import json
import os
import shutil
from dataclasses import dataclass, field, asdict
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional, Any
import threading
import uuid


class SessionStatus(str, Enum):
    """Status of a COACHING session."""
    ACTIVE = "active"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"
    ABANDONED = "abandoned"


@dataclass
class CoachingDecisionRecord:
    """Record of a single coaching decision."""
    step_index: int
    node_id: str
    action_type: str
    decision: str  # accept, reject, correct, manual, skip
    correction: Optional[Dict[str, Any]] = None
    feedback: Optional[str] = None
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    execution_success: Optional[bool] = None

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'CoachingDecisionRecord':
        return cls(**data)


@dataclass
class CoachingSessionState:
    """
    Complete state of a COACHING session.

    This state can be persisted and recovered to resume an interrupted session.
    """
    session_id: str
    workflow_id: str
    execution_id: str
    status: SessionStatus = SessionStatus.ACTIVE
    current_step_index: int = 0
    total_steps: int = 0
    decisions: List[CoachingDecisionRecord] = field(default_factory=list)
    stats: Dict[str, int] = field(default_factory=lambda: {
        'suggestions_made': 0,
        'accepted': 0,
        'rejected': 0,
        'corrected': 0,
        'manual_executions': 0,
        'skipped': 0
    })
    variables: Dict[str, Any] = field(default_factory=dict)
    started_at: str = field(default_factory=lambda: datetime.now().isoformat())
    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
    completed_at: Optional[str] = None
    error_message: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        data = {
            'session_id': self.session_id,
            'workflow_id': self.workflow_id,
            'execution_id': self.execution_id,
            'status': self.status.value if isinstance(self.status, SessionStatus) else self.status,
            'current_step_index': self.current_step_index,
            'total_steps': self.total_steps,
            'decisions': [d.to_dict() for d in self.decisions],
            'stats': self.stats,
            'variables': self.variables,
            'started_at': self.started_at,
            'updated_at': self.updated_at,
            'completed_at': self.completed_at,
            'error_message': self.error_message,
            'metadata': self.metadata
        }
        return data

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'CoachingSessionState':
        decisions = [
            CoachingDecisionRecord.from_dict(d)
            for d in data.get('decisions', [])
        ]
        status = data.get('status', 'active')
        if isinstance(status, str):
            status = SessionStatus(status)

        return cls(
            session_id=data['session_id'],
            workflow_id=data['workflow_id'],
            execution_id=data['execution_id'],
            status=status,
            current_step_index=data.get('current_step_index', 0),
            total_steps=data.get('total_steps', 0),
            decisions=decisions,
            stats=data.get('stats', {}),
            variables=data.get('variables', {}),
            started_at=data.get('started_at', datetime.now().isoformat()),
            updated_at=data.get('updated_at', datetime.now().isoformat()),
            completed_at=data.get('completed_at'),
            error_message=data.get('error_message'),
            metadata=data.get('metadata', {})
        )

    def update_timestamp(self) -> None:
        """Update the updated_at timestamp."""
        self.updated_at = datetime.now().isoformat()

    def add_decision(self, decision: CoachingDecisionRecord) -> None:
        """Add a decision and update stats."""
        self.decisions.append(decision)
        self.stats['suggestions_made'] += 1

        if decision.decision == 'accept':
            self.stats['accepted'] += 1
        elif decision.decision == 'reject':
            self.stats['rejected'] += 1
        elif decision.decision == 'correct':
            self.stats['corrected'] += 1
        elif decision.decision == 'manual':
            self.stats['manual_executions'] += 1
        elif decision.decision == 'skip':
            self.stats['skipped'] += 1

        self.current_step_index += 1
        self.update_timestamp()

    def get_acceptance_rate(self) -> float:
        """Calculate acceptance rate."""
        total = self.stats['accepted'] + self.stats['rejected'] + self.stats['corrected']
        if total == 0:
            return 0.0
        return self.stats['accepted'] / total

    def get_correction_rate(self) -> float:
        """Calculate correction rate."""
        total = self.stats['accepted'] + self.stats['rejected'] + self.stats['corrected']
        if total == 0:
            return 0.0
        return self.stats['corrected'] / total

    def can_resume(self) -> bool:
        """Check if session can be resumed."""
        return self.status in [SessionStatus.ACTIVE, SessionStatus.PAUSED]


class CoachingSessionPersistence:
    """
    Persistence layer for COACHING sessions.

    Handles saving, loading, and managing COACHING session states.
    """

    def __init__(self, storage_path: Optional[Path] = None):
        """
        Initialize persistence layer.

        Args:
            storage_path: Path to store session data. Defaults to data/coaching_sessions
        """
        if storage_path is None:
            storage_path = Path(__file__).parent.parent.parent / 'data' / 'coaching_sessions'
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(parents=True, exist_ok=True)

        # Index file for quick lookup
        self._index_file = self.storage_path / 'sessions_index.json'
        self._index: Dict[str, Dict[str, Any]] = {}
        self._lock = threading.Lock()

        self._load_index()

    def _load_index(self) -> None:
        """Load the sessions index."""
        if self._index_file.exists():
            try:
                with open(self._index_file, 'r') as f:
                    self._index = json.load(f)
            except Exception as e:
                print(f"Warning: Could not load sessions index: {e}")
                self._index = {}

    def _save_index(self) -> None:
        """Save the sessions index."""
        try:
            temp_file = self._index_file.with_suffix('.tmp')
            with open(temp_file, 'w') as f:
                json.dump(self._index, f, indent=2)
            shutil.move(str(temp_file), str(self._index_file))
        except Exception as e:
            print(f"Warning: Could not save sessions index: {e}")

    def _session_file(self, session_id: str) -> Path:
        """Get the path for a session file."""
        return self.storage_path / f"{session_id}.json"

    def create_session(
        self,
        workflow_id: str,
        execution_id: str,
        total_steps: int = 0,
        variables: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> CoachingSessionState:
        """
        Create a new COACHING session.

        Args:
            workflow_id: ID of the workflow being coached
            execution_id: ID of the execution
            total_steps: Total number of steps in the workflow
            variables: Initial variables
            metadata: Additional metadata

        Returns:
            New session state
        """
        session_id = f"coaching_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"

        session = CoachingSessionState(
            session_id=session_id,
            workflow_id=workflow_id,
            execution_id=execution_id,
            total_steps=total_steps,
            variables=variables or {},
            metadata=metadata or {}
        )

        self.save_session(session)
        return session

    def save_session(self, session: CoachingSessionState) -> None:
        """
        Save a session state to disk.

        Args:
            session: Session state to save
        """
        with self._lock:
            session.update_timestamp()

            # Save session file
            session_file = self._session_file(session.session_id)
            try:
                temp_file = session_file.with_suffix('.tmp')
                with open(temp_file, 'w') as f:
                    json.dump(session.to_dict(), f, indent=2)
                shutil.move(str(temp_file), str(session_file))
            except Exception as e:
                raise RuntimeError(f"Failed to save session {session.session_id}: {e}")

            # Update index
            self._index[session.session_id] = {
                'workflow_id': session.workflow_id,
                'execution_id': session.execution_id,
                'status': session.status.value if isinstance(session.status, SessionStatus) else session.status,
                'started_at': session.started_at,
                'updated_at': session.updated_at,
                'current_step': session.current_step_index,
                'total_steps': session.total_steps
            }
            self._save_index()

    def load_session(self, session_id: str) -> Optional[CoachingSessionState]:
        """
        Load a session state from disk.

        Args:
            session_id: ID of the session to load

        Returns:
            Session state or None if not found
        """
        session_file = self._session_file(session_id)
        if not session_file.exists():
            return None

        try:
            with open(session_file, 'r') as f:
                data = json.load(f)
            return CoachingSessionState.from_dict(data)
        except Exception as e:
            print(f"Warning: Could not load session {session_id}: {e}")
            return None

    def delete_session(self, session_id: str) -> bool:
        """
        Delete a session.

        Args:
            session_id: ID of the session to delete

        Returns:
            True if deleted, False otherwise
        """
        with self._lock:
            session_file = self._session_file(session_id)
            if session_file.exists():
                session_file.unlink()

            if session_id in self._index:
                del self._index[session_id]
                self._save_index()
                return True

            return False

    def list_sessions(
        self,
        workflow_id: Optional[str] = None,
        status: Optional[SessionStatus] = None,
        limit: int = 100
    ) -> List[Dict[str, Any]]:
        """
        List sessions with optional filters.

        Args:
            workflow_id: Filter by workflow ID
            status: Filter by status
            limit: Maximum number of sessions to return

        Returns:
            List of session summaries
        """
        sessions = []

        for session_id, info in self._index.items():
            # Apply filters
            if workflow_id and info.get('workflow_id') != workflow_id:
                continue
            if status:
                session_status = info.get('status', 'active')
                if session_status != status.value:
                    continue

            sessions.append({
                'session_id': session_id,
                **info
            })

        # Sort by updated_at descending
        sessions.sort(key=lambda x: x.get('updated_at', ''), reverse=True)

        return sessions[:limit]

    def get_resumable_sessions(self, workflow_id: str) -> List[CoachingSessionState]:
        """
        Get all resumable sessions for a workflow.

        Args:
            workflow_id: Workflow ID to filter by

        Returns:
            List of resumable session states
        """
        resumable = []

        for session_info in self.list_sessions(workflow_id=workflow_id):
            status = session_info.get('status', 'active')
            if status in ['active', 'paused']:
                session = self.load_session(session_info['session_id'])
                if session and session.can_resume():
                    resumable.append(session)

        return resumable

    def pause_session(self, session_id: str) -> bool:
        """
        Pause an active session.

        Args:
            session_id: Session to pause

        Returns:
            True if paused, False otherwise
        """
        session = self.load_session(session_id)
        if session and session.status == SessionStatus.ACTIVE:
            session.status = SessionStatus.PAUSED
            self.save_session(session)
            return True
        return False

    def resume_session(self, session_id: str) -> Optional[CoachingSessionState]:
        """
        Resume a paused session.

        Args:
            session_id: Session to resume

        Returns:
            Resumed session or None if cannot resume
        """
        session = self.load_session(session_id)
        if session and session.can_resume():
            session.status = SessionStatus.ACTIVE
            self.save_session(session)
            return session
        return None

    def complete_session(
        self,
        session_id: str,
        success: bool = True,
        error_message: Optional[str] = None
    ) -> Optional[CoachingSessionState]:
        """
        Mark a session as completed.

        Args:
            session_id: Session to complete
            success: Whether the session completed successfully
            error_message: Error message if failed

        Returns:
            Completed session or None
        """
        session = self.load_session(session_id)
        if session:
            session.status = SessionStatus.COMPLETED if success else SessionStatus.FAILED
            session.completed_at = datetime.now().isoformat()
            session.error_message = error_message
            self.save_session(session)
            return session
        return None

    def abandon_session(self, session_id: str) -> bool:
        """
        Mark a session as abandoned.

        Args:
            session_id: Session to abandon

        Returns:
            True if abandoned
        """
        session = self.load_session(session_id)
        if session:
            session.status = SessionStatus.ABANDONED
            session.completed_at = datetime.now().isoformat()
            self.save_session(session)
            return True
        return False

    def cleanup_old_sessions(self, max_age_days: int = 30) -> int:
        """
        Remove sessions older than max_age_days.

        Args:
            max_age_days: Maximum age in days

        Returns:
            Number of sessions removed
        """
        cutoff = datetime.now().timestamp() - (max_age_days * 24 * 3600)
        removed = 0

        with self._lock:
            to_remove = []
            for session_id, info in self._index.items():
                try:
                    updated_at = datetime.fromisoformat(info.get('updated_at', '')).timestamp()
                    if updated_at < cutoff:
                        to_remove.append(session_id)
                except:
                    pass

            for session_id in to_remove:
                session_file = self._session_file(session_id)
                if session_file.exists():
                    session_file.unlink()
                del self._index[session_id]
                removed += 1

            if removed > 0:
                self._save_index()

        return removed

    def get_statistics(self) -> Dict[str, Any]:
        """
        Get overall statistics about COACHING sessions.

        Returns:
            Statistics dictionary
        """
        total = len(self._index)
        by_status = {}
        total_decisions = 0
        total_accepted = 0
        total_corrected = 0

        for session_id, info in self._index.items():
            status = info.get('status', 'unknown')
            by_status[status] = by_status.get(status, 0) + 1

            # Load full session for detailed stats
            session = self.load_session(session_id)
            if session:
                total_decisions += session.stats.get('suggestions_made', 0)
                total_accepted += session.stats.get('accepted', 0)
                total_corrected += session.stats.get('corrected', 0)

        return {
            'total_sessions': total,
            'by_status': by_status,
            'total_decisions': total_decisions,
            'total_accepted': total_accepted,
            'total_corrected': total_corrected,
            'overall_acceptance_rate': total_accepted / total_decisions if total_decisions > 0 else 0,
            'overall_correction_rate': total_corrected / total_decisions if total_decisions > 0 else 0
        }


# Global instance
_global_persistence: Optional[CoachingSessionPersistence] = None


def get_coaching_persistence(storage_path: Optional[Path] = None) -> CoachingSessionPersistence:
    """
    Get or create the global coaching session persistence instance.

    Args:
        storage_path: Optional custom storage path

    Returns:
        CoachingSessionPersistence instance
    """
    global _global_persistence
    if _global_persistence is None:
        _global_persistence = CoachingSessionPersistence(storage_path)
    return _global_persistence