task_generator.py

"""Seeded task generator for Supergames-style OpenEnv tasks."""

from __future__ import annotations

from dataclasses import dataclass
from random import Random
from typing import List, Tuple

try:
    from .models import Game, GameTitle, Severity, StaffPool, WorkItem, WorkType
except ImportError:
    from models import Game, GameTitle, Severity, StaffPool, WorkItem, WorkType


TaskTuple = Tuple[List[Game], List[WorkItem], StaffPool, int, str]


@dataclass(frozen=True)
class DifficultyPreset:
    game_count: int
    bug_count: int
    feature_count: int
    staff_range: Tuple[int, int]
    sprint_count: int
    crisis: bool = False


PRESETS = {
    "easy": DifficultyPreset(
        game_count=1,
        bug_count=3,
        feature_count=1,
        staff_range=(55, 75),
        sprint_count=1,
    ),
    "medium": DifficultyPreset(
        game_count=4,
        bug_count=4,
        feature_count=4,
        staff_range=(75, 100),
        sprint_count=3,
    ),
    "hard": DifficultyPreset(
        game_count=4,
        bug_count=5,
        feature_count=5,
        staff_range=(130, 160),
        sprint_count=5,
    ),
    "crisis": DifficultyPreset(
        game_count=4,
        bug_count=2,
        feature_count=2,
        staff_range=(90, 115),
        sprint_count=3,
        crisis=True,
    ),
}


GAME_TEMPLATES = [
    {
        "id": "mmo",
        "title": GameTitle.MMO,
        "branches": ["Mumbai", "Bangalore", "Pune"],
        "monthly_revenue": (1_000_000.0, 1_500_000.0),
        "revenue_potential": (2_200_000.0, 3_000_000.0),
        "active_players": (260_000, 340_000),
        "churn_rate": (0.03, 0.05),
    },
    {
        "id": "shooter",
        "title": GameTitle.SHOOTER,
        "branches": ["Chennai", "Hyderabad", "Delhi"],
        "monthly_revenue": (700_000.0, 1_000_000.0),
        "revenue_potential": (1_300_000.0, 1_800_000.0),
        "active_players": (140_000, 200_000),
        "churn_rate": (0.04, 0.06),
    },
    {
        "id": "strat",
        "title": GameTitle.STRAT,
        "branches": ["Bangalore", "Kolkata", "Noida"],
        "monthly_revenue": (450_000.0, 700_000.0),
        "revenue_potential": (1_600_000.0, 2_300_000.0),
        "active_players": (85_000, 125_000),
        "churn_rate": (0.05, 0.07),
    },
    {
        "id": "fighter",
        "title": GameTitle.FIGHTER,
        "branches": ["Hyderabad", "Mumbai", "Gurgaon"],
        "monthly_revenue": (280_000.0, 420_000.0),
        "revenue_potential": (800_000.0, 1_200_000.0),
        "active_players": (55_000, 80_000),
        "churn_rate": (0.06, 0.08),
    },
]


BUG_TITLES = [
    "Payment gateway outage",
    "Login server instability",
    "Save file corruption",
    "Ranked mode desync",
    "Anti-cheat false bans",
    "Progression wipe on update",
    "Guild bank duplication glitch",
    "Crash during matchmaking",
    "Inventory rollback bug",
    "Session token leak",
]

FEATURE_TITLES = [
    "Seasonal content drop",
    "Ranked 2.0 system",
    "Multiplayer co-op mode",
    "Open world expansion",
    "Battle pass refresh",
    "Tournament bracket mode",
    "Crafting system overhaul",
    "New player onboarding revamp",
    "Cross-platform party finder",
    "Guild progression track",
]


def generate_task(difficulty: str = "medium", seed: int = 42) -> TaskTuple:
    """Generate a Supergames-style task tuple for ad hoc evaluation."""
    difficulty_key = difficulty.lower()
    if difficulty_key not in PRESETS:
        supported = ", ".join(sorted(PRESETS))
        raise ValueError(f"Unknown difficulty '{difficulty}'. Supported: {supported}")

    rng = Random(seed)
    preset = PRESETS[difficulty_key]
    games = _generate_games(rng, preset)
    work_queue = _generate_work_queue(rng, preset, games)
    staff_pool = StaffPool(total=rng.randint(*preset.staff_range))
    goal = _build_goal(difficulty_key, preset, games, staff_pool)
    return games, work_queue, staff_pool, preset.sprint_count, goal


def _generate_games(rng: Random, preset: DifficultyPreset) -> List[Game]:
    templates = GAME_TEMPLATES.copy()
    rng.shuffle(templates)
    selected = templates[: preset.game_count]

    return [
        Game(
            id=template["id"],
            title=template["title"],
            branch=rng.choice(template["branches"]),
            monthlyRevenue=_rounded_money(rng, template["monthly_revenue"], 10_000),
            revenuePotential=_rounded_money(rng, template["revenue_potential"], 10_000),
            activePlayers=rng.randrange(
                template["active_players"][0],
                template["active_players"][1] + 1,
                5_000,
            ),
            churnRate=round(rng.uniform(*template["churn_rate"]), 3),
        )
        for template in selected
    ]


