src/data_loader.py

"""
Dataset generation, curation, and split management for 3-class email threat classification:
- Legitimate (Class 0)
- Spam (Class 1)
- Suspicious (Class 2)
"""

import os
import random
import pandas as pd
import numpy as np
from typing import Tuple, Dict, Any
from sklearn.model_selection import train_test_split


CLASS_NAMES = ["Legitimate", "Spam", "Suspicious"]
CLASS_TO_ID = {"Legitimate": 0, "Spam": 1, "Suspicious": 2}
ID_TO_CLASS = {0: "Legitimate", 1: "Spam", 2: "Suspicious"}


# Diverse seed pools representing real-world email threat telemetry
LEGITIMATE_TEMPLATES = [
    (
        "Quarterly Q3 Financial Review Meeting Notes",
        "Hi Team,\n\nPlease find the summary of our Q3 financial review meeting attached below. Key highlights include an 18% increase in operational efficiency across cloud services and steady revenue growth in APAC.\n\nAction items for next sprint:\n1. Dave to finalize infrastructure budget by Thursday.\n2. Sarah to review AWS reserved instance utilization.\n\nLet me know if anyone has questions before our next sync.\n\nBest regards,\nMarcus Vance\nDirector of Operations"
    ),
    (
        "Sprint 24 Retrospective & Planning Agenda",
        "Hello everyone,\n\nOur sprint 24 retrospective will take place tomorrow at 10:00 AM in Conference Room B and via Google Meet. Please review the Jira backlog board prior to the call and update your assigned story points.\n\nTopics to cover:\n- Feature release velocity\n- Resolving CI/CD pipeline build delays\n- ML model evaluation benchmark results\n\nThanks,\nElena Rostova\nScrum Master"
    ),
    (
        "Code Review: PR #412 - Implement Caching Layer in API Gateway",
        "Hey Alex,\n\nI just reviewed your pull request regarding Redis caching. Overall the logic looks solid and tests pass cleanly. I left two minor comments regarding TTL expiration edge cases when the cache key contains special characters.\n\nOnce those are addressed, feel free to merge into staging.\n\nCheers,\nKenji Sato\nStaff Software Engineer"
    ),
    (
        "HR Notice: Upcoming Open Enrollment for Health Benefits",
        "Dear Colleagues,\n\nThis is a reminder that the annual open enrollment period for health, dental, and vision insurance begins on November 1st and closes on November 15th. You can review plan options in the internal Workday portal.\n\nInformational webinars will be held next Tuesday and Thursday. No changes will occur automatically, so please make your selections on time.\n\nWarm regards,\nPeople & Culture Department"
    ),
    (
        "Client Project Update: Phase 2 Milestone Delivery",
        "Hi Samantha,\n\nWe have completed Phase 2 testing for the client portal redesign. The user acceptance testing results were positive with a 94% task completion rate among beta participants.\n\nAttached is the milestone completion sign-off document for your review. We look forward to kicking off Phase 3 next Monday.\n\nSincerely,\nDavid Miller\nAccount Director"
    ),
    (
        "Customer Support Ticket #89211 - Resolved",
        "Hello Arthur,\n\nThank you for reaching out to enterprise support. We investigated the database connection timeout issue you reported on server node-04. The connection pool limit has been adjusted, and monitoring shows normal latency over the past 6 hours.\n\nPlease reply directly to this thread if you experience any further anomalies.\n\nBest,\nSupport Engineering Team"
    ),
    (
        "Lunch and Learn Session: Best Practices in MLflow Experimentation",
        "Hi All,\n\nJoin us this Friday at 12:30 PM for our monthly engineering Lunch and Learn. Jordan will present 'Scaling MLflow Model Registries and Tracking Pipelines'. Pizza and refreshments will be provided in the 4th floor lounge.\n\nCalendar invite has been sent out.\n\nThanks,\nEngineering Community of Practice"
    ),
    (
        "Weekly Cloud Infrastructure Utilization Report",
        "Hi DevOps Team,\n\nAttached is the weekly resource utilization summary across our Kubernetes clusters. CPU utilization averaged 64%, and memory remained within nominal thresholds. Storage provisioning in us-east-1 requires slight optimization before end of month.\n\nFull Grafana dashboards are accessible via our internal VPN.\n\nRegards,\nCloud Platform Reliability Team"
    ),
    (
        "Invitation to Annual Tech Summit 2026",
        "Dear Abdul,\n\nYou have been nominated to attend the 2026 Enterprise Architecture & AI Summit in Austin, Texas. Your travel registration link and hotel booking guidelines have been approved by departmental budget.\n\nPlease confirm your availability by Friday so we can finalize travel accommodations.\n\nBest,\nConference Coordination Desk"
    ),
    (
        "Follow-up on Vendor Contract Agreement",
        "Good morning Rachel,\n\nThank you for the productive discussion yesterday. Our legal team reviewed the master service agreement redlines and accepted clauses 4 through 7. We have slight revisions regarding SLA terms in section 9.\n\nLet me know your availability for a 15-minute sync tomorrow afternoon.\n\nBest,\nJonathan Blake\nProcurement Specialist"
    )
]

