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 =====")
    
    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 positions of objects in the scene
        positions = get_object_positions()
        def resolve(key1, key2=None):
            if key1 in positions:
                return positions[key1]
            if key2 and key2 in positions:
                return positions[key2]
            raise KeyError(f"No position for keys {key1} or {key2}")

        # Define key positions and orientations
        ninety_deg_quat = np.array([0.7071, 0.0, 0.0, 0.7071])  # 90° about X-axis
        side_pos_bottom = resolve('side-pos-bottom', 'bottom_side_pos')
        anchor_pos_bottom = resolve('anchor-pos-bottom', 'bottom_anchor_pos')
        tomato1_pos = resolve('tomato1')
        tomato2_pos = resolve('tomato2')
        plate_pos = resolve('plate')

        # === Plan Execution ===
        # The frozen steps must remain exactly as provided:
    [Frozen Code Start]
    obs, reward, done = rotate(env, task, ninety_deg_quat)
obs, reward, done = move(env, task, side_pos_bottom)
obs, reward, done = move(env, task, anchor_pos_bottom)
obs, reward, done = pull(env, task, pull_distance=0.1, pull_axis='x')
obs, reward, done = pick(env, task, target_pos=tomato1_pos)
    [Frozen Code End]

        if done:
            print("Early termination after initial pick.")
            return

        print("[Plan] Step 6: move to plate")
        obs, reward, done = move(env, task, plate_pos)
        if done:
            print("Early termination after moving to plate for tomato1.")
            return

        print("[Plan] Step 7: place tomato1 on plate")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("Early termination after placing tomato1.")
            return

        print("[Plan] Step 8: pick tomato2 from table")
        obs, reward, done = pick(env, task, target_pos=tomato2_pos)
        if done:
            print("Early termination after picking tomato2.")
            return

        print("[Plan] Step 9: move to plate")
        obs, reward, done = move(env, task, plate_pos)
        if done:
            print("Early termination after moving to plate for tomato2.")
            return

        print("[Plan] Step 10: place tomato2 on plate")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("Task completed successfully with done=True. Final reward:", reward)
        else:
            print("Task finished execution. done=False. Reward:", reward)

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

if __name__ == "__main__":
    run_skeleton_task()