src/data_generator.py

"""
Realistic email data generator for the Email Triage environment.

Produces emails with varied senders, subjects, domains, and ground-truth
folder assignments.  Supports three difficulty tiers and injects ambiguous
cases to test agent reasoning.
"""

from __future__ import annotations

import logging
import random
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple

from src.models import Email

logger = logging.getLogger(__name__)

# ── VIP Sender Profiles ─────────────────────────────────────────────────────

VIP_SENDERS: List[Dict] = [
    {
        "name": "Sarah Chen",
        "email": "sarah.chen@acmecorp.com",
        "domain": "acmecorp.com",
        "role": "CEO",
        "typical_folders": ["work"],
    },
    {
        "name": "James Rodriguez",
        "email": "j.rodriguez@acmecorp.com",
        "domain": "acmecorp.com",
        "role": "CFO",
        "typical_folders": ["finance"],
    },
    {
        "name": "Priya Patel",
        "email": "priya.patel@acmecorp.com",
        "domain": "acmecorp.com",
        "role": "VP Engineering",
        "typical_folders": ["work", "meetings"],
    },
    {
        "name": "Michael Brooks",
        "email": "m.brooks@acmecorp.com",
        "domain": "acmecorp.com",
        "role": "Finance Manager",
        "typical_folders": ["finance"],
    },
    {
        "name": "Lisa Wang",
        "email": "lisa.wang@acmecorp.com",
        "domain": "acmecorp.com",
        "role": "HR Director",
        "typical_folders": ["work", "meetings"],
    },
]

VIP_EMAILS = {v["email"] for v in VIP_SENDERS}

# ── Email Templates ──────────────────────────────────────────────────────────
# Each template: (subject_template, body_template, ground_truth_folder,
#                  priority, has_attachment, category_hint)

WORK_TEMPLATES: List[Dict] = [
    {
        "subject": "Re: Sprint {n} Planning - Action Items",
        "body": "Hi team, following up on our sprint planning session. Please review the attached action items and update your Jira tickets by EOD. Key deliverables for this sprint include the API refactor and the new dashboard components.",
        "folder": "work",
        "priority": False,
        "attachment": False,
        "hint": "project",
    },
    {
        "subject": "Code Review Request: PR #{n} - Auth Module Refactor",
        "body": "I've submitted a pull request for the authentication module refactor. Could you review the changes when you get a chance? Main changes: JWT token rotation, session management improvements, and rate limiting updates.",
        "folder": "work",
        "priority": False,
        "attachment": False,
        "hint": "code_review",
    },
    {
        "subject": "Production Incident - API Latency Spike",
        "body": "URGENT: We're seeing p99 latency spike to 2.3s on the /api/users endpoint. Root cause appears to be a slow database query. I've opened an incident channel. Please join if you're available.",
        "folder": "work",
        "priority": True,
        "attachment": False,
        "hint": "incident",
    },
    {
        "subject": "Updated: Project {project} Technical Specification",
        "body": "I've updated the technical spec for Project {project} based on feedback from the architecture review. Key changes include the migration strategy and the new caching layer design. Please review and comment.",
        "folder": "work",
        "priority": False,
        "attachment": True,
        "hint": "documentation",
    },
    {
        "subject": "Onboarding Checklist for New Team Member",
        "body": "We have a new engineer joining on Monday. Please make sure the following are set up: Git access, Slack channels, dev environment, and CI/CD pipeline permissions. Let me know if you need any help.",
        "folder": "work",
        "priority": False,
        "attachment": True,
        "hint": "hr",
    },
    {
        "subject": "Weekly Status Report - Week {n}",
        "body": "Here's the weekly status update. Completed: Auth module v2, API docs refresh. In progress: Dashboard redesign, performance optimization. Blocked: Waiting on third-party API credentials.",
        "folder": "work",
        "priority": False,
        "attachment": False,
        "hint": "status",
    },
    {
        "subject": "Re: Deployment Schedule for v{n}.0",
        "body": "Confirming the deployment window for version {n}.0 is Thursday 10pm-2am EST. Rollback plan is documented in Confluence. Please confirm your availability for the deploy support rotation.",
        "folder": "work",
        "priority": True,
        "attachment": False,
        "hint": "deployment",
    },
    {
        "subject": "Feature Flag Configuration Update",
        "body": "I've enabled the following feature flags in staging: dark-mode-v2, new-onboarding-flow, experimental-search. Please run smoke tests and report any issues in the #qa channel.",
        "folder": "work",
        "priority": False,
        "attachment": False,
        "hint": "engineering",
    },
]