SPAM_TEMPLATES = [
    (
        "EXCLUSIVE OFFER: Boost Your Website Traffic by 500% Overnight!",
        "Struggling to rank #1 on Google? Look no further!\n\nOur automated SEO backlink blast will drive 50,000 guaranteed unique visitors to your website in just 48 hours. We use proprietary high-authority PBN links to outrank all your competitors.\n\nSpecial introductory price: Only $49.99! Click here to claim your discount coupon now: http://cheap-seo-rankings-fast.biz/promo\n\nAct now! Offer expires at midnight!"
    ),
    (
        "CONGRATULATIONS! You have won $1,000,000 in the International Euro Sweepstakes!",
        "Dear Winner,\n\nYour email address was randomly selected from over 10 million online entries to win the First Prize in our International Cash Bonanza! To claim your $1,000,000 prize payout, simply send your full name, age, phone number, and mailing address to claimdesk@euro-lotto-awards-winner.info.\n\nDon't let your prize expire! Reply today!"
    ),
    (
        "Miracle Weight Loss Formula: Melt 30 lbs of Belly Fat in 14 Days!",
        "Doctors are stunned! This secret all-natural botanical extract burns stubborn abdominal fat without diet or exercise.\n\nOrder today and receive 3 bottles FREE plus free overnight shipping. Over 100,000 satisfied customers across the nation!\n\nOrder now at: http://pure-keto-miracle-melt.com/special-discount\n\n100% money back guarantee!"
    ),
    (
        "Pre-Approved Business Loan: Get $250,000 in Your Account by 5 PM",
        "Dear Business Owner,\n\nYour company has been pre-approved for an unsecured working capital line of credit up to $250,000 with zero collateral and 0% APR for the first 12 months. No credit check required!\n\nApply in 60 seconds at: http://fast-capital-funding-express.net/apply-now\n\nFunds deposited same business day. Terms apply."
    ),
    (
        "Wholesale Luxury Watches - Rolex, Omega, Breitling at 90% OFF!",
        "Looking for high-end luxury timepieces at unbeatable clearance prices?\n\nOur authentic replica inventory must be cleared out this weekend! Swiss movement, sapphire crystal, water-resistant. Free worldwide delivery on all orders over $99.\n\nShop our catalog now: http://timepiece-outlet-luxury-clearance.cc\n\nUse code LUX90 at checkout!"
    ),
    (
        "Crypto Trading Bot: Generate 300% Monthly Passive Yield Guaranteed!",
        "Tired of working a 9-5 job? Let artificial intelligence trade Bitcoin, Ethereum, and Solana for you while you sleep!\n\nOur proprietary neural arbitrage algorithm has produced a 98.7% win rate over the past 3 years. Start with as little as $100 today and watch your portfolio multiply.\n\nSign up today at http://crypto-ai-profit-generator.xyz\n\nSpaces limited to first 50 signups!"
    ),
    (
        "Hot Singles in Your Area Are Looking to Connect Tonight!",
        "Hey there handsome,\n\nWe found 4 matches matching your profile within 5 miles of your location! Profiles are 100% verified and free to message. Don't spend another lonely evening.\n\nView photos and chat now: http://local-dating-matches-finder.top/meet\n\nUnsubscribe link included below."
    ),
    (
        "Urgent Notice: Your Domain Registration is Expiring Soon - Renew Now!",
        "Attention Domain Administrator,\n\nYour domain name listing has reached its renewal term. Failure to renew within 48 hours will cause your domain to be dropped and acquired by third-party bidders.\n\nRenew your listing today for only $89/year at: http://domain-registry-services-renewal.org/pay\n\nProtect your brand name immediately."
    ),
    (
        "Refinance Your Mortgage Today: Rates Dropped to Historic Lows!",
        "Homeowners are saving an average of $750/month with new federal refinancing guidelines. Check your estimated monthly savings in under 2 minutes with no credit hit!\n\nLock in fixed rates starting at 2.9% APR today.\n\nGet your free quote: http://rate-drop-mortgage-quote.online/refi\n\nLimited time federal incentive program."
    ),
    (
        "Massive Warehouse Clearance: 80% Off All Electronics and Gadgets",
        "Overstocked brand-name inventory must go! Smart TVs, noise-canceling headphones, gaming consoles, and tablet computers priced below cost.\n\nFirst come, first served while inventory lasts.\n\nBrowse catalog: http://liquidation-tech-warehouse-deals.biz\n\nFree shipping on all orders today only!"
    )
]

