# Motion control — cmd_vel, obstacle detection, safe movement
import asyncio
import math
import time

from .config import (
    log, _ros2_cache, _ros2_lock, _estop_event,
    OBSTACLE_DISTANCE, OBSTACLE_WAIT_TIMEOUT,
    ROTATION_SCALE, DISTANCE_SCALE,
)
import ai_teammate_ros2_bridge.config as _cfg


LOW_BATTERY_PCT = 10  # refuse motion below this %


def _check_battery() -> tuple[bool, int]:
    """Check battery level. Returns (ok, percentage)."""
    try:
        import psutil
        bat = psutil.sensors_battery()
        pct = bat.percent if bat else 100
        return pct > LOW_BATTERY_PCT, pct
    except Exception:
        return True, 100


def publish_cmd_vel(linear_x: float, angular_z: float, duration: float = 1.0):
    """Publish cmd_vel via rclpy native publisher at 10Hz for duration (blocking).
    Returns 'ok', 'estop' (interrupted), 'low_battery', or 'error'."""
    ok, pct = _check_battery()
    if not ok:
        log.warning(f"Low battery ({pct}%) — motion refused")
        return "low_battery"
    from geometry_msgs.msg import Twist

    pub = _cfg._cmd_vel_pub
    if not pub:
        log.warning("ROS2 node not ready, cannot publish cmd_vel")
        return "error"

    _estop_event.clear()
    rate_hz = 10
    try:
        msg = Twist()
        msg.linear.x = float(linear_x)
        msg.angular.z = float(angular_z)

        interrupted = False
        if duration > 0:
            ticks = int(duration * rate_hz)
            for i in range(ticks):
                if _estop_event.is_set():
                    log.warning(f"E-STOP: interrupted at tick {i}/{ticks}")
                    interrupted = True
                    break
                pub.publish(msg)
                time.sleep(1.0 / rate_hz)
            if not interrupted:
                log.info(f"Published cmd_vel: linear={linear_x}, angular={angular_z} for {duration:.1f}s ({ticks} ticks)")
        else:
            pub.publish(msg)
            log.info(f"Published cmd_vel: linear={linear_x}, angular={angular_z}")

        # Always send STOP after movement
        stop = Twist()
        pub.publish(stop)
        log.info("Published cmd_vel STOP")

        return "estop" if interrupted else "ok"
    except Exception as e:
        log.error(f"cmd_vel publish error: {e}")
        return "error"


async def calibrate_rotation(speed: float = 0.5) -> dict:
    """Rotate 360 degrees and measure actual rotation via IMU yaw accumulation."""
    from geometry_msgs.msg import Twist
    pub = _cfg._cmd_vel_pub
    if not pub:
        return {"status": "error", "message": "ROS2 not ready"}

    # Safety: check minimum distance in all directions (LiDAR overall min)
    with _ros2_lock:
        scan = _ros2_cache.get("scan", {})
    min_dist = scan.get("min_distance")
    if min_dist is not None and min_dist < 0.3:
        return {"status": "error", "message": f"Too close to obstacle ({min_dist:.2f}m) — need 0.3m clearance in all directions"}

    # Apply current calibration scale to commanded rotation
    target_rad = 2 * math.pi * _cfg.ROTATION_SCALE
    duration = target_rad / speed
    rate_hz = 10
    ticks = int(duration * rate_hz)
    dt = 1.0 / rate_hz

    # Use IMU yaw (more reliable than odometry theta on robots without wheel encoders)
    accumulated_rad = 0.0
    with _ros2_lock:
        imu = _ros2_cache.get("imu", {})
    prev_yaw = imu.get("yaw", 0)
    log.info(f"calibrate_rotation: ticks={ticks}, prev_yaw={prev_yaw:.4f}, estop={_estop_event.is_set()}")

    msg = Twist()
    msg.angular.z = float(speed)

    # Force clear estop — replace event to avoid race conditions
    _estop_event.clear()
    estop_check = False  # disable estop during calibration (safety via LiDAR instead)
    for i in range(ticks):
        # Safety: check obstacle during rotation
        with _ros2_lock:
            rt_scan = _ros2_cache.get("scan", {})
        rt_min = rt_scan.get("min_distance")
        if rt_min is not None and rt_min < 0.2:
            log.warning(f"Calibration aborted — obstacle at {rt_min:.2f}m during rotation")
            break
        pub.publish(msg)
        await asyncio.sleep(dt)
        with _ros2_lock:
            imu = _ros2_cache.get("imu", {})
        yaw = imu.get("yaw", 0)
        delta = yaw - prev_yaw
        if delta > math.pi:
            delta -= 2 * math.pi
        elif delta < -math.pi:
            delta += 2 * math.pi
        accumulated_rad += abs(delta)
        if i % 20 == 0:
            log.info(f"  cal tick {i}/{ticks}: yaw={yaw:.4f} delta={delta:.4f} acc={math.degrees(accumulated_rad):.1f}°")
        prev_yaw = yaw

    # Stop
    stop = Twist()
    for _ in range(3):
        pub.publish(stop)

    actual_deg = math.degrees(accumulated_rad)
    scale = 360.0 / actual_deg if actual_deg > 10 else 1.0

    log.info(f"Rotation calibration: commanded=360°, actual={actual_deg:.1f}°, scale={scale:.4f}")
    return {
        "status": "ok",
        "commanded_deg": 360.0,
        "actual_deg": round(actual_deg, 2),
        "rotation_scale": round(scale, 4),
    }


