from env import setup_environment, shutdown_environment
from skill_code import move, pick, place, rotate, pull
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions

def run_skeleton_task():
    '''Task: Unlock the cabinet by opening a drawer, then transfer two tomatoes onto the plate.'''
    print("===== Starting Task: Open Drawer & Move Tomatoes =====")
    env, task = setup_environment()
    try:
        # Reset and set up recording
        descriptions, obs = task.reset()
        init_video_writers(obs)
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # Retrieve positions for drawer and items
        positions = get_object_positions()
        side_pos_key = 'bottom_side_pos'
        anchor_pos_key = 'bottom_anchor_pos'
        item1_key = 'item1'
        item2_key = 'item2'
        plate_key = 'plate'

        # Step 1: Rotate gripper (no-op rotation to satisfy rotate call)
        print("[Step 1] rotate gripper")
        obs = task.get_observation()
        current_quat = obs.gripper_pose[3:7]
[Frozen Code Start]
        obs, reward, done = rotate(env, task, target_quat=current_quat)
[Frozen Code End]
        if done:
            print("Task ended unexpectedly after rotate.")
            return

        # Step 2: Move to the side position of the bottom drawer
        print("[Step 2] move to drawer side position")
        drawer_side_pos = positions[side_pos_key]
        obs, reward, done = move(env, task, target_pos=drawer_side_pos)
        if done:
            print("Task ended unexpectedly after moving to drawer side.")
            return

        # Step 3: Move to the anchor position of the bottom drawer
        print("[Step 3] move to drawer anchor position")
        drawer_anchor_pos = positions[anchor_pos_key]
        obs, reward, done = move(env, task, target_pos=drawer_anchor_pos)
        if done:
            print("Task ended unexpectedly after moving to drawer anchor.")
            return

        # Step 4: Pick the bottom drawer handle
        print("[Step 4] pick bottom drawer")
        obs, reward, done = pick(env, task, target_pos=drawer_anchor_pos)
        if done:
            print("Task ended unexpectedly during pick bottom drawer.")
            return

        # Step 5: Pull drawer open
        print("[Step 5] pull drawer open")
        obs, reward, done = pull(env, task, pull_distance=0.1, pull_axis='x')
        if done:
            print("Task ended unexpectedly during pull.")
            return

        # Step 6: Pick the first tomato
        print("[Step 6] pick first tomato")
        item1_pos = positions[item1_key]
        obs, reward, done = pick(env, task, target_pos=item1_pos)
        if done:
            print("Task ended unexpectedly during pick first tomato.")
            return

        # Step 7: Place the first tomato on the plate
        print("[Step 7] place first tomato on plate")
        plate_pos = positions[plate_key]
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("Task ended unexpectedly during place first tomato.")
            return

        # Step 8: Pick the second tomato
        print("[Step 8] pick second tomato")
        item2_pos = positions[item2_key]
        obs, reward, done = pick(env, task, target_pos=item2_pos)
        if done:
            print("Task ended unexpectedly during pick second tomato.")
            return

        # Step 9: Place the second tomato on the plate
        print("[Step 9] place second tomato on plate")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("Task ended unexpectedly during place second tomato.")
            return

        # Final check
        if reward is not None and reward > 0:
            print("Task completed successfully! Reward:", reward)
        else:
            print("Task finished but did not achieve all goals (done=False or zero reward).")

    except Exception as e:
        print("Error during task execution:", str(e))
    finally:
        shutdown_environment(env)
        print("===== End of Task =====")

if __name__ == "__main__":
    run_skeleton_task()