SUSPICIOUS_TEMPLATES = [
    (
        "URGENT: Microsoft 365 Account Suspension Notice - Action Required",
        "SECURITY ALERT: We detected unauthorized login attempts to your Office 365 mailbox from IP 185.220.101.5 (Moscow, Russia).\n\nAs a security precaution, your email access and OneDrive cloud sync will be permanently TERMINATED within 24 hours unless you re-verify your employee identity.\n\nPlease click here immediately to verify your credentials and restore access: http://192.168.1.105/auth/m365/login?user=account_verification\n\nFailure to authenticate will result in account lockout by your IT administrator.\n\nOffice 365 Security Operations Desk"
    ),
    (
        "URGENT: Confidential Wire Transfer Request - CEO Authorization",
        "Abdul,\n\nAre you currently at your desk? I am currently in an urgent confidential acquisition meeting with external legal counsel and cannot take calls.\n\nWe need to process an immediate wire transfer of $48,500 to our acquisition escrow account before 3:00 PM EST to finalize the agreement. Please confirm your availability to execute this payment immediately and I will forward the beneficiary routing and SWIFT details.\n\nDo not discuss this with anyone on the floor as it is strictly subject to NDA.\n\nSent from my iPhone\nArthur Pendelton\nChief Executive Officer"
    ),
    (
        "Action Required: DocuSign - Please Review and Sign Confidential Termination Agreement",
        "DocuSign Security Envelope #90218-B\n\nHuman Resources has transmitted a time-sensitive confidential document regarding your 2026 Compensation and Stock Option Agreement for mandatory electronic signature.\n\nVIEW DOCUMENT: http://bit.ly/docusign-secure-document-sign-9912\n\nThis link is confidential and will expire in 2 hours. If you fail to review and sign before the deadline, your payroll direct deposit may be placed on temporary administrative hold.\n\nDocuSign Electronic Signature Service"
    ),
    (
        "IT Helpdesk Alert: Mandatory Password Expiration & Multi-Factor Authentication Reset",
        "Dear Employee,\n\nYour corporate domain network password is scheduled to expire in 4 hours as per our updated quarterly cybersecurity compliance policy.\n\nYou must update your network credentials and re-authenticate your Duo/MFA security tokens immediately to prevent interruption to internal VPN and email services.\n\nUpdate your credentials here: http://tinyurl.com/corp-helpdesk-mfa-update\n\nDo not ignore this notice. Unverified workstations will be quarantined from the corporate network.\n\nGlobal IT Helpdesk & Security Operations"
    ),
    (
        "Payroll Alert: Direct Deposit Account Change Notification - Immediate Action Required",
        "Notice of Payroll Account Modification:\n\nYour direct deposit bank routing information was modified via Employee Self-Service on 10/11/2026. If you did NOT request this banking change, your scheduled salary payment will be deposited into the new Wells Fargo account ending in *8821.\n\nIf this was unauthorized, cancel the pending transfer and secure your payroll credentials immediately at: http://company-payroll-portal-auth.link/cancel-fraud\n\nHuman Resources & Payroll Compliance"
    ),
    (
        "Invoice Overdue #INV-99214: Legal Action Notice & Attached Receipt",
        "Attention Accounts Payable,\n\nOur records show invoice INV-99214 totaling $14,250.00 is now 45 days past due. Continued non-payment will result in immediate escalation to third-party collection agencies and legal garnishment proceedings.\n\nPlease find the detailed ledger and wire transfer instructions enclosed in the attached archive: Invoice_Details_Overdue.zip\n\nExecute the overdue wire payment to our routing account within 24 hours to halt legal action.\n\nAccounts Receivable Department"
    ),
    (
        "PayPal Security Notification: Unauthorized Transaction of $1,849.00 Detected",
        "Dear Customer,\n\nWe observed suspicious activity on your account. An unauthorized transaction for Apple Store online in the amount of $1,849.00 USD was initiated from an unrecognized device in Frankfurt, Germany.\n\nIf you did not make this purchase, dispute the transaction and reverse payment within 60 minutes: http://paypal-resolution-center-fraud-dispute.top/login\n\nYour account has been temporarily restricted for security protection.\n\nPayPal Anti-Fraud Division"
    ),
    (
        "Google Workspace Alert: Critical Vulnerability Detected on Your Mailbox",
        "Admin Alert:\n\nA severe zero-day exploit has been detected targeting your Google Workspace mailbox. Your session tokens have been invalidated.\n\nTo restore email sync and download your security patch, sign in to your enterprise administrator portal immediately:\nhttp://bit.ly/google-admin-mfa-security-verify\n\nFailure to authenticate within 30 minutes will result in automated account quarantine.\n\nGoogle Cloud Security Team"
    ),
    (
        "DHL Express Delivery: Package Delivery Exception - Address Verification Required",
        "Tracking ID: #DHL-889104-US\n\nYour international package could not be delivered due to an invalid recipient delivery address and unpaid customs clearance duty ($4.95).\n\nPlease confirm your residential address and submit credit card billing verification at: http://dhl-express-parcel-tracking-update.com/verify-shipment\n\nIf unclaimed within 48 hours, package will be returned to the origin sender or destroyed.\n\nDHL Express Courier Services"
    ),
    (
        "Internal Revenue Service: Final Demand Notice - Tax Audit Penalty Enforcement",
        "NOTICE OF INTENT TO LEVY:\n\nThe Internal Revenue Service has issued a final warrant for outstanding tax penalties totaling $6,420.00 under case file #IRS-9812-TAX. Failure to resolve this balance immediately will lead to bank account asset freezing and local federal sheriff summons.\n\nVerify your taxpayer identification and submit emergency settlement at: http://irs-tax-audit-resolution-gov.info/case-login\n\nRespond within 12 hours of receipt."
    ),
    (
        "URGENT: Microsoft 365 Account Suspension Notice - Action Required",
        "We detected unauthorized login attempts from IP 185.220.101.5. Your account will be permanently terminated within 24 hours unless you verify your employee identity at http://192.168.1.105/auth/login"
    ),
    (
        "URGENT: Confidential Wire Transfer Request - CEO Authorization",
        "Abdul, I am in an urgent acquisition meeting. We need to process an immediate wire transfer of $48,500 before 3 PM. Confirm receipt and I will send routing numbers."
    ),
    (
        "Security Alert: Mandatory Password Reset Required",
        "Your password will expire today. Enter credentials at http://tinyurl.com/reset-pw to unlock account."
    ),
    (
        "DocuSign: Electronic Signature Requested on Payroll Agreement",
        "Please sign document http://bit.ly/docusign-payroll-update immediately to ensure direct deposit."
    )
]


