env/tasks.py

"""
Task generation for AETHER-TaskFlow.

Generates realistic workflow tasks across three difficulty levels:
  easy   – stable, predictable, moderate resources
  medium – dynamic priorities, tighter deadlines
  hard   – scarce resources, high uncertainty, adversarial failures
"""

from __future__ import annotations

import random
from typing import List, Optional, Tuple

from models import TaskInfo


# ---------------------------------------------------------------------------
# Task templates – drawn from realistic enterprise/ops domains
# ---------------------------------------------------------------------------

_TASK_TEMPLATES: List[Tuple[str, str]] = [
    # (name_template, category)
    ("Email triage: {volume} messages", "communication"),
    ("Code review: PR #{pr_id}", "engineering"),
    ("Data pipeline: {dataset} ETL", "data"),
    ("Customer support ticket #{tid}", "support"),
    ("Security audit: {module} module", "security"),
    ("Performance optimization: {service}", "engineering"),
    ("Database backup: {db_name}", "infrastructure"),
    ("Report generation: {report_type}", "analytics"),
    ("Incident response: {severity} alert", "operations"),
    ("Content moderation: batch #{bid}", "moderation"),
    ("API rate-limit review: {api_name}", "infrastructure"),
    ("ML model retraining: {model_name}", "ml"),
    ("Budget reconciliation: {quarter}", "finance"),
    ("Compliance check: {regulation}", "legal"),
    ("System health scan: {region}", "operations"),
]

_FILL_VALUES: dict = {
    "volume": ["50", "120", "300", "500"],
    "pr_id": ["1042", "2381", "9001", "4417"],
    "dataset": ["sales_Q3", "user_events", "inventory", "logs_prod"],
    "tid": ["55123", "10984", "30021", "77654"],
    "module": ["auth", "payments", "admin", "reporting"],
    "service": ["checkout", "search", "recommendations", "notifications"],
    "db_name": ["prod_main", "analytics_dw", "user_db", "logs_archive"],
    "report_type": ["weekly_KPI", "SLA_breach", "revenue_forecast", "churn"],
    "severity": ["P1", "P2", "P3"],
    "bid": ["4401", "8812", "1123"],
    "api_name": ["stripe", "twilio", "sendgrid", "maps"],
    "model_name": ["churn_v3", "fraud_detector", "recommender_v2"],
    "quarter": ["Q3-2025", "Q4-2025", "Q1-2026"],
    "regulation": ["GDPR", "SOC2", "HIPAA", "PCI-DSS"],
    "region": ["us-east-1", "eu-west-2", "ap-southeast-1"],
}


def _fill_template(template: str, rng: random.Random) -> str:
    result = template
    for key, choices in _FILL_VALUES.items():
        placeholder = "{" + key + "}"
        if placeholder in result:
            result = result.replace(placeholder, rng.choice(choices))
    return result


# ---------------------------------------------------------------------------
# Difficulty profiles
# ---------------------------------------------------------------------------

_PROFILES: dict = {
    "easy": {
        "n_tasks": 5,
        "priority_range": (0.4, 0.9),
        "deadline_range": (4, 8),
        "uncertainty_range": (0.05, 0.35),
        "value_range": (8.0, 20.0),
        "energy_cost_range": (0.5, 1.5),
        "budget_cost_range": (1.0, 5.0),
        "initial_time": 10,
        "initial_energy": 12.0,
        "initial_budget": 60.0,
        "max_steps": 10,
    },
    "medium": {
        "n_tasks": 8,
        "priority_range": (0.3, 1.0),
        "deadline_range": (2, 6),
        "uncertainty_range": (0.15, 0.65),
        "value_range": (5.0, 25.0),
        "energy_cost_range": (0.8, 2.5),
        "budget_cost_range": (2.0, 10.0),
        "initial_time": 10,
        "initial_energy": 10.0,
        "initial_budget": 50.0,
        "max_steps": 10,
    },
    "hard": {
        "n_tasks": 12,
        "priority_range": (0.2, 1.0),
        "deadline_range": (1, 4),
        "uncertainty_range": (0.35, 0.95),
        "value_range": (3.0, 30.0),
        "energy_cost_range": (1.2, 4.0),
        "budget_cost_range": (5.0, 20.0),
        "initial_time": 10,
        "initial_energy": 8.0,
        "initial_budget": 40.0,
        "max_steps": 10,
    },
}


def get_profile(difficulty: str) -> dict:
    return _PROFILES[difficulty]


def generate_tasks(difficulty: str, seed: Optional[int] = None) -> List[TaskInfo]:
    """Generate a task queue for the given difficulty level."""
    rng = random.Random(seed)
    profile = _PROFILES[difficulty]

    tasks: List[TaskInfo] = []
    used_templates = rng.choices(range(len(_TASK_TEMPLATES)), k=profile["n_tasks"])

    for i, t_idx in enumerate(used_templates):
        name_template, category = _TASK_TEMPLATES[t_idx]
        name = _fill_template(name_template, rng)

        priority = rng.uniform(*profile["priority_range"])
        deadline = rng.randint(*profile["deadline_range"])
        uncertainty = rng.uniform(*profile["uncertainty_range"])
        value = rng.uniform(*profile["value_range"])
        energy_cost = rng.uniform(*profile["energy_cost_range"])
        budget_cost = rng.uniform(*profile["budget_cost_range"])

        tasks.append(
            TaskInfo(
                task_id=i,
                name=name,
                priority=priority,
                deadline=deadline,
                uncertainty=uncertainty,
                value=value,
                required_energy=energy_cost,
                required_budget=budget_cost,
                category=category,
            )
        )

    # Sort by priority descending so agent sees most urgent first
    tasks.sort(key=lambda t: t.priority, reverse=True)
    # Re-index after sort
    for idx, t in enumerate(tasks):
        t.task_id = idx

    return tasks


def apply_dynamic_updates(
    tasks: List[TaskInfo],
    step: int,
    difficulty: str,
    rng: random.Random,
) -> List[TaskInfo]:
    """
    Apply stochastic dynamic updates to the task queue (medium/hard only).
    - Priority drift
    - Deadline tightening
    - Uncertainty spikes
    """
    if difficulty == "easy":
        return tasks

    for task in tasks:
        if task.status != "pending":
            continue

        # NOTE: deadline countdown is handled centrally in aether_env.py
        # _step_impl() to avoid double-decrement on hard mode.

        # Priority drift ±0.1
        drift = rng.uniform(-0.08, 0.12)
        task.priority = min(1.0, max(0.1, task.priority + drift))

        # Uncertainty spike (hard mode)
        if difficulty == "hard" and rng.random() < 0.15:
            task.uncertainty = min(0.95, task.uncertainty + rng.uniform(0.1, 0.25))

    return tasks