env/scenario_generator.py

"""
Curriculum-aware scenario generator.

Produces randomized but structured scenarios with configurable difficulty:
    - Easy:   3–5 tasks, few conflicts, no hidden items.
    - Medium: 5–8 tasks, some conflicts, some hidden items.
    - Hard:   8–12 tasks, many conflicts, many hidden items.
"""

import random
from typing import Optional

from env.state import State
from env.utils import (
    TIME_SLOTS,
    DURATIONS,
    MEETING_TITLES,
    WORK_TITLES,
    PERSONAL_TITLES,
    MESSAGE_CONTENTS_URGENT,
    MESSAGE_CONTENTS_NORMAL,
    SENDERS,
)


# ─── Difficulty Profiles ─────────────────────────────────────────────────────

DIFFICULTY_PROFILES = {
    "easy": {
        "task_range": (3, 5),
        "inbox_range": (2, 3),
        "conflict_probability": 0.1,
        "hidden_ratio": 0.0,
        "delayed_ratio": 0.0,
        "high_priority_ratio": 0.2,
        "urgent_message_ratio": 0.2,
        "duration_choices": [30],
    },
    "medium": {
        "task_range": (5, 8),
        "inbox_range": (3, 5),
        "conflict_probability": 0.3,
        "hidden_ratio": 0.2,
        "delayed_ratio": 0.2,
        "high_priority_ratio": 0.3,
        "urgent_message_ratio": 0.3,
        "duration_choices": [30, 60],
    },
    "hard": {
        "task_range": (8, 12),
        "inbox_range": (5, 8),
        "conflict_probability": 0.5,
        "hidden_ratio": 0.3,
        "delayed_ratio": 0.3,
        "high_priority_ratio": 0.4,
        "urgent_message_ratio": 0.4,
        "duration_choices": [30, 60, 90, 120],
    },
}


class ScenarioGenerator:
    """Generates diverse, curriculum-aware simulation scenarios."""

    def __init__(self, difficulty: str = "medium", seed: Optional[int] = None):
        """Initialize the generator.

        Args:
            difficulty: One of 'easy', 'medium', 'hard'.
            seed: Optional random seed for reproducibility.
        """
        self.difficulty = difficulty
        self.profile = DIFFICULTY_PROFILES.get(difficulty, DIFFICULTY_PROFILES["medium"])
        if seed is not None:
            random.seed(seed)

    def set_difficulty(self, difficulty: str):
        """Update the difficulty level (for curriculum learning)."""
        self.difficulty = difficulty
        self.profile = DIFFICULTY_PROFILES.get(difficulty, DIFFICULTY_PROFILES["medium"])

    def generate(self) -> State:
        """Generate a complete scenario with tasks, inbox, and optional hidden elements.

        Returns:
            A new State object.
        """
        profile = self.profile

        # Generate visible tasks
        num_tasks = random.randint(*profile["task_range"])
        all_tasks = self._generate_tasks(num_tasks, profile)

        # Split into visible and hidden
        hidden_count = int(len(all_tasks) * profile["hidden_ratio"])
        random.shuffle(all_tasks)
        visible_tasks = all_tasks[hidden_count:]
        hidden_tasks = all_tasks[:hidden_count]

        # Set reveal times for hidden tasks
        for ht in hidden_tasks:
            reveal_slot = random.choice(TIME_SLOTS[2:8])  # Reveal between 09:00–11:30
            ht["reveal_at"] = reveal_slot

        # Generate inbox messages
        num_messages = random.randint(*profile["inbox_range"])
        all_messages = self._generate_messages(num_messages, profile)

        # Split into immediate and delayed
        delayed_count = int(len(all_messages) * profile["delayed_ratio"])
        random.shuffle(all_messages)
        visible_messages = all_messages[delayed_count:]
        delayed_messages = all_messages[:delayed_count]

        # Set reveal times for delayed messages
        for dm in delayed_messages:
            reveal_slot = random.choice(TIME_SLOTS[2:10])
            dm["reveal_at"] = reveal_slot

        # Generate user preferences
        preferences = self._generate_preferences()

        return State(
            current_time="08:00",
            tasks=visible_tasks,
            inbox=visible_messages,
            preferences=preferences,
            hidden_tasks=hidden_tasks,
            delayed_inbox=delayed_messages,
        )

    def _generate_tasks(self, count: int, profile: dict) -> list:
        """Generate a list of task dicts."""
        tasks = []
        used_titles = set()

        for i in range(count):
            task_type = random.choices(
                ["meeting", "work", "personal"],
                weights=[0.4, 0.45, 0.15],
                k=1,
            )[0]

            # Pick title based on type
            title_pool = {
                "meeting": MEETING_TITLES,
                "work": WORK_TITLES,
                "personal": PERSONAL_TITLES,
            }[task_type]

            available = [t for t in title_pool if t not in used_titles]
            if not available:
                available = title_pool
            title = random.choice(available)
            used_titles.add(title)

            # Time slot — sometimes force conflicts
            if random.random() < profile["conflict_probability"] and tasks:
                # Reuse an existing task's time to create a conflict
                time_slot = random.choice(tasks)["time"]
            else:
                time_slot = random.choice(TIME_SLOTS[:14])  # 08:00–14:30 range

            # Priority
            if random.random() < profile["high_priority_ratio"]:
                priority = "high"
            else:
                priority = random.choice(["medium", "low"])

            # Duration
            duration = random.choice(profile["duration_choices"])

            tasks.append({
                "id": i,
                "title": title,
                "time": time_slot,
                "duration": duration,
                "priority": priority,
                "type": task_type,
                "status": "pending",
            })

        return tasks

    def _generate_messages(self, count: int, profile: dict) -> list:
        """Generate a list of inbox message dicts."""
        messages = []
        used_senders = set()

        for i in range(count):
            # Sender
            available_senders = [s for s in SENDERS if s not in used_senders]
            if not available_senders:
                available_senders = SENDERS
            sender = random.choice(available_senders)
            used_senders.add(sender)

            # Urgency
            if random.random() < profile["urgent_message_ratio"]:
                urgency = "high"
                content = random.choice(MESSAGE_CONTENTS_URGENT)
            else:
                urgency = random.choice(["medium", "low"])
                content = random.choice(MESSAGE_CONTENTS_NORMAL)

            messages.append({
                "id": i,
                "sender": sender,
                "content": content,
                "urgency": urgency,
                "replied": False,
            })

        return messages

    def _generate_preferences(self) -> dict:
        """Generate a randomized user preference profile."""
        # Pick 2–4 preferred meeting times
        preferred_times = random.sample(TIME_SLOTS[2:12], k=random.randint(2, 4))

        # Pick 1–2 focus hour blocks
        focus_start = random.choice(TIME_SLOTS[4:10])
        focus_idx = TIME_SLOTS.index(focus_start)
        focus_hours = TIME_SLOTS[focus_idx:focus_idx + 2]

        return {
            "preferred_meeting_times": sorted(preferred_times),
            "focus_hours": focus_hours,
            "priority_weight": {"high": 3, "medium": 2, "low": 1},
            "max_meetings_per_day": random.choice([4, 5, 6]),
            "preferred_break_after": random.choice([2, 3]),
        }

    def __repr__(self) -> str:
        return f"ScenarioGenerator(difficulty={self.difficulty})"