env.py

"""
Enterprise Email Triage Simulator - OpenEnv Compliant RL Environment

This environment simulates a corporate email inbox where an AI agent must triage
incoming emails based on their content and urgency. The agent must learn to
appropriately handle different types of emails using three available actions.

Author: Senior Python AI Engineer
Purpose: Meta PyTorch Hackathon - Theme 3.1: Professional Tasks (Enterprise Workflows)
"""

from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
from enum import Enum
import random
import json
from pydantic import BaseModel, Field
import openenv.core as oe
from typing import Generic, TypeVar
from reward_system import reward_system, RewardBreakdown


class EmailIntent(Enum):
    """Enumeration of email intent types for classification."""
    ROUTINE_PASSWORD_RESET = "routine_password_reset"
    ANGRY_CLIENT_REFUND = "angry_client_refund"
    VIP_SERVER_OUTAGE = "vip_server_outage"
    GENERAL_INQUIRY = "general_inquiry"
    SPAM = "spam"
    INVOICE_DISCREPANCY = "invoice_discrepancy"
    HR_SENSITIVE = "hr_sensitive"
    SPEAR_PHISHING = "spear_phishing"
    FEATURE_REQUEST = "feature_request"
    MIXED_CHURN = "mixed_churn"


class ActionType(Enum):
    """Enumeration of available agent actions."""
    AUTO_REPLY = "auto_reply"
    ROUTE_TO_HUMAN = "route_to_human"
    ASK_FOR_CLARIFICATION = "ask_for_clarification"


@dataclass
class Email:
    """Represents an email in the corporate inbox."""
    id: str
    sender: str
    subject: str
    body: str
    intent: EmailIntent
    priority: int  # 1-5 scale, 5 being highest
    is_vip: bool = False
    department: Optional[str] = None


class EmailState(BaseModel):
    """State representation for the RL environment."""
    current_email: Optional[Dict[str, Any]] = None
    processed_emails: List[str] = Field(default_factory=list)
    available_actions: List[str] = Field(default_factory=lambda: [action.value for action in ActionType])
    step_count: int = 0
    total_reward: float = 0.0


# Type variables for Generic Environment
ActType = TypeVar('ActType')
ObsType = TypeVar('ObsType') 
StateType = TypeVar('StateType')