FINANCE_TEMPLATES: List[Dict] = [
    {
        "subject": "Invoice #{n} - Cloud Infrastructure Services",
        "body": "Please find attached Invoice #{n} for cloud infrastructure services for the period March 1-31. Total: $12,450.00. Payment terms: Net 30. Remit to the bank details on the invoice.",
        "folder": "finance",
        "priority": False,
        "attachment": True,
        "hint": "invoice",
    },
    {
        "subject": "Expense Report Approval Required - Q{q} Travel",
        "body": "Your expense report for Q{q} business travel ($3,287.50) has been submitted and requires your approval. Includes: flights, hotel, meals, and ground transportation. Please review in Concur.",
        "folder": "finance",
        "priority": False,
        "attachment": True,
        "hint": "expense",
    },
    {
        "subject": "Purchase Order #{n} Confirmation",
        "body": "Your purchase order #{n} for 50 developer licenses of JetBrains IntelliJ IDEA has been approved. Total cost: $7,450.00. Licenses will be activated within 2 business days.",
        "folder": "finance",
        "priority": False,
        "attachment": True,
        "hint": "purchase_order",
    },
    {
        "subject": "Q{q} Budget Variance Report",
        "body": "Attached is the Q{q} budget variance report for the Engineering department. We are currently 8% under budget on headcount and 12% over on infrastructure costs. Action items included.",
        "folder": "finance",
        "priority": True,
        "attachment": True,
        "hint": "budget",
    },
    {
        "subject": "Annual Subscription Renewal - GitHub Enterprise",
        "body": "Your GitHub Enterprise subscription is up for renewal on May 15. Current plan: 200 seats at $21/seat/month. Renewal quote attached. Please approve or request changes by April 30.",
        "folder": "finance",
        "priority": False,
        "attachment": True,
        "hint": "subscription",
    },
    {
        "subject": "Payroll Adjustment Notice",
        "body": "This is to confirm the salary adjustment effective April 1. Updated compensation details are available in Workday. Please review and contact HR if you have questions.",
        "folder": "finance",
        "priority": True,
        "attachment": False,
        "hint": "payroll",
    },
]

MEETING_TEMPLATES: List[Dict] = [
    {
        "subject": "Calendar: Weekly Team Standup - {day}",
        "body": "Recurring meeting: Weekly Team Standup\nTime: {day} 9:30 AM - 9:45 AM EST\nLocation: Zoom (link in calendar invite)\nAgenda: Yesterday's progress, today's plan, blockers",
        "folder": "meetings",
        "priority": False,
        "attachment": False,
        "hint": "meeting",
    },
    {
        "subject": "Meeting Invite: Architecture Review - {project}",
        "body": "You're invited to an architecture review for {project}.\nDate: Thursday 2:00 PM - 3:30 PM EST\nLocation: Conference Room B / Zoom\nPlease prepare a 10-minute overview of your component.",
        "folder": "meetings",
        "priority": False,
        "attachment": False,
        "hint": "meeting",
    },
    {
        "subject": "RSVP: All-Hands Meeting - Q{q} Kickoff",
        "body": "Please RSVP for the Q{q} All-Hands Meeting.\nDate: Monday, April 15 at 2:00 PM EST\nAgenda: Company updates, product roadmap, team celebrations\nCatering will be provided for in-office attendees.",
        "folder": "meetings",
        "priority": False,
        "attachment": False,
        "hint": "meeting",
    },
    {
        "subject": "1:1 Rescheduled to {day} at 3pm",
        "body": "Hey, I need to reschedule our 1:1 to {day} at 3pm. Same Zoom link. I'd like to discuss your career development goals and the upcoming project assignments.",
        "folder": "meetings",
        "priority": False,
        "attachment": False,
        "hint": "meeting",
    },
    {
        "subject": "Reminder: Design Sprint Workshop Tomorrow",
        "body": "Reminder that the Design Sprint Workshop is tomorrow from 10am-4pm in the Innovation Lab. Bring your laptop and Figma access. Lunch will be provided.",
        "folder": "meetings",
        "priority": True,
        "attachment": False,
        "hint": "meeting",
    },
    {
        "subject": "Meeting Notes: Product Roadmap Review",
        "body": "Attached are the meeting notes from today's Product Roadmap Review. Key decisions: Q3 focus on mobile app, postpone enterprise features to Q4, hire 2 more designers. Action items assigned in Asana.",
        "folder": "meetings",
        "priority": False,
        "attachment": True,
        "hint": "meeting_notes",
    },
]