def _generate_work_queue(
    rng: Random,
    preset: DifficultyPreset,
    games: List[Game],
) -> List[WorkItem]:
    work_queue: List[WorkItem] = []

    for index in range(1, preset.bug_count + 1):
        severity = _bug_severity(rng, index)
        work_queue.append(
            WorkItem(
                id=f"b{index}",
                gameId=rng.choice(games).id,
                workType=WorkType.BUG,
                title=_pick_title(rng, BUG_TITLES, index),
                severity=severity,
                effort=_bug_effort(rng, severity),
                revenueImpact=_bug_revenue_impact(rng, severity),
                impactDelay=0 if severity >= Severity.CRITICAL else rng.choice([0, 1]),
                churnReduction=_churn_reduction(severity),
            )
        )

    for index in range(1, preset.feature_count + 1):
        severity = rng.choice([Severity.MEDIUM, Severity.HIGH, Severity.HIGH])
        work_queue.append(
            WorkItem(
                id=f"f{index}",
                gameId=rng.choice(games).id,
                workType=WorkType.FEATURE,
                title=_pick_title(rng, FEATURE_TITLES, index),
                severity=severity,
                effort=rng.randrange(260, 801, 20),
                revenueImpact=round(rng.uniform(90.0, 420.0), 1),
                impactDelay=rng.choice([1, 1, 2]),
                churnReduction=0.0,
            )
        )

    if preset.crisis:
        crisis_game = rng.choice(games)
        work_queue.append(
            WorkItem(
                id="crisis-1",
                gameId=crisis_game.id,
                workType=WorkType.BUG,
                title=f"CRISIS: {crisis_game.title.value} player data exposed",
                severity=Severity.BLOCKER,
                effort=rng.randrange(320, 421, 10),
                revenueImpact=round(rng.uniform(500.0, 700.0), 1),
                impactDelay=0,
                churnReduction=0.5,
                crisis=True,
            )
        )

    return work_queue


def _build_goal(
    difficulty: str,
    preset: DifficultyPreset,
    games: List[Game],
    staff_pool: StaffPool,
) -> str:
    title_count = "title" if len(games) == 1 else "titles"
    sprint_count = "sprint" if preset.sprint_count == 1 else "sprints"
    game_names = ", ".join(game.title.value for game in games)

    if preset.crisis:
        return (
            f"You manage {game_names}. You have {staff_pool.total} staff and "
            f"{preset.sprint_count} {sprint_count}. Sprint 1 is normal, but a "
            "crisis blocker becomes active in sprint 2. Prioritise urgent bugs, "
            "balance feature payoff, and maximise total revenue while containing "
            "the crisis."
        )

    if difficulty == "easy":
        return (
            f"You are the engineering lead for {game_names}. You have "
            f"{staff_pool.total} staff and {preset.sprint_count} sprint. Assign "
            "staff to the highest-impact bugs and features to maximise revenue."
        )

    return (
        f"You manage engineering across {len(games)} Super Games {title_count}: "
        f"{game_names}. You have {staff_pool.total} staff and "
        f"{preset.sprint_count} {sprint_count}. Allocate staff to bugs and "
        "features to maximise total revenue. Unresolved critical and blocker "
        "bugs increase churn each sprint, while features take longer but can "
        "unlock higher long-term payoff."
    )


def _bug_severity(rng: Random, index: int) -> Severity:
    if index == 1:
        return Severity.BLOCKER
    if index == 2:
        return Severity.CRITICAL
    return rng.choice(
        [
            Severity.MEDIUM,
            Severity.HIGH,
            Severity.HIGH,
            Severity.CRITICAL,
            Severity.BLOCKER,
        ]
    )


def _bug_effort(rng: Random, severity: Severity) -> int:
    if severity == Severity.BLOCKER:
        return rng.randrange(260, 381, 10)
    if severity == Severity.CRITICAL:
        return rng.randrange(180, 301, 10)
    if severity == Severity.HIGH:
        return rng.randrange(120, 221, 10)
    return rng.randrange(80, 161, 10)


def _bug_revenue_impact(rng: Random, severity: Severity) -> float:
    ranges = {
        Severity.BLOCKER: (250.0, 480.0),
        Severity.CRITICAL: (150.0, 300.0),
        Severity.HIGH: (70.0, 160.0),
        Severity.MEDIUM: (30.0, 90.0),
        Severity.LOW: (10.0, 40.0),
    }
    return round(rng.uniform(*ranges[severity]), 1)


def _churn_reduction(severity: Severity) -> float:
    reductions = {
        Severity.BLOCKER: 0.35,
        Severity.CRITICAL: 0.2,
        Severity.HIGH: 0.1,
        Severity.MEDIUM: 0.0,
        Severity.LOW: 0.0,
    }
    return reductions[severity]


def _pick_title(rng: Random, titles: List[str], index: int) -> str:
    if index <= len(titles):
        return rng.sample(titles, k=len(titles))[index - 1]
    return rng.choice(titles)


def _rounded_money(rng: Random, value_range: Tuple[float, float], step: int) -> float:
    low = int(value_range[0] // step)
    high = int(value_range[1] // step)
    return float(rng.randint(low, high) * step)


__all__ = ["PRESETS", "TaskTuple", "generate_task"]