class EmergencyEnv:
    def __init__(self, difficulty="easy"):
        self.difficulty = difficulty
        self.grid_size = 5
        self.reset()

    def reset(self):
        self.agent = [0, 0]
        self.resources = 3

        if self.difficulty == "easy":
            self.requests = [
                {"pos": [4, 4], "priority": "normal"}
            ]
            self.time_left = 20

        elif self.difficulty == "medium":
            self.requests = [
                {"pos": [4, 4], "priority": "critical"},
                {"pos": [2, 3], "priority": "normal"}
            ]
            self.time_left = 30

        else:
            self.requests = [
                {"pos": [4, 4], "priority": "critical"},
                {"pos": [2, 3], "priority": "critical"},
                {"pos": [1, 4], "priority": "normal"}
            ]
            self.time_left = 40

        return self.get_state()

    def get_state(self):
        return {
            "agent": self.agent,
            "requests": self.requests,
            "num_requests": len(self.requests),
            "resources_left": self.resources,
            "time_left": self.time_left
        }

    def step(self, action):
        reward = -0.5   # reduced penalty
        done = False

        # movement
        if action == 0: self.agent[0] -= 1
        elif action == 1: self.agent[0] += 1
        elif action == 2: self.agent[1] -= 1
        elif action == 3: self.agent[1] += 1

        # boundary
        self.agent[0] = max(0, min(4, self.agent[0]))
        self.agent[1] = max(0, min(4, self.agent[1]))

        # allocation
        if action == 4 and self.resources > 0:
            for req in self.requests:
                if self.agent == req["pos"]:
                    reward += 20 if req["priority"] == "critical" else 10
                    self.requests.remove(req)
                    self.resources -= 1
                    break

        self.time_left -= 1

        if self.time_left <= 0:
            reward -= 20
            done = True

        if len(self.requests) == 0:
            reward += 30
            done = True

        return self.get_state(), reward, done, {}