"""Behavior for the Littlehand robot."""

from __future__ import annotations

import os
from typing import Callable, Optional

from ...core.behavior import ActionContext, RobotBehavior, _log
from .actions import LITTLEHAND_ACTIONS
from .signature import LittlehandSignature

_XY_MATCH_RADIUS_MM = float(os.getenv("BAJAR_XY_RADIUS_MM", "40.0"))
_STACK_CLEARANCE_MM = float(os.getenv("STACK_CLEARANCE_MM", "5.0"))  # Clearance when placing on top of objects

class LittlehandBehavior(RobotBehavior):
    """Littlehand behavior using the default pick-and-place actions."""

    ACTIONS = LITTLEHAND_ACTIONS
    CommandSignature = LittlehandSignature

    def action_handlers(self) -> dict[str, Callable[[ActionContext], bool]]:
        return {
            "subir": self.action_subir,
            "bajar": self.action_bajar,
            "ir": self.action_ir,
            "tomar": self.action_tomar,
            "soltar": self.action_soltar,
            "reiniciar": self.action_reiniciar,
        }

    def action_subir(self, ctx: ActionContext) -> bool:
        """Move to the home height (initial Z)."""
        if ctx.pose is None:
            return False
        target_z = ctx.home_pose.get("z", ctx.pose[2])
        return self._queue_move(ctx, ctx.pose[0], ctx.pose[1], target_z)

    def action_bajar(self, ctx: ActionContext) -> bool:
        """Move down by step_mm or to top of object under the tool.

        If holding an object, accounts for its height when placing on obstacles.

        Note: position_mm[2] from the camera represents the TOP surface of the object
        (camera looks down, so it sees the top). We use height_mm only for the HELD
        object to calculate placement position.
        """
        obstacle = self._find_object_under_pose(ctx)

        # Get held object height for stack-aware placement
        held_height = self._get_held_object_height(ctx)

        if obstacle is not None:
            # obstacle.position_mm[2] is the TOP surface of the obstacle
            obstacle_top_z = obstacle.position_mm[2]

            if held_height > 0:
                # Stack-aware: place held object on top of obstacle
                # When vacuum releases, bottom of held object should be at obstacle_top
                # TCP needs to be at: obstacle_top + held_height + tcp_offset
                target_z = obstacle_top_z + held_height + _STACK_CLEARANCE_MM + ctx.config.tcp_offset_mm
                _log(
                    f"bajar: STACK-AWARE placement on '{obstacle.object_type}' color={obstacle.color} "
                    f"obstacle_top_z={obstacle_top_z:.1f} held_height={held_height:.1f} "
                    f"clearance={_STACK_CLEARANCE_MM:.1f} tcp_offset={ctx.config.tcp_offset_mm:.1f} "
                    f"target_z={target_z:.1f}"
                )
            else:
                # Not holding anything: move TCP to object top (for grabbing)
                target_z = obstacle_top_z + ctx.config.tcp_offset_mm
                _log(
                    f"bajar: move to object '{obstacle.object_type}' color={obstacle.color} "
                    f"obstacle_top_z={obstacle_top_z:.1f} tcp_offset={ctx.config.tcp_offset_mm:.1f} "
                    f"target_z={target_z:.1f}"
                )
            return self._queue_move(ctx, ctx.pose[0], ctx.pose[1], target_z)

        target_z = ctx.pose[2] - self.config.step_mm
        _log(f"bajar: no object under pose, step to z={target_z:.1f}")
        return self._queue_move(ctx, ctx.pose[0], ctx.pose[1], target_z)

    def _get_held_object_height(self, ctx: ActionContext) -> float:
        """Get the height of the currently held object.

        Uses configured height from config.toml [object_heights] section,
        falling back to the detected height_mm.
        """
        held_id = ctx.shared_state.get_held_object_id()
        if not held_id:
            return 0.0

        held_obj = ctx.scene.get(held_id)
        if not held_obj:
            _log(f"bajar: held object id={held_id} not found in scene")
            return 0.0

        # Use configured height based on object type and size
        height = self._get_configured_height(held_obj.object_type, held_obj.size)
        _log(f"bajar: holding object id={held_id} type={held_obj.object_type} size={held_obj.size} height={height:.1f}mm")
        return height

    def _get_configured_height(self, object_type: str, size: str) -> float:
        """Get configured height for object type/size from environment or defaults."""
        # Try specific key like OBJECT_HEIGHT_CUBE_BIG
        key = f"OBJECT_HEIGHT_{object_type.upper()}_{size.upper()}"
        height_str = os.getenv(key)
        if height_str:
            try:
                return float(height_str)
            except ValueError:
                pass

        # Try generic key like OBJECT_HEIGHT_CUBE
        key = f"OBJECT_HEIGHT_{object_type.upper()}"
        height_str = os.getenv(key)
        if height_str:
            try:
                return float(height_str)
            except ValueError:
                pass

        # Default height (configurable via OBJECT_HEIGHT_DEFAULT)
        default = float(os.getenv("OBJECT_HEIGHT_DEFAULT", "40.0"))
        return default

    def action_ir(self, ctx: ActionContext) -> bool:
        """Move to object X/Y while keeping current Z."""
        if ctx.pose is None or ctx.target is None:
            return False
        pos = ctx.target.position_mm
        return self._queue_move(ctx, pos[0], pos[1], ctx.pose[2])

    def action_tomar(self, ctx: ActionContext) -> bool:
        """Activate tool (low-level grab) and track held object."""
        self._queue_steps(ctx, self.robot_adapter.grab())

        # Find the object under current pose to track for stack-aware placement
        # (ctx.target may be None since requires_object=False in littlehand)
        target_obj = ctx.target or self._find_object_under_pose(ctx)
        if target_obj is not None:
            ctx.shared_state.set_held_object_id(target_obj.id)
            height = self._get_configured_height(target_obj.object_type, target_obj.size)
            _log(f"tomar: now holding object id={target_obj.id} type={target_obj.object_type} height={height:.1f}mm")
        else:
            _log("tomar: no object found under pose, not tracking held object")
        return True

    def action_soltar(self, ctx: ActionContext) -> bool:
        """Deactivate tool (low-level release) and clear held object."""
        self._queue_steps(ctx, self.robot_adapter.release())
        # Clear the held object tracking
        held_id = ctx.shared_state.get_held_object_id()
        if held_id:
            _log(f"soltar: released object id={held_id}")
        ctx.shared_state.clear_held_object()
        return True

    def action_reiniciar(self, ctx: ActionContext) -> bool:
        """Reset: release tool, move home, clear objects and held state."""
        self._queue_steps(ctx, self.robot_adapter.reset_tool())
        self._queue_steps(ctx, self.robot_adapter.move(ctx.home_pose))
        ctx.scene.clear_detected()
        ctx.shared_state.clear_held_object()
        _log("reiniciar: cleared held object state")
        return True

    def _find_object_under_pose(self, ctx: ActionContext) -> Optional["SceneObject"]:
        """Find the topmost object near the current pose x,y (mm).

        Note: position_mm[2] is treated as the TOP surface of the object
        (camera looks down, sees the top surface).
        """
        if ctx.pose is None:
            _log("bajar: missing pose, cannot find object under tool")
            return None
        px, py, pz = ctx.pose[0], ctx.pose[1], ctx.pose[2]
        _log(
            "bajar: current pose x={:.1f} y={:.1f} z={:.1f}".format(px, py, pz)
        )
        candidates = []
        for obj in ctx.scene.query():
            dx = px - obj.position_mm[0]
            dy = py - obj.position_mm[1]
            dist2 = dx * dx + dy * dy
            if dist2 > _XY_MATCH_RADIUS_MM * _XY_MATCH_RADIUS_MM:
                continue
            # position_mm[2] IS the top surface (camera sees top of object)
            top_surface_z = obj.position_mm[2]
            candidates.append((top_surface_z, obj))
            _log(
                "bajar: near id={} type={} color={} center=({:.1f},{:.1f}) "
                "dist_xy={:.1f} top_z={:.1f}".format(
                    obj.id,
                    obj.object_type,
                    obj.color,
                    obj.position_mm[0],
                    obj.position_mm[1],
                    (dist2 ** 0.5),
                    top_surface_z,
                )
            )
        if not candidates:
            _log(f"bajar: no objects within radius={_XY_MATCH_RADIUS_MM:.1f}mm")
            return None
        candidates.sort(key=lambda item: item[0], reverse=True)
        return candidates[0][1]