SPAM_TEMPLATES: List[Dict] = [
    {
        "subject": "Congratulations! You've Won a $500 Gift Card!",
        "body": "CLICK HERE to claim your exclusive $500 Amazon Gift Card! Limited time offer - expires in 24 hours! No purchase necessary. *Terms and conditions apply.",
        "folder": "spam",
        "priority": False,
        "attachment": False,
        "hint": None,
    },
    {
        "subject": "Urgent: Your Account Has Been Compromised",
        "body": "We detected suspicious activity on your account. Click the link below to verify your identity immediately or your account will be suspended within 24 hours. http://totally-legit-security.ru/verify",
        "folder": "spam",
        "priority": False,
        "attachment": False,
        "hint": None,
    },
    {
        "subject": "Limited Time: 90% Off Premium Software Bundle",
        "body": "Get the complete Adobe Creative Suite, Microsoft Office, and 50+ premium tools for just $29.99! This exclusive deal won't last. Click now to download instantly!",
        "folder": "spam",
        "priority": False,
        "attachment": False,
        "hint": None,
    },
    {
        "subject": "You've been selected for an exclusive investment opportunity",
        "body": "Dear valued investor, We have a once-in-a-lifetime opportunity for guaranteed 500% returns in just 30 days. Reply now to secure your spot. Minimum investment: $1,000.",
        "folder": "spam",
        "priority": False,
        "attachment": False,
        "hint": None,
    },
    {
        "subject": "Free Bitcoin - Claim Your Reward Now",
        "body": "You've been airdropped 0.5 BTC! Connect your wallet at crypto-free-airdrop.xyz to claim. Offer expires in 12 hours. Don't miss out on this incredible opportunity!",
        "folder": "spam",
        "priority": False,
        "attachment": False,
        "hint": None,
    },
    {
        "subject": "SALE ALERT: Unbeatable Deals Inside!",
        "body": "Shop now and save up to 80% on designer brands! Free shipping on orders over $25. Use code SAVE80 at checkout. Unsubscribe: click here",
        "folder": "spam",
        "priority": False,
        "attachment": False,
        "hint": None,
    },
    {
        "subject": "Your prescription is ready - Order online",
        "body": "Canadian Pharmacy Online - Best prices on all medications. No prescription required. Fast discreet shipping. Order now and save 70%!",
        "folder": "spam",
        "priority": False,
        "attachment": False,
        "hint": None,
    },
    {
        "subject": "Make $5000/week working from home!",
        "body": "Thousands of people are already earning $5000+ per week with our simple online system. No experience needed! Start today and see results tomorrow. Click here to learn more.",
        "folder": "spam",
        "priority": False,
        "attachment": False,
        "hint": None,
    },
]

