import random

class TaskGenerator:
    """
    Generates tasks dynamically for the agent to solve.
    Every episode gets a fresh task.
    """

    def __init__(self):
        self.tasks = [
            {"question": "What is 2 + 2?",     "answer": "4"},
            {"question": "What is 5 + 3?",     "answer": "8"},
            {"question": "What is 10 - 4?",    "answer": "6"},
            {"question": "What is 3 x 3?",     "answer": "9"},
            {"question": "What is 20 / 4?",    "answer": "5"},
            {"question": "Capital of India?",   "answer": "delhi"},
            {"question": "Color of sky?",       "answer": "blue"},
            {"question": "How many days in a week?", "answer": "7"},
        ]

    def generate(self):
        """Return a random task"""
        task = random.choice(self.tasks)
        return task

    def generate_batch(self, n=5):
        """Return multiple tasks"""
        return random.sample(self.tasks, min(n, len(self.tasks)))


if __name__ == "__main__":
    gen = TaskGenerator()

    print("=== Testing TaskGenerator ===")
    for i in range(3):
        task = gen.generate()
        print(f"Task {i+1}: {task}")