def expand_templates(templates: list, target_count: int, class_name: str) -> list:
    """
    Expands base seed templates with realistic contextual variations,
    subject line perturbations, timestamps, references, and greetings
    to build a high-fidelity synthetic benchmark corpus.
    """
    perturbations = [
        " Please review the attached details.",
        " Confirm receipt of this message as soon as possible.",
        " All team members must adhere to this protocol.",
        " Note that this is time sensitive.",
        " Let us know if you need any adjustments.",
        " Reference ticket #ID-{}.",
        " Timestamp verification token: [AUTH-{}]"
    ]

    generated = []
    base_len = len(templates)
    
    for i in range(target_count):
        base_subject, base_body = templates[i % base_len]
        rand_token = random.randint(10000, 99999)
        pert = perturbations[i % len(perturbations)].format(rand_token)
        
        # Introduce variation
        if i >= base_len:
            if class_name == "Suspicious" and i % 3 == 0:
                subj = f"[ACTION REQUIRED] {base_subject} (Ref: #{rand_token})"
            elif class_name == "Spam" and i % 2 == 0:
                subj = f"*** SPECIAL NOTICE *** {base_subject} [#{rand_token}]"
            else:
                subj = f"{base_subject} - Update #{rand_token}"
            
            # Diverse lengths: some concise, some perturbed
            if i % 5 == 0:
                body = base_body  # Short concise variant
            else:
                body = f"{base_body}\n\n{pert}"
        else:
            subj = base_subject
            body = base_body

        generated.append({
            "subject": subj,
            "body": body,
            "label": class_name,
            "label_id": CLASS_TO_ID[class_name]
        })

    return generated