ARCHIVE_TEMPLATES: List[Dict] = [
    {
        "subject": "Your monthly newsletter - Tech Digest March",
        "body": "This month in tech: AI advances, new framework releases, industry trends. Read the full newsletter at our website. Manage your subscription preferences below.",
        "folder": "archive",
        "priority": False,
        "attachment": False,
        "hint": "newsletter",
    },
    {
        "subject": "Automated Report: Weekly Analytics Summary",
        "body": "Your weekly analytics summary is ready. Key metrics: 12,450 page views (+5%), 3,200 unique visitors (+8%), 2.3% conversion rate (-0.1%). Full report available in the dashboard.",
        "folder": "archive",
        "priority": False,
        "attachment": True,
        "hint": "report",
    },
    {
        "subject": "GitHub: New release published - project-lib v2.{n}.0",
        "body": "A new release has been published for project-lib. Version 2.{n}.0 includes bug fixes and performance improvements. See the changelog for details.",
        "folder": "archive",
        "priority": False,
        "attachment": False,
        "hint": "notification",
    },
    {
        "subject": "Your order #{n} has been delivered",
        "body": "Your Amazon order #{n} has been delivered. If you didn't receive your package, please contact customer support. We'd love your feedback - leave a review!",
        "folder": "archive",
        "priority": False,
        "attachment": False,
        "hint": "shipping",
    },
    {
        "subject": "Slack digest: 47 unread messages in #engineering",
        "body": "You have 47 unread messages in #engineering. Top threads: deployment pipeline discussion, new hire introductions, Friday social planning.",
        "folder": "archive",
        "priority": False,
        "attachment": False,
        "hint": "notification",
    },
]

# ── Ambiguous Templates ──────────────────────────────────────────────────────
# Emails that could legitimately go to 2+ folders

AMBIGUOUS_TEMPLATES: List[Dict] = [
    {
        "subject": "Invoice for Q{q} Tech Conference Registration",
        "body": "Attached is the invoice for the Q{q} Tech Conference registration fee ($1,200). This includes the 3-day conference pass and workshop access. Receipt for expense reimbursement attached.",
        "folder": "finance",
        "acceptable": ["finance", "work"],
        "priority": False,
        "attachment": True,
        "hint": "invoice",
    },
    {
        "subject": "Team Lunch Budget Approval - Friday",
        "body": "Can you approve the budget for Friday's team lunch? Estimated cost: $350 for 15 people at the Italian place downtown. I'll submit the expense report after.",
        "folder": "meetings",
        "acceptable": ["meetings", "finance"],
        "priority": False,
        "attachment": False,
        "hint": "meeting",
    },
    {
        "subject": "Contract Review: SaaS Vendor Agreement",
        "body": "Legal has flagged several items in the SaaS vendor agreement that need engineering input. Could you review sections 4.2 (data handling) and 5.1 (SLA terms) and provide your technical assessment?",
        "folder": "work",
        "acceptable": ["work", "finance"],
        "priority": True,
        "attachment": True,
        "hint": "contract",
    },
    {
        "subject": "Reminder: Training Session - Cloud Security Best Practices",
        "body": "Reminder: The mandatory Cloud Security Best Practices training is scheduled for Wednesday 2pm. This is required for SOC 2 compliance. Please complete the pre-work in the shared drive.",
        "folder": "meetings",
        "acceptable": ["meetings", "work"],
        "priority": True,
        "attachment": False,
        "hint": "training",
    },
    {
        "subject": "Re: Hardware Upgrade Request - Development Laptops",
        "body": "IT has approved the hardware upgrade request. 5 new MacBook Pro M3 Max units will be ordered. Total cost: $17,495. Please confirm the shipping addresses for each team member.",
        "folder": "work",
        "acceptable": ["work", "finance"],
        "priority": False,
        "attachment": False,
        "hint": "procurement",
    },
    {
        "subject": "Planning Offsite: Q{q} Strategy Session",
        "body": "We're organizing a Q{q} strategy offsite at the lake house. Dates: May 10-12. Budget: $8,000. Agenda includes roadmap planning, team building, and retrospective. RSVP by Friday.",
        "folder": "meetings",
        "acceptable": ["meetings", "work", "finance"],
        "priority": False,
        "attachment": True,
        "hint": "planning",
    },
]

# ── Sender Pools ─────────────────────────────────────────────────────────────

