# run_skeleton_task.py  (Completed)

import sys
import time
import types
import inspect
import traceback
from typing import Dict, Any, Tuple, List

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

# --- Environment / helpers ( *do not modify these imports* ) -----------------
from env import setup_environment, shutdown_environment

from skill_code import *          # noqa: F401  (we rely on the skills provided in the package)
from video import (               # noqa: F401
    init_video_writers,
    recording_step,
    recording_get_observation,
)
from object_positions import get_object_positions

# -----------------------------------------------------------------------------
# Helper utilities
# -----------------------------------------------------------------------------
_AVAILABLE_SKILLS = {
    "move": move,
    "pick": pick,
    "place": place,
}

# Optional skills – these might or might not exist in skill_code depending on
# the benchmark / task definition.  We defensively fetch them if present.
for _opt in ("rotate", "pull"):
    if _opt in globals() and isinstance(globals()[_opt], types.FunctionType):
        _AVAILABLE_SKILLS[_opt] = globals()[_opt]


def _safe_call_skill(skill_name: str,
                     env,
                     task,
                     **kwargs) -> Tuple[Any, Any, bool]:
    """
    Wrapper that introspects the skill function’s signature and
    gracefully ignores super-fluous keyword arguments.

    Returns (obs, reward, done).  If the underlying skill does not
    follow exactly that convention we return (None, None, False) so
    that the caller can continue execution.
    """
    if skill_name not in _AVAILABLE_SKILLS:
        print(f"[WARN] Skill <{skill_name}> not available – skipping.")
        return None, None, False

    fn = _AVAILABLE_SKILLS[skill_name]
    try:
        sig = inspect.signature(fn)
        call_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
        # Skills are assumed to expect at least (env, task, …)
        res = fn(env, task, **call_kwargs)
        # Many RLBench skills return (obs, reward, done) but some might
        # only return None.  Normalise here.
        if isinstance(res, tuple) and len(res) == 3:
            return res
        return res, None, False
    except Exception as exc:          # noqa: BLE001
        print(f"[ERROR] Exception in skill <{skill_name}>: {exc}")
        traceback.print_exc()
        # Continue execution even if the single skill fails; we do not
        # want the whole exploration phase to crash.
        return None, None, False


# -----------------------------------------------------------------------------
#   Very small “world-state” blackboard so we can remember exploration facts
# -----------------------------------------------------------------------------
class Blackboard(dict):
    """A dict that exposes .get_predicate()/ .set_predicate() helpers."""

    def set_predicate(self, predicate: str, key: str) -> None:
        self[(predicate, key)] = True

    def get_predicate(self, predicate: str, key: str) -> bool:
        return self.get((predicate, key), False)


# -----------------------------------------------------------------------------
#   Exploration routine – tries to reveal hidden predicates mentioned in the
#   exploration domain description (identified, temperature-known, …).
#   This phase is completely generic and does not assume a concrete task goal.
# -----------------------------------------------------------------------------
def _exploration_phase(env, task, positions: Dict[str, np.ndarray]) -> None:
    print("-----  Exploration Phase  -----")

    bb = Blackboard()
    robot_placeholder = "robot_0"

    # 1)  IDENTIFICATION (approach every object position)
    for obj_name, pos in positions.items():
        print(f"[Explore] Move close to <{obj_name}> for identification.")
        _safe_call_skill("move", env=env, task=task, target_pos=pos + np.array([0, 0, 0.10]))
        bb.set_predicate("identified", obj_name)

    # 2)  WEIGHT-KNOWLEDGE & DURABILITY – attempt a pick on light objects
    #     We simply try to pick every object once.  Failures are ignored.
    for obj_name, pos in positions.items():
        print(f"[Explore] Attempt pick on <{obj_name}>.")
        obs, reward, done = _safe_call_skill(
            "pick",
            env=env,
            task=task,
            target_pos=pos,
        )
        if done:
            print("[Explore] Environment signalled ‘done’ during pick – resetting.")
            task.reset()

        # Mark potential discoveries
        bb.set_predicate("weight-known", obj_name)
        bb.set_predicate("durability-known", obj_name)

        # If pick succeeded we need to place the object back
        if "holding" in str(type(obs)).lower():  # heuristic
            _safe_call_skill(
                "place",
                env=env,
                task=task,
                target_pos=pos + np.array([0.05, 0, 0]),
            )

    # 3)  Potential drawer exploration (rotation + pull)
    drawer_keys = [k for k in positions if "drawer" in k.lower()]
    for drawer_name in drawer_keys:
        print(f"[Explore] Try rotating / pulling drawer <{drawer_name}>.")
        # We do not know exact anchor/side positions.  Simple heuristic:
        drawer_pos = positions[drawer_name]
        _safe_call_skill("move", env=env, task=task, target_pos=drawer_pos + np.array([0, 0, 0.12]))
        _safe_call_skill("rotate", env=env, task=task, angle=np.pi / 2)
        _safe_call_skill("pull", env=env, task=task)

    print("----- End Exploration Phase -----")
    # We could analyse the blackboard here if we needed to “deduce” a
    # missing predicate, but for the purposes of this skeleton we only
    # demonstrate how the exploration might be carried out.


# -----------------------------------------------------------------------------
#   MAIN  – Generic task runner
# -----------------------------------------------------------------------------
def run_skeleton_task() -> None:
    """Generic skeleton for running any task in the simulation."""
    print("\n==============================")
    print("   Starting Skeleton Task")
    print("==============================")

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

        # -------- Optional: enable video capture ----------------------
        init_video_writers(obs)                       # noqa: F405
        task.step = recording_step(task.step)         # noqa: F405
        task.get_observation = recording_get_observation(task.get_observation)  # noqa: F405

        # ------------------------------------------------ Obtain object positions
        positions: Dict[str, np.ndarray] = get_object_positions()
        print(f"[Info] Found {len(positions)} objects from helper.")

        # ------------------------------------------------ Exploration –
        # This phase tries to reveal the “missing predicates” mentioned
        # in the problem description.
        _exploration_phase(env, task, positions)

        # ------------------------------------------------ Main task plan
        # NOTE: Without a concrete goal description we perform a
        #       demonstrative behaviour: pick the first movable object,
        #       move it to a hard-coded drop-zone, and finish.
        movable_objects: List[str] = [k for k in positions.keys()
                                      if "drawer" not in k.lower()]

        if not movable_objects:
            print("[Plan] No obvious movable objects – skipping main plan.")
        else:
            main_obj = movable_objects[0]
            main_pos = positions[main_obj]
            print(f"[Plan] Executing simple demonstration plan on <{main_obj}>.")

            # Approach & pick
            _safe_call_skill("move", env=env, task=task, target_pos=main_pos + np.array([0, 0, 0.10]))
            _safe_call_skill("pick", env=env, task=task, target_pos=main_pos)

            # Move to “drop zone” (arbitrary offset) and place
            drop_zone = main_pos + np.array([0.20, 0.20, 0])
            _safe_call_skill("move", env=env, task=task, target_pos=drop_zone + np.array([0, 0, 0.15]))
            _safe_call_skill("place", env=env, task=task, target_pos=drop_zone)

            print("[Plan] Demonstration plan completed.")

        # ------------------------------------------------ (END)

    finally:
        # Always ensure the simulator terminates cleanly
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()