# run_skeleton_task.py (Completed Executable)

import time
import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# Import every predefined skill exactly as delivered in skill_code
from skill_code import *

from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


def _safe_get(dictionary, key, default):
    """Utility: return dictionary[key] if it exists, otherwise default."""
    if key in dictionary:
        return dictionary[key]
    print(f"[WARN] Requested key '{key}' missing – substituting default value.")
    return default


def _attempt_skill_call(func, *args, **kwargs):
    """Call a skill and swallow exceptions so that exploration can continue."""
    try:
        return func(*args, **kwargs)
    except Exception as exc:  # noqa: broad-exception-caught
        print(f"[ERROR] {func.__name__} failed with: {exc}")
        return None, 0.0, False


def _exploration_phase(env, task, positions):
    """
    Try the available skills on every known object.
    The purpose is to discover which predicate / capability might be missing.
    """
    print("===== [Exploration] Phase START =====")

    # ------------
    # 1) Rotation‐based test (does gripper reach 90°?)
    # ------------
    target_quat = [0.0, 0.0, 0.7071068, 0.7071068]  # 90° about Z
    _attempt_skill_call(rotate, env, task, target_quat)

    # ------------
    # 2) Iterate over every registered object and try pick / place
    # ------------
    for obj_name, obj_pos in positions.items():
        print(f"[Exploration] Trying pick/place on '{obj_name}' @ {obj_pos}")
        obs, reward, done = _attempt_skill_call(
            pick,
            env,
            task,
            target_pos=obj_pos,
            approach_distance=0.15,
            max_steps=120,
            threshold=0.01,
            approach_axis="z",
            timeout=8.0,
        )

        if done:
            print("[Exploration] Task ended while picking – stopping exploration.")
            return

        # Place it back (use same target for simplicity)
        _attempt_skill_call(
            place,
            env,
            task,
            target_pos=obj_pos,
            approach_distance=0.15,
            max_steps=120,
            threshold=0.02,
            approach_axis="z",
            timeout=8.0,
        )

    # ------------
    # 3) Attempt to pull anything that looks like a drawer
    # ------------
    drawer_candidates = [name for name in positions if "drawer" in name.lower()]
    for drawer_name in drawer_candidates:
        print(f"[Exploration] Trying pull on '{drawer_name}'")
        _attempt_skill_call(pull, env, task, drawer_name)

    print("===== [Exploration] Phase END =====")


def run_skeleton_task():
    """Generic skeleton for running any task in your simulation."""
    print("===== Starting Skeleton Task =====")

    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset the task
        descriptions, obs = task.reset()

        # (Optional) Video init
        init_video_writers(obs)

        # Wrap task.step/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()  # {name: (x, y, z), ...}

        # Feedback indicated 'rubbish' may be missing – ensure we have it
        if "rubbish" not in positions:
            # Use gripper location as a best-guess placeholder
            rubbish_guess = tuple(obs.gripper_pose[:3])
            print(
                "[Feedback] 'rubbish' object missing from list – "
                f"injecting placeholder at {rubbish_guess}"
            )
            positions["rubbish"] = rubbish_guess

        # === EXPLORATION PHASE ===
        _exploration_phase(env, task, positions)

        # === MAIN ORACLE PLAN PLACEHOLDER ===
        # Insert high-level task logic here if known.
        # For now we simply log completion of exploration.
        print("===== [Task] No additional scripted plan – finished. =====")

    finally:
        # Ensure environment shutdown
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()