INTERNAL_SENDERS = [
    ("Alex Kim", "alex.kim@acmecorp.com", "acmecorp.com"),
    ("Jordan Taylor", "j.taylor@acmecorp.com", "acmecorp.com"),
    ("Morgan Lee", "morgan.lee@acmecorp.com", "acmecorp.com"),
    ("Casey Wilson", "casey.w@acmecorp.com", "acmecorp.com"),
    ("Riley Johnson", "riley.j@acmecorp.com", "acmecorp.com"),
    ("Avery Brown", "a.brown@acmecorp.com", "acmecorp.com"),
    ("Dakota Smith", "d.smith@acmecorp.com", "acmecorp.com"),
]

EXTERNAL_SENDERS = [
    ("Vendor Support", "support@cloudservices.io", "cloudservices.io"),
    ("AWS Billing", "billing@amazonaws.com", "amazonaws.com"),
    ("GitHub", "noreply@github.com", "github.com"),
    ("Jira Notifications", "jira@atlassian.net", "atlassian.net"),
    ("Slack", "notifications@slack.com", "slack.com"),
    ("DocuSign", "dse@docusign.net", "docusign.net"),
    ("Zoom", "no-reply@zoom.us", "zoom.us"),
]

SPAM_SENDERS = [
    ("Prize Center", "winner@prize-claims.ru", "prize-claims.ru"),
    ("Security Alert", "alert@secure-verify.xyz", "secure-verify.xyz"),
    ("Deal Finder", "deals@mega-savings.biz", "mega-savings.biz"),
    ("Crypto Rewards", "airdrop@free-bitcoin.club", "free-bitcoin.club"),
    ("Easy Money", "info@work-from-home-now.net", "work-from-home-now.net"),
    ("Pharma Online", "order@canada-pharmacy.top", "canada-pharmacy.top"),
    ("Best Deals", "sales@unbeatable-offers.shop", "unbeatable-offers.shop"),
]

PROJECT_NAMES = [
    "Phoenix", "Horizon", "Atlas", "Nebula", "Catalyst",
    "Vertex", "Quantum", "Apex", "Prism", "Zenith",
]

DAYS = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday"]


# ── Generator Functions ──────────────────────────────────────────────────────

def _fill_template(template: Dict, rng: random.Random) -> Dict:
    """Fill in template placeholders with random values."""
    n = rng.randint(100, 9999)
    q = rng.randint(1, 4)
    project = rng.choice(PROJECT_NAMES)
    day = rng.choice(DAYS)

    subject = template["subject"].format(n=n, q=q, project=project, day=day)
    body = template["body"].format(n=n, q=q, project=project, day=day)

    return {
        "subject": subject,
        "body": body,
        "folder": template["folder"],
        "priority": template["priority"],
        "attachment": template["attachment"],
        "hint": template.get("hint"),
        "acceptable": template.get("acceptable", []),
    }


def _pick_sender(
    folder: str,
    is_vip: bool,
    rng: random.Random,
) -> Tuple[str, str, bool]:
    """Choose an appropriate sender for the email category."""
    if is_vip:
        vip = rng.choice(VIP_SENDERS)
        return vip["email"], vip["domain"], True

    if folder == "spam":
        _, email, domain = rng.choice(SPAM_SENDERS)
        return email, domain, False

    if folder in ("archive",):
        _, email, domain = rng.choice(EXTERNAL_SENDERS)
        return email, domain, False

    # work, finance, meetings — mostly internal
    if rng.random() < 0.75:
        _, email, domain = rng.choice(INTERNAL_SENDERS)
    else:
        _, email, domain = rng.choice(EXTERNAL_SENDERS)
    return email, domain, False


