import os
from env.delivery_env import DeliveryEnv
from graders.grader import compute_score

TASK = os.getenv("TASK", "easy")

env = DeliveryEnv()
state = env.reset(TASK)

print(f"[START] task={TASK} env=food_delivery model=smart-urgency-agent")

rewards = []
step_count = 0

while True:
    orders = [o for o in state["orders"] if not o["assigned"]]
    riders = state["riders"]

    if not orders:
        break

    best_score = float("inf")
    best_action = None

    # 🔥 Evaluate all (order, rider) pairs
    for order in orders:
        for rider in riders:
            distance = abs(rider["location"] - order["location"])
            urgency = order["deadline"] - state["time"]

            score = distance + urgency  # lower is better

            if score < best_score:
                best_score = score
                best_action = (order, rider)

    order, rider = best_action
    action = f"assign_{order['id']}_to_{rider['id']}"

    state, reward, done, info = env.step(action)

    rewards.append(reward)
    step_count += 1

    error = info["error"] if info["error"] else "null"

    print(f"[STEP] step={step_count} action={action} reward={reward:.2f} done={str(done).lower()} error={error}")

    if done:
        break

max_possible = len(rewards) * 1.5  # adjusted for bonus rewards
score = compute_score(sum(rewards), max_possible)

print(f"[END] success=true steps={step_count} score={score:.2f} rewards={','.join(f'{r:.2f}' for r in rewards)}")