async def calibrate_distance(target_m: float = 1.0, speed: float = 0.3) -> dict:
    """Move forward and measure actual distance via LiDAR front distance change."""
    from geometry_msgs.msg import Twist
    pub = _cfg._cmd_vel_pub
    if not pub:
        return {"status": "error", "message": "ROS2 not ready"}

    # Check obstacle — need wall/obstacle in front for LiDAR reference
    blocked, front_before = _check_front_obstacle(100.0)  # any distance
    if front_before is None:
        return {"status": "error", "message": "No LiDAR data — cannot calibrate distance"}
    if front_before < target_m + 0.5:
        return {"status": "error", "message": f"Obstacle too close ({front_before:.2f}m) — need {target_m + 0.5:.1f}m clearance"}

    # Move forward with obstacle checking (safe_move)
    calibrated_m = target_m * _cfg.DISTANCE_SCALE
    move_result = await _handle_safe_move(calibrated_m, speed, 0.2, 0.3)
    if move_result["status"] not in ("ok",):
        return {"status": "error", "message": f"Move failed: {move_result.get('status')}", "detail": move_result}

    await asyncio.sleep(0.5)  # settle

    # Measure front distance after move
    _, front_after = _check_front_obstacle(100.0)
    if front_after is None:
        return {"status": "error", "message": "LiDAR lost after move"}

    actual_m = front_before - front_after
    scale = target_m / actual_m if abs(actual_m) > 0.05 else 1.0

    # Move back (also safe)
    await _handle_safe_move(-calibrated_m, speed, 0.2, 0.3)

    log.info(f"Distance calibration: commanded={target_m}m, actual={actual_m:.3f}m, scale={scale:.4f}")
    return {
        "status": "ok" if result == "ok" else result,
        "commanded_m": target_m,
        "actual_m": round(actual_m, 3),
        "distance_scale": round(scale, 4),
        "front_before": round(front_before, 3),
        "front_after": round(front_after, 3),
    }