def generate_realistic_emails(
    count: int,
    difficulty: str = "easy",
    seed: Optional[int] = None,
) -> List[Email]:
    """Generate a batch of realistic emails with ground-truth labels.

    Args:
        count: Number of emails to generate.
        difficulty: One of "easy", "medium", "hard".
            - easy:   Clear work vs. spam split, no ambiguity.
            - medium: 4 folders, ~15% ambiguous cases.
            - hard:   6 folders, ~25% ambiguous, VIP senders, urgency flags.
        seed: Random seed for reproducibility.

    Returns:
        List of Email objects sorted by timestamp (oldest first).

    Example:
        >>> emails = generate_realistic_emails(10, difficulty="medium", seed=42)
        >>> len(emails)
        10
        >>> all(e.ground_truth_folder in AVAILABLE_FOLDERS for e in emails)
        True
    """
    rng = random.Random(seed)
    base_time = datetime(2024, 4, 12, 8, 0, 0)
    emails: List[Email] = []

    # Decide template mix based on difficulty
    if difficulty == "easy":
        # Only work and spam
        templates_pool: List[Tuple[List[Dict], float]] = [
            (WORK_TEMPLATES, 0.5),
            (SPAM_TEMPLATES, 0.5),
        ]
        ambiguous_ratio = 0.0
        vip_ratio = 0.0
    elif difficulty == "medium":
        templates_pool = [
            (WORK_TEMPLATES, 0.30),
            (FINANCE_TEMPLATES, 0.25),
            (MEETING_TEMPLATES, 0.25),
            (SPAM_TEMPLATES, 0.20),
        ]
        ambiguous_ratio = 0.15
        vip_ratio = 0.10
    else:  # hard
        templates_pool = [
            (WORK_TEMPLATES, 0.25),
            (FINANCE_TEMPLATES, 0.20),
            (MEETING_TEMPLATES, 0.15),
            (SPAM_TEMPLATES, 0.15),
            (ARCHIVE_TEMPLATES, 0.15),
            (AMBIGUOUS_TEMPLATES, 0.10),
        ]
        ambiguous_ratio = 0.25
        vip_ratio = 0.20

    # Pre-compute the number of ambiguous and VIP emails
    num_ambiguous = max(1, int(count * ambiguous_ratio)) if ambiguous_ratio > 0 else 0
    num_vip = max(1, int(count * vip_ratio)) if vip_ratio > 0 else 0

    # Generate indices for special emails
    all_indices = list(range(count))
    rng.shuffle(all_indices)
    ambiguous_indices = set(all_indices[:num_ambiguous])
    vip_indices = set(all_indices[num_ambiguous:num_ambiguous + num_vip])

    for i in range(count):
        is_ambiguous = i in ambiguous_indices
        is_vip = i in vip_indices

        # Choose template
        if is_ambiguous and AMBIGUOUS_TEMPLATES:
            template = rng.choice(AMBIGUOUS_TEMPLATES)
        else:
            # Weighted random template selection
            r = rng.random()
            cumulative = 0.0
            chosen_pool = WORK_TEMPLATES  # fallback
            for pool, weight in templates_pool:
                cumulative += weight
                if r < cumulative:
                    chosen_pool = pool
                    break
            template = rng.choice(chosen_pool)

        filled = _fill_template(template, rng)
        sender, domain, sender_is_vip = _pick_sender(
            filled["folder"], is_vip, rng
        )

        timestamp = base_time + timedelta(
            hours=rng.randint(0, 8),
            minutes=rng.randint(0, 59),
            seconds=i,  # ensure unique ordering
        )

        acceptable = filled.get("acceptable", [])
        email_is_ambiguous = bool(acceptable) and len(acceptable) > 1

        email = Email(
            id=i,
            subject=filled["subject"],
            sender=sender,
            sender_domain=domain,
            timestamp=timestamp,
            body_preview=filled["body"][:300],
            priority_flag=filled["priority"] or is_vip,
            has_attachment=filled["attachment"],
            is_vip_sender=sender_is_vip or is_vip,
            category_hint=filled["hint"] if difficulty != "hard" else None,
            ground_truth_folder=filled["folder"],
            is_ambiguous=email_is_ambiguous,
            acceptable_folders=acceptable if email_is_ambiguous else [filled["folder"]],
        )
        emails.append(email)

    # Sort by timestamp for realistic inbox ordering
    emails.sort(key=lambda e: e.timestamp)

    # Re-assign IDs after sorting
    for idx, email in enumerate(emails):
        email.id = idx

    logger.info(
        "Generated %d emails (difficulty=%s, ambiguous=%d, vip=%d)",
        len(emails), difficulty, num_ambiguous, num_vip,
    )
    return emails