import numpy as np
from scipy.spatial.transform import Rotation as R

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

# === Skill Imports ===
#   Do NOT redefine these – they already exist in skill_code.
from skill_code import rotate, move, pull, pick, place


def _safe_pos(positions_dict, obj_name):
    """Utility: fetch an object's 3-D position, raise a clear error if missing."""
    if obj_name not in positions_dict:
        raise KeyError(f"[run_task] Could not find position for '{obj_name}'. "
                       f"Make sure the name is correct and the object exists.")
    return positions_dict[obj_name]


def run_task():
    """Run the full oracle plan specified in the specification."""
    print("===== [run_task]  Starting =====")

    # ------------------------------------------------------------
    # 1)  Environment setup & reset
    # ------------------------------------------------------------
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()

        # -------------  video helpers (optional) -------------
        init_video_writers(obs)
        task.step            = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # -------------  Fetch all object positions -------------
        positions = get_object_positions()

        # Drawer related key names (bottom drawer chosen)
        side_key   = 'bottom_side_pos'
        anchor_key = 'bottom_anchor_pos'

        # Tomato & plate key names  (item1 & item2 are used as tomatoes)
        tomato1_key = 'item1'
        tomato2_key = 'item2'
        plate_key   = 'plate'

        # Resolve positions with helpful error messages
        side_pos     = _safe_pos(positions, side_key)
        anchor_pos   = _safe_pos(positions, anchor_key)
        tomato1_pos  = _safe_pos(positions, tomato1_key)
        tomato2_pos  = _safe_pos(positions, tomato2_key)
        plate_pos    = _safe_pos(positions, plate_key)

        # ------------------------------------------------------------
        # ===  Oracle Plan – step-by-step  ===========================
        # ------------------------------------------------------------
        # [Frozen Code Start]
        current_quat = task.get_observation().gripper_pose[3:7]
        target_quat  = (R.from_quat(current_quat) * R.from_euler('z', 90, degrees=True)).as_quat()
        obs, reward, done = rotate(env, task, target_quat)
        obs, reward, done = move(env, task, side_pos)
        obs, reward, done = move(env, task, anchor_pos)
        obs, reward, done = pull(env, task, pull_distance=0.20, pull_axis='-x')
        # [Frozen Code End]

        if done:
            print("[run_task] Task ended unexpectedly during drawer opening.")
            return

        # ------------------------------------------------------------
        # Pick tomato #1
        # ------------------------------------------------------------
        print("\n[Step-5] pick   → item1 (tomato #1)")
        obs, reward, done = pick(
            env, task,
            target_pos=tomato1_pos,
            approach_distance=0.15,
            approach_axis='-z'
        )
        if done:
            print("[run_task] Task ended unexpectedly after picking item1.")
            return

        # ------------------------------------------------------------
        # Place tomato #1 on plate
        # ------------------------------------------------------------
        print("\n[Step-6] place  → item1 onto plate")
        obs, reward, done = place(
            env, task,
            target_pos=plate_pos,
            approach_distance=0.15,
            approach_axis='-z'
        )
        if done:
            print("[run_task] Task ended unexpectedly after placing item1.")
            return

        # ------------------------------------------------------------
        # Pick tomato #2
        # ------------------------------------------------------------
        print("\n[Step-7] pick   → item2 (tomato #2)")
        obs, reward, done = pick(
            env, task,
            target_pos=tomato2_pos,
            approach_distance=0.15,
            approach_axis='-z'
        )
        if done:
            print("[run_task] Task ended unexpectedly after picking item2.")
            return

        # ------------------------------------------------------------
        # Place tomato #2 on plate
        # ------------------------------------------------------------
        print("\n[Step-8] place  → item2 onto plate")
        obs, reward, done = place(
            env, task,
            target_pos=plate_pos,
            approach_distance=0.15,
            approach_axis='-z'
        )
        if done:
            print("[run_task] Task ended unexpectedly after placing item2.")
            return

        # ------------------------------------------------------------
        #  Goal achieved!
        # ------------------------------------------------------------
        print("\n===== [run_task]  Task completed successfully! =====")
    except Exception as e:
        # Provide a clear message but always attempt clean shutdown
        print("=== [run_task]  ERROR:", str(e))
        raise
    finally:
        # Ensure the simulator shuts down no matter what
        shutdown_environment(env)
        print("===== [run_task]  Environment shut down =====")


if __name__ == "__main__":
    run_task()