tasks.py

import copy
from typing import List, Dict, Any, Tuple
from models import Email, RewardInfo

class Task:
    def __init__(self, name: str, description: str):
        self.name = name
        self.description = description

    def get_initial_emails(self) -> List[Email]:
        raise NotImplementedError

    def initialize_metrics(self) -> Dict[str, Any]:
        return {}

    def grade(self, emails: List[Email], metrics: Dict[str, Any], action: Dict[str, Any]) -> RewardInfo:
        raise NotImplementedError


class EasyTask(Task):
    def __init__(self):
        super().__init__(
            name="easy",
            description="Move the email with the subject 'Invoice' to the 'Finance' folder."
        )

    def get_initial_emails(self) -> List[Email]:
        return [
            Email(id="email_1", sender="vendor@supplies.com", subject="Invoice #1024", body="Please find attached the invoice for your recent order.", folder="inbox"),
            Email(id="email_2", sender="team@company.com", subject="Weekly Sync", body="Let's meet at 10 AM.", folder="inbox")
        ]

    def grade(self, emails: List[Email], metrics: Dict[str, Any], action: Dict[str, Any]) -> RewardInfo:
        invoice_email = next((e for e in emails if "Invoice" in e.subject), None)
        is_done = action.get("action_type") == "submit"
        
        # Grading
        score = 0.0
        if invoice_email and invoice_email.folder == "Finance":
            score = 1.0
            
        # Penalize if they deleted it
        if not invoice_email:
            score = 0.0
            is_done = True
        
        return RewardInfo(score=score, is_done=is_done, metrics={"invoice_moved": score == 1.0})


class MediumTask(Task):
    def __init__(self):
        super().__init__(
            name="medium",
            description="Find the unread email from 'boss@company.com', read it, and reply with exactly 'Will do.'."
        )

    def get_initial_emails(self) -> List[Email]:
        return [
            Email(id="email_1", sender="newsletter@tech.com", subject="Daily Tech News", body="Here is the news...", folder="inbox", is_read=False),
            Email(id="email_2", sender="boss@company.com", subject="Project Update Needed", body="Please send me the project update by EOD.", folder="inbox", is_read=False)
        ]

    def initialize_metrics(self) -> Dict[str, Any]:
        return {"boss_email_read": False, "replied": False}

    def grade(self, emails: List[Email], metrics: Dict[str, Any], action: Dict[str, Any]) -> RewardInfo:
        boss_email = next((e for e in emails if e.sender == "boss@company.com"), None)
        
        if boss_email and boss_email.is_read:
            metrics["boss_email_read"] = True
            
        if action.get("action_type") == "reply" and action.get("email_id") == getattr(boss_email, "id", ""):
            if action.get("reply_body", "").strip() == "Will do.":
                metrics["replied"] = True
                
        is_done = action.get("action_type") == "submit"
        
        score = 0.0
        if metrics["boss_email_read"]:
            score += 0.5
        if metrics["replied"]:
            score += 0.5
            
        if not boss_email:
            score = 0.0
            is_done = True
            
        return RewardInfo(score=score, is_done=is_done, metrics=metrics)


class HardTask(Task):
    def __init__(self):
        super().__init__(
            name="hard",
            description="Inbox triage: 1) Delete all spam (from 'spam@scam.com'). 2) Move Support emails to 'Support' folder. 3) Flag any email with 'Urgent' in the subject."
        )
        self.total_spam = 2
        self.total_support = 2
        self.total_urgent = 1

    def get_initial_emails(self) -> List[Email]:
        return [
            Email(id="email_1", sender="spam@scam.com", subject="Win a free iPhone!", body="Click here to win.", folder="inbox"),
            Email(id="email_2", sender="customer1@gmail.com", subject="Need help with login", body="I can't login to my account.", folder="inbox"),
            Email(id="email_3", sender="boss@company.com", subject="Urgent: Fix production issue", body="The server is down!", folder="inbox", is_flagged=False),
            Email(id="email_4", sender="spam@scam.com", subject="Cheap meds", body="Buy now.", folder="inbox"),
            Email(id="email_5", sender="customer2@yahoo.com", subject="Refund request", body="I want a refund.", folder="inbox")
        ]

    def grade(self, emails: List[Email], metrics: Dict[str, Any], action: Dict[str, Any]) -> RewardInfo:
        # Check spam (should be deleted, meaning not in emails)
        spam_emails = [e for e in emails if e.sender == "spam@scam.com"]
        spam_deleted = self.total_spam - len(spam_emails)
        
        # Check support (can use a simple heuristic: customer in email)
        support_emails = [e for e in emails if "customer" in e.sender and e.folder == "Support"]
        
        # Check urgent
        urgent_emails = [e for e in emails if "Urgent" in e.subject and e.is_flagged]
        
        score_spam = (spam_deleted / self.total_spam) * 0.4
        score_support = (len(support_emails) / self.total_support) * 0.4
        score_urgent = (len(urgent_emails) / self.total_urgent) * 0.2
        
        total_score = score_spam + score_support + score_urgent
        
        # Penalties: Did they delete the boss's email?
        boss_email = next((e for e in emails if e.sender == "boss@company.com"), None)
        if not boss_email:
            total_score -= 1.0 # Huge penalty
            
        total_score = max(0.0, min(1.0, total_score))
        
        is_done = action.get("action_type") == "submit" or (not boss_email)
        
        return RewardInfo(
            score=total_score, 
            is_done=is_done, 
            metrics={
                "spam_deleted": spam_deleted,
                "support_moved": len(support_emails),
                "urgent_flagged": len(urgent_emails)
            }
        )

TASKS = {
    "easy": EasyTask(),
    "medium": MediumTask(),
    "hard": HardTask()
}

def get_task(task_name: str) -> Task:
    if task_name not in TASKS:
        raise ValueError(f"Unknown task: {task_name}. Available: {list(TASKS.keys())}")
    return TASKS[task_name]