"""
ExecutionContext class for maintaining workflow state and data.
"""

import json
import uuid
from datetime import datetime
from typing import Dict, Any, Optional


class ExecutionContext:
    """
    Maintains workflow execution state and data throughout the workflow lifecycle.
    
    Provides JSON serialization capabilities for persistence and context passing
    between nodes. Tracks execution metadata and user-specific information.
    """
    
    def __init__(self, workflow_id: str, user_id: str):
        """
        Initialize execution context.
        
        Args:
            workflow_id: Unique identifier for the workflow being executed
            user_id: Identifier for the user executing the workflow
        """
        self.workflow_id = workflow_id
        self.user_id = user_id
        self.execution_id = str(uuid.uuid4())
        self.start_time = datetime.utcnow()
        self.data: Dict[str, Any] = {}
        self.variables: Dict[str, Any] = {}
        self.error_handling: str = "continue"  # continue|stop|retry
        self.current_node_id: Optional[str] = None
        self.previous_node_id:Optional[str] = None
    
    def set_data(self, key: str, value: Any) -> None:
        """
        Set data in the execution context.
        
        Args:
            key: Data key
            value: Data value (must be JSON serializable)
        """
        self.data[key] = value
    
    def get_data(self, key: str, default: Any = None) -> Any:
        """
        Get data from the execution context.
        
        Args:
            key: Data key to retrieve
            default: Default value if key not found
            
        Returns:
            Data value or default
        """
        return self.data.get(key, default)
    
    def set_variable(self, key: str, value: Any) -> None:
        """
        Set a variable in the execution context.
        
        Args:
            key: Variable key
            value: Variable value (must be JSON serializable)
        """
        self.variables[key] = value
    
    def get_variable(self, key: str, default: Any = None) -> Any:
        """
        Get a variable from the execution context.
        
        Args:
            key: Variable key to retrieve
            default: Default value if key not found
            
        Returns:
            Variable value or default
        """
        return self.variables.get(key, default)
    
    def merge_data(self, data: Dict[str, Any]) -> None:
        """
        Merge additional data into the context.
        
        Args:
            data: Dictionary of data to merge
        """
        self.data.update(data)
    
    def to_json(self) -> str:
        """
        Serialize context to JSON string.
        
        Returns:
            JSON string representation of the context
        """
        context_dict = {
            'workflow_id': self.workflow_id,
            'user_id': self.user_id,
            'execution_id': self.execution_id,
            'start_time': self.start_time.isoformat(),
            'data': self.data,
            'variables': self.variables,
            'error_handling': self.error_handling,
            'current_node_id': self.current_node_id
        }
        return json.dumps(context_dict, default=str)
    
    @classmethod
    def from_json(cls, json_str: str) -> 'ExecutionContext':
        """
        Deserialize context from JSON string.
        
        Args:
            json_str: JSON string representation of context
            
        Returns:
            ExecutionContext instance
        """
        context_dict = json.loads(json_str)
        
        # Create new instance
        context = cls(
            workflow_id=context_dict['workflow_id'],
            user_id=context_dict['user_id']
        )
        
        # Restore state
        context.execution_id = context_dict['execution_id']
        context.start_time = datetime.fromisoformat(context_dict['start_time'])
        context.data = context_dict.get('data', {})
        context.variables = context_dict.get('variables', {})
        context.error_handling = context_dict.get('error_handling', 'continue')
        context.current_node_id = context_dict.get('current_node_id')
        
        return context
    
    def copy(self) -> 'ExecutionContext':
        """
        Create a copy of the execution context.
        
        Returns:
            New ExecutionContext instance with copied data
        """
        new_context = ExecutionContext(self.workflow_id, self.user_id)
        new_context.execution_id = self.execution_id
        new_context.start_time = self.start_time
        new_context.data = self.data.copy()
        new_context.variables = self.variables.copy()
        new_context.error_handling = self.error_handling
        new_context.current_node_id = self.current_node_id
        return new_context
    
    def __str__(self) -> str:
        """String representation of the context."""
        return f"ExecutionContext(workflow_id={self.workflow_id}, execution_id={self.execution_id})"