async def calibrate_full(max_rounds: int = 3, speed_rot: float = 0.5, speed_dist: float = 0.3) -> dict:
    """Multi-round calibration: rotate + distance, refine until converged."""
    import ai_teammate_ros2_bridge.config as cfg

    rounds = []
    for i in range(max_rounds):
        log.info(f"=== Calibration round {i+1}/{max_rounds} ===")

        # Rotation
        rot_result = await calibrate_rotation(speed_rot)
        rot_scale = rot_result.get("rotation_scale", 1.0)
        rot_actual = rot_result.get("actual_deg", 0)

        if 0.5 <= rot_scale <= 2.0 and rot_scale != 1.0:
            cfg.ROTATION_SCALE *= rot_scale  # compound correction
            # Clamp
            cfg.ROTATION_SCALE = max(0.5, min(2.0, cfg.ROTATION_SCALE))
            log.info(f"  Round {i+1} rotation: actual={rot_actual:.1f}°, compound_scale={cfg.ROTATION_SCALE:.4f}")

        await asyncio.sleep(1)

        # Distance (skip if no wall in front)
        dist_result = await calibrate_distance(1.0, speed_dist)
        dist_scale = dist_result.get("distance_scale", 1.0)
        dist_actual = dist_result.get("actual_m", 0)

        if dist_result.get("status") == "ok" and 0.5 <= dist_scale <= 2.0 and dist_scale != 1.0:
            cfg.DISTANCE_SCALE *= dist_scale
            cfg.DISTANCE_SCALE = max(0.5, min(2.0, cfg.DISTANCE_SCALE))
            log.info(f"  Round {i+1} distance: actual={dist_actual:.3f}m, compound_scale={cfg.DISTANCE_SCALE:.4f}")

        rounds.append({
            "round": i + 1,
            "rotation": {"actual_deg": rot_actual, "scale": rot_scale},
            "distance": {"actual_m": dist_actual, "scale": dist_scale, "status": dist_result.get("status")},
        })

        # Converged if both scales are close to 1.0 (within 2%)
        if abs(rot_scale - 1.0) < 0.02 and (dist_result.get("status") != "ok" or abs(dist_scale - 1.0) < 0.02):
            log.info(f"  Converged at round {i+1}")
            break

        await asyncio.sleep(1)

    return {
        "status": "ok",
        "rotation_scale": round(cfg.ROTATION_SCALE, 4),
        "distance_scale": round(cfg.DISTANCE_SCALE, 4),
        "rounds": rounds,
        "converged": len(rounds) < max_rounds or (abs(rounds[-1]["rotation"]["scale"] - 1.0) < 0.02),
    }


def _check_front_obstacle(threshold: float) -> tuple[bool, float | None]:
    """Check front obstacle from LiDAR cache. Returns (blocked, distance)."""
    with _ros2_lock:
        scan = _ros2_cache.get("scan", {})
    front = scan.get("front_min_distance")
    if front is None:
        return False, None
    return front < threshold, front


def _wait_for_clear(threshold: float) -> str:
    """Block until obstacle clears or E-STOP or timeout. Returns 'clear', 'estop', 'timeout'."""
    start = time.time()
    log.info(f"Waiting for obstacle to clear (threshold={threshold}m, timeout={OBSTACLE_WAIT_TIMEOUT}s)...")
    while True:
        if _estop_event.is_set():
            return "estop"
        if time.time() - start > OBSTACLE_WAIT_TIMEOUT:
            return "timeout"
        blocked, dist = _check_front_obstacle(threshold)
        if not blocked:
            log.info(f"Obstacle cleared (front={dist}m)")
            return "clear"
        time.sleep(0.2)


async def _handle_safe_move(distance: float, speed: float, step: float, threshold: float) -> dict:
    """Move with periodic obstacle checks every `step` meters."""
    loop = asyncio.get_event_loop()
    direction = 1.0 if distance > 0 else -1.0
    remaining = abs(distance)
    moved = 0.0

    while remaining > 0.01:
        if _estop_event.is_set():
            return {"status": "estop", "moved": round(moved, 3), "remaining": round(remaining, 3)}

        if direction > 0:
            blocked, front_dist = _check_front_obstacle(threshold)
            if blocked:
                log.warning(f"SAFE_MOVE: obstacle at {front_dist:.2f}m after {moved:.2f}m — waiting for clear")
                wait_result = await loop.run_in_executor(None, _wait_for_clear, threshold)
                if wait_result == "estop":
                    return {"status": "estop", "moved": round(moved, 3), "remaining": round(remaining, 3)}
                if wait_result == "timeout":
                    return {"status": "obstacle_timeout", "moved": round(moved, 3), "remaining": round(remaining, 3),
                            "front_distance": front_dist, "threshold": threshold,
                            "message": f"Obstacle not cleared after {OBSTACLE_WAIT_TIMEOUT}s"}

        chunk = min(step, remaining)
        duration = chunk / speed
        result = await loop.run_in_executor(None, publish_cmd_vel, direction * speed, 0.0, duration)

        if result == "estop":
            return {"status": "estop", "moved": round(moved + chunk, 3), "remaining": round(remaining - chunk, 3)}

        moved += chunk
        remaining -= chunk

    return {"status": "ok", "moved": round(moved, 3), "distance": distance, "speed": speed}


