from env import setup_environment, shutdown_environment
from skill_code import pick, place, move
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 get_observation for recording
        original_step = task.step
        task.step = recording_step(original_step)
        original_get_obs = task.get_observation
        task.get_observation = recording_get_observation(original_get_obs)

        # === Retrieve Object Positions ===
        positions = get_object_positions()
        if 'plate' not in positions:
            raise KeyError("Plate object not found in positions.")
        plate_pos = positions['plate']

        # Identify all tomatoes in the scene
        tomato_names = sorted([name for name in positions.keys() if name.startswith('tomato')])
        if not tomato_names:
            print("[Task] No tomatoes found to pick and place.")
        else:
            # === Execute Oracle Plan: pick and place each tomato onto the plate ===
            for tomato in tomato_names:
                tomato_pos = positions[tomato]
                print(f"[Task] Picking up {tomato} at: {tomato_pos}")
                try:
                    obs, reward, done = pick(
                        env,
                        task,
                        target_pos=np.array(tomato_pos),
                        approach_distance=0.15,
                        max_steps=100,
                        threshold=0.01,
                        approach_axis='z',
                        timeout=10.0
                    )
                except Exception as e:
                    print(f"[Error] Exception during pick of {tomato}: {e}")
                    break
                if done:
                    print("[Task] Task ended during pick!")
                    return

                print(f"[Task] Placing {tomato} on plate at: {plate_pos}")
                try:
                    obs, reward, done = place(
                        env,
                        task,
                        target_pos=np.array(plate_pos),
                        approach_distance=0.15,
                        max_steps=100,
                        threshold=0.01,
                        approach_axis='z',
                        timeout=10.0
                    )
                except Exception as e:
                    print(f"[Error] Exception during place of {tomato}: {e}")
                    break
                if done:
                    print("[Task] Task ended during place!")
                    return

            print(f"[Task] Completed dropping {len(tomato_names)} tomato(s) onto the plate. Final reward: {reward}")

    except Exception as e:
        print(f"[Task] Encountered exception: {e}")
    finally:
        # Always ensure the environment is properly shutdown
        shutdown_environment(env)

    print("===== End of Skeleton Task =====")

if __name__ == "__main__":
    run_skeleton_task()