def generate_curated_dataset(samples_per_class: int = 1200, random_seed: int = 42) -> pd.DataFrame:
    """
    Generates a curated 3-class dataset containing legitimate, spam,
    and suspicious threat emails with balanced distribution.
    """
    random.seed(random_seed)
    np.random.seed(random_seed)

    legit_samples = expand_templates(LEGITIMATE_TEMPLATES, samples_per_class, "Legitimate")
    spam_samples = expand_templates(SPAM_TEMPLATES, samples_per_class, "Spam")
    susp_samples = expand_templates(SUSPICIOUS_TEMPLATES, samples_per_class, "Suspicious")

    all_samples = legit_samples + spam_samples + susp_samples
    df = pd.DataFrame(all_samples)
    
    # Shuffle thoroughly
    df = df.sample(frac=1.0, random_state=random_seed).reset_index(drop=True)
    return df


def split_and_save_data(
    df: pd.DataFrame,
    output_dir: str,
    train_ratio: float = 0.70,
    val_ratio: float = 0.15,
    test_ratio: float = 0.15,
    random_seed: int = 42
) -> Dict[str, pd.DataFrame]:
    """
    Performs stratified train/val/test splits preventing data leakage
    and saves resulting CSV datasets.
    """
    os.makedirs(output_dir, exist_ok=True)

    # First split: train vs (val + test)
    val_test_ratio = val_ratio + test_ratio
    train_df, val_test_df = train_test_split(
        df,
        test_size=val_test_ratio,
        random_state=random_seed,
        stratify=df["label_id"]
    )

    # Second split: val vs test
    val_proportion = val_ratio / val_test_ratio
    val_df, test_df = train_test_split(
        val_test_df,
        test_size=(1.0 - val_proportion),
        random_state=random_seed,
        stratify=val_test_df["label_id"]
    )

    # Save to disk
    train_path = os.path.join(output_dir, "train.csv")
    val_path = os.path.join(output_dir, "val.csv")
    test_path = os.path.join(output_dir, "test.csv")

    train_df.to_csv(train_path, index=False)
    val_df.to_csv(val_path, index=False)
    test_df.to_csv(test_path, index=False)

    print(f"Stratified data split complete:")
    print(f"  Training set:   {len(train_df)} samples -> {train_path}")
    print(f"  Validation set: {len(val_df)} samples -> {val_path}")
    print(f"  Test set:       {len(test_df)} samples -> {test_path}")

    return {
        "train": train_df,
        "val": val_df,
        "test": test_df
    }


def load_datasets(processed_dir: str) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """
    Loads train, val, and test splits from processed data directory.
    """
    train_path = os.path.join(processed_dir, "train.csv")
    val_path = os.path.join(processed_dir, "val.csv")
    test_path = os.path.join(processed_dir, "test.csv")

    if not all(os.path.exists(p) for p in [train_path, val_path, test_path]):
        raise FileNotFoundError(f"Processed splits not found in {processed_dir}. Run prepare_data.py first.")

    train_df = pd.read_csv(train_path)
    val_df = pd.read_csv(val_path)
    test_df = pd.read_csv(test_path)

    return train_df, val_df, test_df