async def _handle_safe_navigate_path(waypoints: list, speed: float, step: float, threshold: float) -> dict:
    """Navigate waypoints with obstacle checks between segments."""
    loop = asyncio.get_event_loop()
    results = []
    prev = waypoints[0]

    for i, wp in enumerate(waypoints[1:], 1):
        if _estop_event.is_set():
            results.append({"wp": i, "status": "estop"})
            break

        dx = wp[0] - prev[0]
        dy = wp[1] - prev[1]
        distance = math.sqrt(dx * dx + dy * dy)

        if distance > 0.01:
            target_angle = math.atan2(dy, dx)
            angle_deg = math.degrees(target_angle)
            if abs(angle_deg) > 1:
                turn_duration = abs(math.radians(angle_deg)) / 0.5
                result = await loop.run_in_executor(None, publish_cmd_vel, 0.0, (0.5 if angle_deg > 0 else -0.5), turn_duration)
                if result == "estop":
                    results.append({"wp": i, "status": "estop"})
                    break

            move_result = await _handle_safe_move(distance, speed, step, threshold)
            results.append({"wp": i, "x": wp[0], "y": wp[1], **move_result})

            if move_result["status"] in ("obstacle", "estop"):
                break
        else:
            results.append({"wp": i, "status": "skipped"})

        prev = wp

    completed = all(r.get("status") == "ok" for r in results)
    return {"status": "ok" if completed else "partial", "waypoints_executed": len(results),
            "total": len(waypoints) - 1, "results": results}


async def _handle_navigate_path(params: dict) -> dict:
    """Navigate through a sequence of (x, y) waypoints using cmd_vel dead reckoning."""
    waypoints = params.get("waypoints", [])
    speed = float(params.get("speed", 0.3))
    if not waypoints or len(waypoints) < 2:
        return {"status": "error", "message": "at least 2 waypoints required"}

    loop = asyncio.get_event_loop()
    results = []
    prev = waypoints[0]

    for i, wp in enumerate(waypoints[1:], 1):
        if _estop_event.is_set():
            results.append({"wp": i, "status": "estop"})
            break

        dx = wp[0] - prev[0]
        dy = wp[1] - prev[1]
        distance = math.sqrt(dx * dx + dy * dy)

        if distance > 0.01:
            target_angle = math.atan2(dy, dx)
            angle_deg = math.degrees(target_angle)
            if abs(angle_deg) > 1:
                turn_duration = abs(math.radians(angle_deg)) / 0.5
                result = await loop.run_in_executor(None, publish_cmd_vel, 0.0, (0.5 if angle_deg > 0 else -0.5), turn_duration)
                if result == "estop":
                    results.append({"wp": i, "status": "estop"})
                    break

            duration = distance / speed
            result = await loop.run_in_executor(None, publish_cmd_vel, speed, 0.0, duration)
            results.append({"wp": i, "x": wp[0], "y": wp[1], "distance": round(distance, 3), "status": result})
        else:
            results.append({"wp": i, "x": wp[0], "y": wp[1], "status": "skipped"})

        prev = wp

    completed = all(r.get("status") == "ok" for r in results)
    return {"status": "ok" if completed else "partial", "waypoints_executed": len(results), "total": len(waypoints) - 1, "results": results}
