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

def run_skeleton_task():
    print("===== Starting Skeleton Task =====")
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()
        init_video_writers(obs)
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # Retrieve object positions
        positions = get_object_positions()

        # Step 1: Rotate gripper orientation (from zero_deg to ninety_deg)
        obs = task.get_observation()
        current_quat = obs.gripper_pose[3:7]
        print("[Task] Rotating gripper to desired orientation.")
        [Frozen Code Start]
        obs, reward, done = rotate(env, task, target_quat=current_quat)
        [Frozen Code End]
        if done:
            print("[Task] Episode ended after rotate.")
            return

        # Step 2: Move to bottom side position
        bottom_side = np.array(positions.get('bottom_side_pos', [0.0, 0.0, 0.0]))
        print(f"[Task] Moving to side of bottom drawer at {bottom_side}.")
        obs, reward, done = move(env, task, target_pos=bottom_side)
        if done:
            print("[Task] Episode ended after move to side.")
            return

        # Step 3: Move to bottom anchor position
        bottom_anchor = np.array(positions.get('bottom_anchor_pos', [0.0, 0.0, 0.0]))
        print(f"[Task] Moving to anchor of bottom drawer at {bottom_anchor}.")
        obs, reward, done = move(env, task, target_pos=bottom_anchor)
        if done:
            print("[Task] Episode ended after move to anchor.")
            return

        # Step 4: Pull the drawer open
        print("[Task] Pulling bottom drawer open.")
        obs, reward, done = pull(env, task, pull_distance=0.1, pull_axis='x')
        if done:
            print("[Task] Episode ended after pull.")
            return

        # Step 5: Move to first tomato (item1)
        item1_pos = np.array(positions.get('item1', [0.0, 0.0, 0.0]))
        print(f"[Task] Moving to first tomato at {item1_pos}.")
        obs, reward, done = move(env, task, target_pos=item1_pos)
        if done:
            print("[Task] Episode ended after move to item1.")
            return

        # Step 6: Pick the first tomato
        print("[Task] Picking up first tomato.")
        obs, reward, done = pick(env, task, target_pos=item1_pos)
        if done:
            print("[Task] Episode ended after pick.")
            return

        # Step 7: Move to plate
        plate_pos = np.array(positions.get('plate', [0.0, 0.0, 0.0]))
        print(f"[Task] Moving to plate at {plate_pos}.")
        obs, reward, done = move(env, task, target_pos=plate_pos)
        if done:
            print("[Task] Episode ended after move to plate.")
            return

        # Step 8: Place the first tomato on the plate
        print("[Task] Placing first tomato on plate.")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("[Task] Episode ended after place.")
            return

        # Step 9: Move to second tomato (item2)
        item2_pos = np.array(positions.get('item2', [0.0, 0.0, 0.0]))
        print(f"[Task] Moving to second tomato at {item2_pos}.")
        obs, reward, done = move(env, task, target_pos=item2_pos)
        if done:
            print("[Task] Episode ended after move to item2.")
            return

        # Step 10: Pick the second tomato
        print("[Task] Picking up second tomato.")
        obs, reward, done = pick(env, task, target_pos=item2_pos)
        if done:
            print("[Task] Episode ended after pick.")
            return

        # Step 11: Move back to plate
        print(f"[Task] Moving back to plate at {plate_pos}.")
        obs, reward, done = move(env, task, target_pos=plate_pos)
        if done:
            print("[Task] Episode ended after move to plate.")
            return

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

        print("[Task] All plan steps executed successfully.")

    finally:
        shutdown_environment(env)
        print("===== End of Skeleton Task =====")


if __name__ == "__main__":
    run_skeleton_task()