class EnterpriseEmailEnv(oe.Environment[ActType, ObsType, StateType]):
    """
    OpenEnv-compliant enterprise email triage environment.
    
    The agent must learn to appropriately handle different types of emails
    using three available actions: auto_reply, route_to_human, and ask_for_clarification.
    
    State Space:
        - Current email details (sender, subject, body, metadata)
        - History of processed emails
        - Available actions
    
    Action Space:
        - 0: auto_reply(email_id, message)
        - 1: route_to_human(email_id, department)
        - 2: ask_for_clarification(email_id)
    
    Reward Structure:
        - +1.0: Correctly route VIP server outage to human
        - +1.0: Correctly auto-reply to routine password reset
        - -1.0: Auto-reply to angry client or VIP (must route to human)
        - -0.5: Invalid tool selection or looping behavior
    """
    
    def __init__(self, max_steps: int = 10):
        """
        Initialize the enterprise email environment.
        
        Args:
            max_steps: Maximum number of steps per episode
        """
        super().__init__()
        self.max_steps = max_steps
        self.current_step = 0
        self.emails = self._load_emails()
        self.processed_emails = []
        self.current_email = None
        self.total_reward = 0.0
        self.last_action = None
        self.action_history = []
        
        # Note: OpenEnv Environment doesn't use traditional gym spaces
        # State is managed through the state property and observations returned from reset/step
        
    def _load_emails(self) -> List[Email]:
        """
        Load emails from dataset.json file, with fallback to hardcoded emails.
        
        Returns:
            List of Email objects representing the corporate inbox
        """
        try:
            with open("dataset.json", "r") as f:
                data = json.load(f)
                emails = []
                
                for email_data in data:
                    # Convert string intent to EmailIntent enum
                    intent_str = email_data.get("intent", "general_inquiry")
                    try:
                        intent = EmailIntent(intent_str)
                    except ValueError:
                        print(f"Warning: Unknown intent '{intent_str}', using general_inquiry")
                        intent = EmailIntent.GENERAL_INQUIRY
                    
                    email = Email(
                        id=email_data["id"],
                        sender=email_data["sender"],
                        subject=email_data["subject"],
                        body=email_data["body"],
                        intent=intent,
                        priority=email_data.get("priority", 3),
                        is_vip=email_data.get("is_vip", False),
                        department=email_data.get("department", "General")
                    )
                    emails.append(email)
                
                print(f"✅ Successfully loaded {len(emails)} emails from dataset.json")
                return emails
                
        except FileNotFoundError:
            print("⚠️  Warning: dataset.json not found, using fallback emails")
            return self._create_fallback_emails()
        except json.JSONDecodeError as e:
            print(f"⚠️  Warning: Error parsing dataset.json: {e}, using fallback emails")
            return self._create_fallback_emails()
        except Exception as e:
            print(f"⚠️  Warning: Error loading dataset.json: {e}, using fallback emails")
            return self._create_fallback_emails()

    def _create_fallback_emails(self) -> List[Email]:
        """
        Create hardcoded list of 5 realistic dummy emails with varying intents.
        
        Returns:
            List of Email objects representing the corporate inbox
        """
        emails = [
            Email(
                id="email_001",
                sender="john.doe@acmecorp.com",
                subject="Password Reset Request",
                body="Hi IT team,\n\nI forgot my password for the corporate portal. Can you reset it for user john.doe? I need access to complete my quarterly reports.\n\nThanks,\nJohn",
                intent=EmailIntent.ROUTINE_PASSWORD_RESET,
                priority=2,
                is_vip=False,
                department="IT"
            ),
            
            # Angry Client Refund Request
            Email(
                id="email_002",
                sender="angrycustomer@retailclient.com",
                subject="REFUND MY MONEY NOW - TERRIBLE SERVICE",
                body="This is outrageous! I ordered a premium widget last month and it arrived broken. Your customer service has been useless and I demand a full refund immediately. I've been a loyal customer for 5 years and this is how you treat me?\n\nI want my money back NOW or I'm switching to your competitor.\n\nFurious,\nMike Chen",
                intent=EmailIntent.ANGRY_CLIENT_REFUND,
                priority=4,
                is_vip=False,
                department="Customer Service"
            ),
            
            # VIP Server Outage (Critical)
            Email(
                id="email_003",
                sender="ceo.vip@fortuneglobal.com",
                subject="🚨 CRITICAL: ALL SERVERS DOWN - ENTERPRISE AT RISK",
                body="This is an emergency. Our entire server infrastructure across all data centers is currently offline. This is affecting millions of users and causing massive revenue loss. We need immediate technical intervention.\n\nThis is our highest priority incident - please escalate to emergency technical teams immediately.\n\nRegards,\nGlobal Infrastructure Team",
                intent=EmailIntent.VIP_SERVER_OUTAGE,
                priority=5,
                is_vip=True,
                department="Emergency Support"
            ),
            
            # General Inquiry
            Email(
                id="email_004",
                sender="potential.client@startup.com",
                subject="Question about partnership opportunities",
                body="Hello,\n\nI'm impressed with your company's work in the enterprise space and would like to explore potential partnership opportunities. Could you provide more information about your B2B collaboration programs?\n\nBest regards,\nSarah Kim",
                intent=EmailIntent.GENERAL_INQUIRY,
                priority=3,
                is_vip=False,
                department="Business Development"
            ),
            
            # Spam
            Email(
                id="email_005",
                sender="winner@lottery-scam.net",
                subject="You've won $1,000,000!!!",
                body="CONGRATULATIONS! You are our lucky winner! Click here immediately to claim your $1,000,000 prize. This offer expires in 24 hours. ACT FAST!\n\nURGENT! URGENT! URGENT!\n\nClaim now!",
                intent=EmailIntent.SPAM,
                priority=1,
                is_vip=False,
                department="Security"
            ),
            
            # EDGE CASE 2: The "Highly Sensitive HR Issue" (Tests empathy vs. automation risk)
            # A massive enterprise risk. If the AI auto-replies to this, it's a disaster. Must go to HR immediately.
            Email(
                id="email_007",
                sender="sarah.jenkins@acmecorp.com",
                subject="CONFIDENTIAL: Formal complaint regarding workplace conduct",
                body="I am writing to formally report an incident of inappropriate conduct by my team lead during yesterday's offsite. I have documented the comments and feel extremely uncomfortable returning to the office. I need to speak with an employee relations representative as soon as possible.\n\nSarah Jenkins",
                intent=EmailIntent.HR_SENSITIVE,
                priority=5,
                is_vip=False,
                department="Human Resources"
            ),
            
            # EDGE CASE 3: The "Spear Phishing" attack (Tests security awareness)
            # It looks like a routine IT ticket, but the domain is fake ("acmecorp-secure.com"). 
            Email(
                id="email_008",
                sender="admin@acmecorp-secure.com",
                subject="ACTION REQUIRED: Update your Okta SSO Credentials",
                body="Dear Employee,\n\nWe are migrating our enterprise single sign-on system. Your current Okta password will expire in 2 hours. Please log in immediately at http://acmecorp-secure-login.net/auth to retain access to your corporate email and GitHub repositories.\n\nIT Support Team",
                intent=EmailIntent.SPEAR_PHISHING,
                priority=4,
                is_vip=False,
                department="Security"
            ),
            
            # EDGE CASE 4: The "Low Priority Feedback" (Tests safe auto-reply usage)
            # This is the perfect candidate for the AI to safely use the `auto_reply` tool to save human time.
            Email(
                id="email_009",
                sender="beta.tester@client.com",
                subject="Suggestion for the new analytics dashboard",
                body="Hey team, I've been using the new beta version of the analytics dashboard. It's pretty good, but it would be awesome if you added a dark mode toggle and the ability to export the charts directly to PDF. Just a thought!\n\nCheers.",
                intent=EmailIntent.FEATURE_REQUEST,
                priority=1,
                is_vip=False,
                department="Product"
            ),
            
            # EDGE CASE 5: The "Ambiguous Multi-Intent" (Tests the Ask For Clarification tool)
            # The client wants to downgrade (Sales/Retention) but also has a technical question (IT/Support). 
            Email(
                id="email_010",
                sender="confused.user@smallbiz.com",
                subject="Downgrading my account + question about data export",
                body="Hello, we are cutting costs and want to drop down from the Enterprise tier to the Basic tier next month. Before we do that, does the Basic tier still allow us to use the API to export our historical data? If not, we might just cancel entirely. Let me know.",
                intent=EmailIntent.MIXED_CHURN,
                priority=3,
                is_vip=False,
                department="Customer Success"
            )
        ]
        
        return emails
    
    @property
    def state(self) -> Dict[str, Any]:
        """
        Get the current environment state (LLM-friendly).
        
        Returns:
            Current environment state as flat dictionary
        """
        return self._get_llm_observation()

    def reset(self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        Reset the environment to initial state.
        
        Args:
            seed: Optional seed for random number generation
            episode_id: Optional episode ID for tracking
        
        Returns:
            Initial state observation
        """
        self.current_step = 0
        self.processed_emails = []
        self.total_reward = 0.0
        self.last_action = None
        self.action_history = []
        
        # Reset emails and select first one
        self.emails = self._load_emails()
        self.current_email = random.choice(self.emails)
        
        return self._get_llm_observation()
    
    def step(self, action: Dict[str, Any]) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
        """
        Execute one step in the environment using LLM tool call format.
        
        Args:
            action: Dictionary containing tool name and arguments
                   Format: {"tool": "tool_name", "arguments": {...}}
            
        Returns:
            Tuple of (observation, reward, done, info) where observation is LLM-friendly
        """
        if self.current_email is None:
            raise ValueError("No current email available. Call reset() first.")
        
        if self.current_email.id in self.processed_emails:
            raise ValueError("Current email has already been processed.")
        
        # Capture current email intent before updating current_email
        processed_intent = self.current_email.intent.value if self.current_email else None
        
        # Parse LLM tool call
        tool_name = action.get("tool")
        arguments = action.get("arguments", {})
        
        # Validate tool name
        valid_tools = [action.value for action in ActionType]
        if tool_name not in valid_tools:
            reward = -0.5
            done = True
            info = {"error": f"Invalid tool '{tool_name}'. Valid tools: {valid_tools}"}
            return self._get_llm_observation(), reward, done, info
        
        # Convert to ActionType
        action_type = ActionType(tool_name)
        
        # Validate arguments based on tool type
        validation_result = self._validate_tool_arguments(action_type, arguments)
        if not validation_result["valid"]:
            reward = -0.5
            done = True
            info = {"error": validation_result["error"]}
            return self._get_llm_observation(), reward, done, info
        
        # Calculate reward with enhanced business logic
        reward = self._calculate_llm_reward(action_type, arguments)
        
        # Update state
        self.processed_emails.append(self.current_email.id)
        self.action_history.append(action_type.value)
        self.total_reward += reward
        self.last_action = action_type
        self.current_step += 1
        
        # Check if episode is done
        done = self.current_step >= self.max_steps or len(self.processed_emails) >= len(self.emails)
        
        # Select next email if available
        if not done:
            remaining_emails = [email for email in self.emails if email.id not in self.processed_emails]
            if remaining_emails:
                self.current_email = random.choice(remaining_emails)
            else:
                done = True
        
        info = {
            "tool_used": action_type.value,
            "arguments": arguments,
            "email_intent": processed_intent,
            "step": self.current_step,
            "total_reward": self.total_reward,
            "validation_passed": True
        }
        
        return self._get_llm_observation(), reward, done, info
    
    def _validate_tool_arguments(self, action_type: ActionType, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """
        Validate tool arguments based on tool type and current email context.
        
        Args:
            action_type: The type of action being performed
            arguments: Dictionary of arguments provided by the LLM
            
        Returns:
            Dictionary with validation result and error message if invalid
        """
        if action_type == ActionType.AUTO_REPLY:
            # Auto-reply requires a message argument
            if "message" not in arguments or not arguments["message"].strip():
                return {
                    "valid": False,
                    "error": "auto_reply tool requires a 'message' argument with non-empty content"
                }
            
            # Check if email_id is provided and matches current email
            if "email_id" in arguments and arguments["email_id"] != self.current_email.id:
                return {
                    "valid": False,
                    "error": f"email_id '{arguments['email_id']}' does not match current email '{self.current_email.id}'"
                }
        
        elif action_type == ActionType.ROUTE_TO_HUMAN:
            # Route to human requires a department argument
            if "department" not in arguments or not arguments["department"].strip():
                return {
                    "valid": False,
                    "error": "route_to_human tool requires a 'department' argument"
                }
            
            # Check if email_id is provided and matches current email
            if "email_id" in arguments and arguments["email_id"] != self.current_email.id:
                return {
                    "valid": False,
                    "error": f"email_id '{arguments['email_id']}' does not match current email '{self.current_email.id}'"
                }
        
        elif action_type == ActionType.ASK_FOR_CLARIFICATION:
            # Ask for clarification requires email_id
            if "email_id" in arguments and arguments["email_id"] != self.current_email.id:
                return {
                    "valid": False,
                    "error": f"email_id '{arguments['email_id']}' does not match current email '{self.current_email.id}'"
                }
        
        return {"valid": True, "error": None}
    
    def _calculate_llm_reward(self, action_type: ActionType, arguments: Dict[str, Any]) -> float:
        """
        Calculate reward using enhanced reward system with multiple independent components.
        
        Args:
            action_type: The type of action being performed
            arguments: Dictionary of arguments provided by the LLM
            
        Returns:
            Total reward value from enhanced reward system
        """
        if self.current_email is None:
            return -0.5
        
        # Convert email to dict format for reward system
        email_dict = {
            "id": self.current_email.id,
            "sender": self.current_email.sender,
            "subject": self.current_email.subject,
            "body": self.current_email.body,
            "intent": self.current_email.intent.value,
            "priority": self.current_email.priority,
            "is_vip": self.current_email.is_vip
        }
        
        # Convert action to dict format for reward system
        action_dict = {
            "tool": action_type.value,
            "arguments": arguments
        }
        
        # Use enhanced reward system
        reward_breakdown = reward_system.calculate_reward(email_dict, action_dict)
        
        # Store breakdown for debugging/monitoring
        if not hasattr(self, 'last_reward_breakdown'):
            self.last_reward_breakdown = reward_breakdown
        else:
            self.last_reward_breakdown = reward_breakdown
        
        return reward_breakdown.total_reward
    
    def get_last_reward_breakdown(self) -> Optional[RewardBreakdown]:
        """Get the detailed breakdown of the last reward calculation"""
        return getattr(self, 'last_reward_breakdown', None)
    
    def _get_llm_observation(self) -> Dict[str, Any]:
        """
        Get LLM-friendly observation as a flat dictionary.
        
        Returns:
            Dictionary with current state information optimized for LLM consumption
        """
        email_data = {}
        if self.current_email:
            email_data = {
                "email_id": self.current_email.id,
                "sender": self.current_email.sender,
                "subject": self.current_email.subject,
                "body": self.current_email.body,
                "intent": self.current_email.intent.value,
                "priority": self.current_email.priority,
                "is_vip": self.current_email.is_vip,
                "suggested_department": self.current_email.department
            }
        
        return {
            "current_email": email_data,
            "processed_emails": self.processed_emails.copy(),
            "available_tools": [
                {
                    "name": "auto_reply",
                    "description": "Automatically reply to an email",
                    "parameters": {
                        "email_id": "string (required)",
                        "message": "string (required)"
                    }
                },
                {
                    "name": "route_to_human", 
                    "description": "Route email to human agent",
                    "parameters": {
                        "email_id": "string (required)",
                        "department": "string (required)"
                    }
                },
                {
                    "name": "ask_for_clarification",
                    "description": "Request more information about email",
                    "parameters": {
                        "email_id": "string (required)"
                    }
                }
            ],
            "step_count": self.current_step,
            "total_reward": self.total_reward,
            "last_action": self.last_action.value if self.last_action else None
        }
    
        
        
    def render(self, mode: str = 'human') -> None:
        """
        Render the current state of the environment.
        
        Args:
            mode: Rendering mode ('human' for text output)
        """
        if mode != 'human':
            return
        
        print("\n" + "="*80)
        print(f"ENTERPRISE EMAIL TRIAGE SIMULATOR - Step {self.current_step}")
        print("="*80)
        
        if self.current_email:
            print(f"\nCURRENT EMAIL:")
            print(f"From: {self.current_email.sender}")
            print(f"Subject: {self.current_email.subject}")
            print(f"Priority: {self.current_email.priority}/5 {'(VIP)' if self.current_email.is_vip else ''}")
            print(f"Body: {self.current_email.body[:200]}...")
            print(f"Intent: {self.current_email.intent.value}")
        
        print(f"\nProcessed Emails: {len(self.processed_emails)}/{len(self.emails)}")
        print(f"Total Reward: {self.total_reward:.2f}")
        print(f"Last Action: {self.last_action.value if self.last_action else 'None'}")
        print("="*80)
    
        
    def close(self) -> None:
        """Clean up environment resources."""
        pass


# Testing function for terminal validation
def test_environment():
    """
    Test function to validate the environment works correctly with LLM tool calls.
    Run this function to test the environment in a terminal.
    """
    print("Initializing Enterprise Email Triage Environment...")
    env = EnterpriseEmailEnv()
    
    print("\nTesting environment reset...")
    obs = env.reset()
    print("Initial observation:")
    print(f"Current email: {obs['current_email']['subject']}")
    print(f"Available tools: {[tool['name'] for tool in obs['available_tools']]}")
    
    print("\nTesting LLM tool call actions...")
    
    # Test 1: Auto-reply to password reset
    print("\n=== Test 1: Auto-reply to password reset ===")
    action1 = {
        "tool": "auto_reply",
        "arguments": {
            "email_id": obs['current_email']['email_id'],
            "message": "I can help you reset your password. Please follow our standard password reset procedure."
        }
    }
    obs, reward, done, info = env.step(action1)
    print(f"Tool: {info['tool_used']}")
    print(f"Reward: {reward:.2f}")
    print(f"Done: {done}")
    
    if not done:
        # Test 2: Route VIP outage to emergency support
        print("\n=== Test 2: Route VIP outage to emergency support ===")
        action2 = {
            "tool": "route_to_human",
            "arguments": {
                "email_id": obs['current_email']['email_id'],
                "department": "Emergency Support"
            }
        }
        obs, reward, done, info = env.step(action2)
        print(f"Tool: {info['tool_used']}")
        print(f"Department: {info['arguments']['department']}")
        print(f"Reward: {reward:.2f}")
        print(f"Done: {done}")
    
    if not done:
        # Test 3: Route angry client to customer service
        print("\n=== Test 3: Route angry client to customer service ===")
        action3 = {
            "tool": "route_to_human",
            "arguments": {
                "email_id": obs['current_email']['email_id'],
                "department": "Customer Service"
            }
        }
        obs, reward, done, info = env.step(action3)
        print(f"Tool: {info['tool_used']}")
        print(f"Department: {info['arguments']['department']}")
        print(f"Reward: {reward:.2f}")
        print(f"Done: {done}")
    
    if not done:
        # Test 4: Ask for clarification on general inquiry
        print("\n=== Test 4: Ask for clarification on general inquiry ===")
        action4 = {
            "tool": "ask_for_clarification",
            "arguments": {
                "email_id": obs['current_email']['email_id']
            }
        }
        obs, reward, done, info = env.step(action4)
        print(f"Tool: {info['tool_used']}")
        print(f"Reward: {reward:.2f}")
        print(f"Done: {done}")
    
    if not done:
        # Test 5: Route spam to security
        print("\n=== Test 5: Route spam to security ===")
        action5 = {
            "tool": "route_to_human",
            "arguments": {
                "email_id": obs['current_email']['email_id'],
                "department": "Security"
            }
        }
        obs, reward, done, info = env.step(action5)
        print(f"Tool: {info['tool_used']}")
        print(f"Department: {info['arguments']['department']}")
        print(f"Reward: {reward:.2f}")
        print(f"Done: {done}")
    
    print(f"\nFinal total reward: {env.total_reward:.2f}")
    env.close()
    print("Environment test completed successfully!")


def test_validation_errors():
    """
    Test function to validate error handling for invalid tool calls.
    """
    print("\n" + "="*60)
    print("TESTING VALIDATION ERRORS")
    print("="*60)
    
    env = EnterpriseEmailEnv()
    obs = env.reset()
    
    # Test 1: Invalid tool name
    print("\n=== Test 1: Invalid tool name ===")
    invalid_action = {
        "tool": "invalid_tool",
        "arguments": {"email_id": obs['current_email']['email_id']}
    }
    obs, reward, done, info = env.step(invalid_action)
    print(f"Error: {info['error']}")
    print(f"Reward: {reward:.2f}")
    
    # Test 2: Missing required arguments
    print("\n=== Test 2: Auto-reply without message ===")
    env.reset()
    obs = env.reset()
    invalid_action2 = {
        "tool": "auto_reply",
        "arguments": {"email_id": obs['current_email']['email_id']}
    }
    obs, reward, done, info = env.step(invalid_action2)
    print(f"Error: {info['error']}")
    print(f"Reward: {reward:.2f}")
    
    # Test 3: Wrong email_id
    print("\n=== Test 3: Wrong email_id ===")
    env.reset()
    obs = env.reset()
    invalid_action3 = {
        "tool": "route_to_human",
        "arguments": {
            "email_id": "wrong_email_id",
            "department": "IT"
        }
    }
    obs, reward, done, info = env.step(invalid_action3)
    print(f"Error: {info['error']}")
    print(f"Reward: {reward:.2f}")
    
    env.close()
    print("\nValidation error tests completed!")


if __name__ == "__main__":
    test_environment()
    test_validation_errors()