"""
main_loop.py

The state machine that actually runs the challenge, end to end:

    INIT -> LOCATE -> PICK -> TRACE -> RETURN -> DONE
              \_________________________________/
                           on any exception -> ERROR (safe park)

Order and why:
  INIT    load the camera calibration, model, arm connection; self-checks.
  LOCATE  with the arm parked out of view: scan the paper's edges, find the
          line, find the obstacle if any, and compute the FULL trace plan
          (trace_planner.plan_trace) -- this already decides, per stretch
          of line, whether to go around the obstacle or (last resort) over
          it. Also locate the marker's pickup position here, in the same
          out-of-view frame. The plan is fully decided before the arm
          moves at all, which is also why PICK happens after LOCATE, not
          before: the arm commits to a grip only once it already knows
          the whole route it's about to run.
  PICK    elevate to Z_APPROACH above the marker, descend to Z_GRASP, close
          the claw, lift to Z_TRAVEL.
  TRACE   stream the plan's waypoints (trace_with_replanning() below).
          Every TRACE_CHECK_EVERY_MM of travel it grabs a fresh frame and
          cheaply checks whether anything now blocks the REMAINING path (a
          hand, an object placed after LOCATE ran) -- only if that trips
          does it actually replan (re-detect obstacles, reroute just
          what's left) and keep going from the current position. A clear
          path never pauses for this at all.
  RETURN  put the marker back at its original pickup spot, release, go home.
  ERROR   any exception at any stage -> open the claw, home. Never leave
          the arm holding position mid-fault.

Arm control: closed-form IK (arm_config's measured L1/L2/Z_SHOULDER) over serial to
arm_serial_ptp.ino on the ESP32 -- see src/perception/arm_serial.py. Every ArmSerial call
blocks until the firmware ACKs that the move actually finished, so there's no separate
Python-side pacing/sleep needed here; the blocking call itself paces the trace.

Run from the repo root:
    python3 src/control/main_loop.py
"""
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "perception"))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "planning"))

import numpy as np  # noqa: E402

import camera_calib as cc  # noqa: E402
import trace_planner as tp  # noqa: E402
import arm_serial as ars  # noqa: E402
import arm_config as ac  # noqa: E402
import pick_place as pp  # noqa: E402
import marker_position as mpos  # noqa: E402

TRACE_CHECK_EVERY_MM = 20.0  # how often (mm of travel) TRACE checks for a new obstacle mid-run


class ChallengeError(RuntimeError):
    pass


def self_checks(cam):
    """INIT: fail loudly before anything moves, rather than partway through a run."""
    problems = []
    if list(cam.get("frame", [1280, 720])) != [1280, 720]:
        problems.append(f"calib_camera.json was captured at {cam.get('frame')}, expected [1280, 720]")
    if not ars.reachable(ac.L1, 0.0, ac.Z_TRACE):   # a point roughly mid-ring, straight out from the base
        problems.append(f"(x={ac.L1:.0f}, y=0) at Z_TRACE isn't reachable -- "
                         "check L1/L2/Z_SHOULDER/Z_TRACE in arm_config.py before trusting anything else")
    if problems:
        raise ChallengeError("self-check failed: " + "; ".join(problems))


def locate(cap, cam, model, cls_id, correction):
    """Arm must already be parked out of view. Returns (plan, marker_xy_mm)."""
    frame = cc.grab(cap, n=15)
    plan = tp.plan_trace(frame, cam)
    for w in plan.warnings:
        print("plan warning:", w)

    marker_xy = mpos.locate_marker_stable(cap, model, cls_id, cam, correction)
    if marker_xy is None:
        raise ChallengeError("could not get a stable marker position during LOCATE")
    import json
    json.dump({"table_mm": np.asarray(marker_xy).tolist(), "time": time.strftime("%Y-%m-%d %H:%M:%S")},
               open(mpos.PICKUP_PATH, "w"), indent=1)
    return plan, marker_xy


