"""Peg-in-hole: grasp near the top, align over the socket, insert straight down."""
import numpy as np

from ..motion.arm import OPEN, CLOSE, fillet
from ..envs.scene import TABLE_TOP
from .base import approach_and_grasp


def solve(env, obj, target, arm="right", grasp_top=0.03, max_gap=0.03):
    a = env.arms[arm]
    rec = env.recorder
    res, ext = approach_and_grasp(env, a, obj, grasp_top=grasp_top, max_gap=max_gap)
    if not res.ok:
        rec.phase = "ABORT: no grasp"
        return {"grasped": False, "reason": res.reason}
    reg = env.scene.regions[target]
    above = a.to_root(np.array([reg["xy"][0], reg["xy"][1], TABLE_TOP+0.20], np.float32))
    # the peg was grasped while standing on the table, so returning the fingers to that height
    # puts its base back at table level -- fully seated, since the hole's floor IS the table
    seat = np.array([above[0], above[1], float(res.hold_pose[2])], np.float32)
    rec.phase = "4. CARRY + ALIGN over the hole"
    a.flow(fillet([a._seg_start(), res.hold_pose+np.array([0, 0, 0.16], np.float32), above])[1:], CLOSE)

    # VISUAL-SERVO the correction: the part slips in the jaws during the carry, so the gripper
    # being over the hole is not the same as the PART being over it. Measure the part's live xy
    # and shift the descent by the residual -- without this the ring lands beside the post.
    live = env.scene.object_pos(obj)
    delta = np.array([reg["xy"][0]-float(live[0]), reg["xy"][1]-float(live[1])], float)
    if float(np.linalg.norm(delta)) > 0.002:
        rec.phase = "4b. FINE-ALIGN on the part's own pose"
        above = above+np.array([delta[0], delta[1], 0.0], np.float32)
        seat = seat+np.array([delta[0], delta[1], 0.0], np.float32)
        a.move_to(above, CLOSE, tol=0.005, max_steps=90)
        print(f"[solver] fine-align: part was off by {np.round(delta,4)} m", flush=True)

    rec.phase = "5. INSERT"
    err = a.move_to(seat, CLOSE, tol=0.006, max_steps=170)
    rec.phase = "6. RELEASE"
    a.move_to(seat, OPEN, tol=0.02, max_steps=45)
    rec.phase = "7. RETREAT"
    a.flow([seat+np.array([0, 0, 0.16], np.float32)], OPEN)
    fin = env.scene.object_pos(obj)
    miss = float(np.hypot(fin[0]-reg["xy"][0], fin[1]-reg["xy"][1]))
    print(f"[solver] {obj} settled at ({fin[0]:+.3f},{fin[1]:+.3f},{fin[2]:.3f}), "
          f"{miss:.3f} m from {target}", flush=True)
    return {"grasped": True, "insert_err": err, "miss": miss}