import random
from .models import Email, Action, Observation, ActionType

class GmailEnvironment:
    def __init__(self):
        self.max_steps = 5
        self.reset()

    def reset(self):
        self.current_step = 0
        # Synthetic email data
        self.emails = [
            Email(id="1", sender="hr@company.com", subject="Interview", body="Schedule a time.", is_spam_truth=False, is_priority_truth=True),
            Email(id="2", sender="bot@spam.net", subject="WIN CASH", body="Click link now.", is_spam_truth=True, is_priority_truth=False),
            Email(id="3", sender="mom@home.com", subject="Dinner?", body="Are you coming?", is_spam_truth=False, is_priority_truth=False),
        ]
        return self._get_obs()

    def _get_obs(self):
        # We create a clean version of emails for the agent by hiding the 'truth' fields
        clean_inbox = []
        for e in self.emails:
            clean_email = Email(
                id=e.id,
                sender=e.sender,
                subject=e.subject,
                body=e.body
                # We skip setting is_spam_truth and is_priority_truth here
            )
            clean_inbox.append(clean_email)
            
        return Observation(
            inbox=clean_inbox, 
            steps_taken=self.current_step, 
            max_steps=self.max_steps
        )

    def step(self, action: Action):
        reward = 0
        target = next((e for e in self.emails if e.id == action.email_id), None)
        
        if target:
            if action.action_type == ActionType.MARK_PRIORITY:
                reward = 10 if target.is_priority_truth else -5
            
            elif action.action_type == ActionType.MARK_SPAM:
                reward = 15 if target.is_spam_truth else -20
            
            elif action.action_type == ActionType.AUTO_REPLY:
                if not target.is_spam_truth and not target.is_priority_truth:
                    reward = 5
                else:
                    reward = -2
            
            elif action.action_type == ActionType.ARCHIVE:
                reward = 0 # Neutral action

        self.current_step += 1
        done = self.current_step >= self.max_steps
        return self._get_obs(), reward, done