def trace_with_replanning(arm, cap, cam, plan, check_every_mm=20.0):
    """
    Streams plan.path_mm one waypoint at a time. Every check_every_mm of actual travel, grabs a
    fresh frame and runs tp.path_blocked() -- cheap, no re-detection, just "does anything in
    this frame hit the remaining path". Only when that trips does it call the heavier tp.replan()
    (re-detects obstacles, reroutes just the untraced part) and continue from there. A clear path
    never pauses for this check at all -- matches the "never stop-and-recompute unless something
    genuinely changed" design.

    Returns the final Plan actually traced (may differ from the one passed in, if it replanned).
    """
    path, kind = plan.path_mm, plan.kind
    arm.move_xy(*path[0], ac.Z_TRAVEL)    # hover above the start first

    traveled_since_check = 0.0
    i = 0
    while i < len(path):
        x_mm, y_mm = path[i]
        z = ac.Z_LIFT if kind[i] == tp.OVER else ac.Z_TRACE
        arm.move_xy(x_mm, y_mm, z)   # blocks until the firmware ACKs the move is actually done

        traveled_since_check += 0.0 if i == 0 else float(np.linalg.norm(path[i] - path[i - 1]))
        if traveled_since_check >= check_every_mm and i < len(path) - 1:
            traveled_since_check = 0.0
            ok, frame = cap.read()
            if ok and tp.path_blocked(plan, frame, cam, from_index=i):
                print(f"TRACE: new obstacle detected at waypoint {i}/{len(path)} -- replanning")
                plan = tp.replan(plan, (x_mm, y_mm), frame, cam)
                for w in plan.warnings:
                    print("replan warning:", w)
                path, kind = plan.path_mm, plan.kind
                i = 0
                continue
        i += 1

    arm.move_xy(*path[-1], ac.Z_TRAVEL)   # lift off the end
    return plan


def safe_park(arm):
    try:
        arm.set_gripper(closed=False)
        arm.home()
    except Exception as e:
        print("safe_park itself failed -- arm may need a manual reset:", e)


def run():
    print("INIT")
    cam = cc.load()
    self_checks(cam)

    print("loading marker model...")
    from ultralytics import YOLO
    model = YOLO(mpos.MODEL_PATH)
    cls_id = mpos.marker_class_id(model)
    correction = mpos.load_correction()

    cap = mpos.open_camera()
    state = "LOCATE"
    plan = marker_xy = None

    try:
        with ars.ArmSerial() as arm:
            while state != "DONE":
                if state == "LOCATE":
                    print("LOCATE: scanning paper edges, line, obstacle, and marker position")
                    plan, marker_xy = locate(cap, cam, model, cls_id, correction)
                    state = "PICK"

                elif state == "PICK":
                    print(f"PICK: marker at {np.round(marker_xy, 1)} mm")
                    pp.pick_marker(arm, marker_xy)
                    state = "TRACE"

                elif state == "TRACE":
                    n_over = int((plan.kind == tp.OVER).sum())
                    if n_over:
                        print(f"TRACE: plan includes {n_over} OVER waypoints -- "
                              f"line will have a gap where it clears the obstacle")
                    print("TRACE: streaming the plan, checking every "
                          f"{TRACE_CHECK_EVERY_MM:.0f}mm for a new obstacle")
                    plan = trace_with_replanning(arm, cap, cam, plan,
                                                  check_every_mm=TRACE_CHECK_EVERY_MM)
                    state = "RETURN"

                elif state == "RETURN":
                    print("RETURN: placing the marker back at its pickup position")
                    pp.place_marker(arm, marker_xy)
                    arm.home()
                    state = "DONE"

            print("DONE")

    except Exception as e:
        print("ERROR:", e)
        try:
            # The with-block above has already closed the original connection by the time we
            # get here (its __exit__ ran while the exception unwound through it), so reopening
            # is safe -- but note it resets the ESP32 (same as any fresh connect), so this is a
            # "get it to a safe position from scratch" park, not a resume of in-progress motion.
            with ars.ArmSerial() as arm:
                safe_park(arm)
        except Exception as e2:
            print("could not reach the arm to park it:", e2)
        raise
    finally:
        cap.release()


if __name__ == "__main__":
    run()