import random

class DeliveryEnv:
    def __init__(self):
        self.grid_size = 5
        self.reset()

    def reset(self):
        self.agent_pos = [0, 0]
        self.pickup = [2, 2]
        self.drop = [4, 4]
        self.has_item = False
        self.steps = 0
        return self.state()

    def state(self):
        return {
            "agent_pos": self.agent_pos,
            "has_item": self.has_item,
            "pickup": self.pickup,
            "drop": self.drop
        }

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

        # movement
        if action == "up":
            self.agent_pos[0] = max(0, self.agent_pos[0] - 1)
        elif action == "down":
            self.agent_pos[0] = min(self.grid_size - 1, self.agent_pos[0] + 1)
        elif action == "left":
            self.agent_pos[1] = max(0, self.agent_pos[1] - 1)
        elif action == "right":
            self.agent_pos[1] = min(self.grid_size - 1, self.agent_pos[1] + 1)

        # pickup
        elif action == "pickup":
            if self.agent_pos == self.pickup and not self.has_item:
                self.has_item = True
                reward += 10

        # drop
        elif action == "drop":
            if self.agent_pos == self.drop and self.has_item:
                reward += 20
                done = True

        self.steps += 1

        if self.steps > 50:
            done = True

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