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

def run_skeleton_task():
    print("===== Starting Skeleton Task =====")
    
    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset the task to its initial state
        descriptions, obs = task.reset()

        # Initialize video writers (optional)
        init_video_writers(obs)

        # Wrap step and observation for recording
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # === Retrieve Object Positions ===
        positions = get_object_positions()
        # Extract relevant positions
        # Drawer interaction positions
        bottom_side_pos = np.array(positions['bottom_side_pos'])
        bottom_anchor_pos = np.array(positions['bottom_anchor_pos'])
        # Table items
        item1_pos = np.array(positions['item1'])
        item2_pos = np.array(positions['item2'])
        # Plate
        plate_pos = np.array(positions['plate'])

        # === Oracle Plan Execution ===

        # 1) Open the bottom drawer
        print("[Plan] Moving to drawer side position")
        obs, reward, done = move(env, task, target_pos=bottom_side_pos)
        if done:
            print("[Plan] Terminated unexpectedly during move-to-side.")
            return

        print("[Plan] Moving to drawer anchor position")
        obs, reward, done = move(env, task, target_pos=bottom_anchor_pos)
        if done:
            print("[Plan] Terminated unexpectedly during move-to-anchor.")
            return

        print("[Plan] Pulling drawer to open")
        obs, reward, done = pull(env, task, pull_distance=0.1, pull_axis='x')
        if done:
            print("[Plan] Task ended after drawer pull.")
            return

        # 2) Pick and toss the two tomatoes onto the plate
        for idx, item_pos in enumerate([item1_pos, item2_pos], start=1):
            print(f"[Plan] Picking item{idx} at position {item_pos}")
            obs, reward, done = pick(
                env,
                task,
                target_pos=item_pos,
                approach_distance=0.15,
                max_steps=100,
                threshold=0.01,
                approach_axis='z',
                timeout=10.0
            )
            if done:
                print(f"[Plan] Task ended during pick of item{idx}.")
                return

            print(f"[Plan] Placing item{idx} on plate at position {plate_pos}")
            obs, reward, done = place(
                env,
                task,
                target_pos=plate_pos,
                approach_distance=0.15,
                max_steps=100,
                threshold=0.01,
                approach_axis='z',
                timeout=10.0
            )
            if done:
                print(f"[Plan] Task ended during place of item{idx}.")
                return

        print("[Plan] Completed tossing tomatoes onto the plate.")
    finally:
        shutdown_environment(env)
        print("===== End of Skeleton Task =====")

if __name__ == "__main__":
    run_skeleton_task()