"""
BBD202 Stage Diagnostic Tests

Diagnostic tests for the Thorlabs MLS203/BBD202 motion controller to:
1. Check if axes are properly homed
2. Move in a 15mm square pattern while outputting position data
3. Detect stage hangs through position monitoring
4. Output response information for debugging
"""

import sys
import time
from typing import Optional
from hardware.bbd202 import MotionController, MotorStatusBits


def format_status_bits(bits: Optional[int]) -> str:
    """Format status bits as a readable string with key flags."""
    if bits is None:
        return "None (no status received)"

    flags = []
    status = MotorStatusBits(bits)

    # Key flags to check
    if MotorStatusBits.HOMED in status:
        flags.append("HOMED")
    if MotorStatusBits.HOMING in status:
        flags.append("HOMING")
    if MotorStatusBits.ENABLED in status:
        flags.append("ENABLED")
    if MotorStatusBits.SETTLED in status:
        flags.append("SETTLED")
    if MotorStatusBits.TRACKING in status:
        flags.append("TRACKING")
    if MotorStatusBits.INMOTIONCW in status:
        flags.append("INMOTIONCW")
    if MotorStatusBits.INMOTIONCCW in status:
        flags.append("INMOTIONCCW")
    if MotorStatusBits.POWEROK in status:
        flags.append("POWEROK")
    if MotorStatusBits.ERROR in status:
        flags.append("ERROR")
    if MotorStatusBits.POSITIONERROR in status:
        flags.append("POSITIONERROR")
    if MotorStatusBits.OVERTEMP in status:
        flags.append("OVERTEMP")
    if MotorStatusBits.COMMUTATIONERROR in status:
        flags.append("COMMUTATIONERROR")

    flag_str = ", ".join(flags) if flags else "NO_FLAGS"
    return f"0x{bits:08X} [{flag_str}]"


def test_homing_status(mc: MotionController) -> bool:
    """
    Test 1: Check if both axes are properly homed.

    Returns True if both axes are homed, False otherwise.
    """
    print("\n" + "="*70)
    print("TEST 1: Checking Homing Status")
    print("="*70)

    # Request fresh status from both axes
    print("\nRequesting status update from both axes...")
    mc.request_status_update(mc.DEST_X_AXIS)
    mc.request_status_update(mc.DEST_Y_AXIS)
    time.sleep(0.2)  # Allow time for response

    # Also poll status to ensure we have latest
    mc.poll_status()
    time.sleep(0.1)

    # Check X-axis
    x_bits = mc.status_bits_x
    x_homed = mc.is_homed_x
    print(f"\nX-Axis Status:")
    print(f"  Raw status bits: {format_status_bits(x_bits)}")
    print(f"  is_homed_x:      {x_homed}")

    # Check Y-axis
    y_bits = mc.status_bits_y
    y_homed = mc.is_homed_y
    print(f"\nY-Axis Status:")
    print(f"  Raw status bits: {format_status_bits(y_bits)}")
    print(f"  is_homed_y:      {y_homed}")

    # Check for errors
    last_err = mc.last_error
    if last_err:
        print(f"\n** Last Error: {last_err}")

    error_check = mc.check_for_errors()
    if error_check:
        print(f"** Error condition detected: {error_check}")

    # Summary
    print(f"\n--- Homing Status Summary ---")
    if x_homed and y_homed:
        print("PASS: Both axes are homed")
        return True
    else:
        print("FAIL: One or both axes are NOT homed")
        if not x_homed:
            print("  -> X-axis needs homing")
        if not y_homed:
            print("  -> Y-axis needs homing")
        return False


def test_square_pattern(mc: MotionController, side_length: float = 15.0) -> bool:
    """
    Test 2: Move the stage in a square pattern.

    Moves in a 15mm square starting from current position, outputting
    position data at each step to detect hangs.

    Pattern: Start -> +X -> +Y -> -X -> -Y (back to start)

    Uses the new wait_until_settled() method for reliable motion detection.
    """
    print("\n" + "="*70)
    print(f"TEST 2: Square Pattern Movement ({side_length}mm)")
    print("="*70)
    print("Using wait_until_settled() for motion detection")

    # Get starting position
    mc.poll_positions()
    time.sleep(0.1)
    start_x = mc.position_x
    start_y = mc.position_y

    if start_x is None or start_y is None:
        print("ERROR: Could not get initial position")
        return False

    print(f"\nStarting position: X={start_x:.4f}mm, Y={start_y:.4f}mm")

    # Define the square corners (relative to start)
    corners = [
        (start_x + side_length, start_y),               # Move +X
        (start_x + side_length, start_y + side_length), # Move +Y
        (start_x, start_y + side_length),               # Move -X
        (start_x, start_y),                             # Move -Y (back to start)
    ]

    corner_names = ["+X edge", "+Y edge (diagonal)", "-X edge", "Return to start"]

    all_passed = True

    for i, (target_x, target_y) in enumerate(corners):
        print(f"\n--- Move {i+1}/4: {corner_names[i]} ---")
        print(f"Target: X={target_x:.4f}mm, Y={target_y:.4f}mm")

        # Clear any previous error
        mc.clear_last_error()

        # Record start time
        move_start = time.time()

        # Issue move command (non-blocking)
        mc.move_to_fast(x=target_x, y=target_y)

        # Use poll_until_idle() in a loop to show progress
        print(f"\n{'Time':>8} {'X_pos':>12} {'Y_pos':>12} {'X_err':>10} {'Y_err':>10} {'Pending'}")
        print("-" * 70)

        sample_count = 0
        while not mc.poll_until_idle(tolerance=0.005, timeout=0.05):
            elapsed = time.time() - move_start

            curr_x = mc.position_x
            curr_y = mc.position_y

            if curr_x is not None and curr_y is not None:
                x_err = curr_x - target_x
                y_err = curr_y - target_y

                # Print every 5th sample
                if sample_count % 5 == 0:
                    pending = mc.get_pending_targets()
                    pending_str = ", ".join([f"{'X' if d==0x21 else 'Y'}" for d in pending.keys()])
                    print(f"{elapsed:>7.3f}s {curr_x:>12.4f} {curr_y:>12.4f} {x_err:>+10.4f} {y_err:>+10.4f} {pending_str:>8}")

            sample_count += 1

            # Timeout check
            if elapsed > 30.0:
                print(f"\n** TIMEOUT after 30s - stage may be hung!")
                all_passed = False
                mc.clear_pending_moves()
                break

        # Move complete - show final status
        elapsed = time.time() - move_start
        curr_x = mc.position_x
        curr_y = mc.position_y

        if curr_x is not None and curr_y is not None:
            x_err = curr_x - target_x
            y_err = curr_y - target_y
            print(f"{elapsed:>7.3f}s {curr_x:>12.4f} {curr_y:>12.4f} {x_err:>+10.4f} {y_err:>+10.4f} ** DONE **")
            print(f"\nMove completed in {elapsed:.3f}s")
            print(f"Final position: X={curr_x:.4f}mm, Y={curr_y:.4f}mm")
            print(f"Position error: X={abs(x_err)*1000:.1f}um, Y={abs(y_err)*1000:.1f}um")

        # Check for errors after move
        err = mc.check_for_errors()
        if err:
            print(f"** Error after move: {err}")
            all_passed = False

        last_err = mc.last_error
        if last_err:
            print(f"** Last error: {last_err}")
            all_passed = False

        # Short delay between moves
        time.sleep(0.1)

    # Final summary
    print(f"\n--- Square Pattern Summary ---")
    mc.poll_positions()
    time.sleep(0.1)
    final_x = mc.position_x
    final_y = mc.position_y
    if final_x is not None and final_y is not None:
        total_x_err = abs(final_x - start_x)
        total_y_err = abs(final_y - start_y)
        print(f"Final position:  X={final_x:.4f}mm, Y={final_y:.4f}mm")
        print(f"Start position:  X={start_x:.4f}mm, Y={start_y:.4f}mm")
        print(f"Return error:    X={total_x_err*1000:.1f}um, Y={total_y_err*1000:.1f}um")

        if total_x_err > 0.05 or total_y_err > 0.05:
            print("FAIL: Did not return to start position accurately (>50um)")
            all_passed = False

    if all_passed:
        print("PASS: Square pattern completed successfully")
    else:
        print("FAIL: Issues detected during square pattern")

    return all_passed


def print_diagnostic_info(mc: MotionController):
    """Print comprehensive diagnostic information about the controller."""
    print("\n" + "="*70)
    print("CONTROLLER DIAGNOSTIC INFO")
    print("="*70)

    try:
        hw_info = mc.get_hw_info()
        print(f"\nHardware Info:")
        print(f"  Serial Number: {hw_info['serial_number']}")
        print(f"  Model:         {hw_info['model']}")
        print(f"  Firmware:      {hw_info['firmware_version']}")
        print(f"  HW Version:    {hw_info['hw_version']}")
        print(f"  Channels:      {hw_info['num_channels']}")
        print(f"  Notes:         {hw_info['notes']}")
    except Exception as e:
        print(f"  Error getting hardware info: {e}")

    # Request and display velocity params
    print(f"\nVelocity Parameters:")
    mc.send_command(0x0414, param1=0x01, dest=mc.DEST_X_AXIS)  # REQ_VELPARAMS
    mc.send_command(0x0414, param1=0x01, dest=mc.DEST_Y_AXIS)
    time.sleep(0.1)

    x_vel = mc.velocity_params_x
    y_vel = mc.velocity_params_y
    if x_vel:
        print(f"  X-Axis: max_vel={x_vel['max_velocity']:.2f}mm/s, accel={x_vel['acceleration']:.2f}mm/s²")
    if y_vel:
        print(f"  Y-Axis: max_vel={y_vel['max_velocity']:.2f}mm/s, accel={y_vel['acceleration']:.2f}mm/s²")


def test_bay_and_channel_status(mc: MotionController) -> dict:
    """
    Test bay occupancy and channel enable states.

    Returns dict with diagnostic info about each axis.
    """
    print("\n" + "="*70)
    print("BAY AND CHANNEL DIAGNOSTICS")
    print("="*70)

    results = {'x': {}, 'y': {}}

    # Query bay status using MGMSG_RACK_REQ_BAYUSED (0x0060)
    print("\n--- Querying Bay Status (RACK_REQ_BAYUSED 0x0060) ---")
    print("Sending to controller (0x11)...")

    # Send bay request to controller
    mc.send_command(0x0060, param1=0x00, param2=0x00, dest=0x11, source=0x01)
    time.sleep(0.2)

    # Check for response in queue
    msgs = mc.get_all_messages()
    bay_response = None
    for msg in msgs:
        print(f"  Received: msg_id=0x{msg.msg_id:04X}, source=0x{msg.source:02X}, "
              f"param1=0x{msg.param1:02X}, param2=0x{msg.param2:02X}, data={msg.data.hex() if msg.data else 'none'}")
        if msg.msg_id == 0x0061:  # RACK_GET_BAYUSED
            bay_response = msg

    if bay_response:
        # param1 contains bay_ident (which bay), param2 contains state
        print(f"\nBay status response: bay={bay_response.param1}, state=0x{bay_response.param2:02X}")
    else:
        print("  No RACK_GET_BAYUSED response received")

    # Check channel enable state for X-axis (0x21)
    print("\n--- X-Axis (Bay 1 / dest=0x21) Channel Status ---")
    print(f"  Destination address: 0x{mc.DEST_X_AXIS:02X}")

    try:
        x_enabled = mc.get_channel_enable_state(mc.DEST_X_AXIS, timeout=2.0)
        print(f"  Channel enabled: {x_enabled}")
        results['x']['enabled'] = x_enabled
    except Exception as e:
        print(f"  Error querying channel state: {e}")
        results['x']['enabled'] = None

    # Get X position to verify communication
    x_pos = mc.get_position(mc.DEST_X_AXIS, timeout=2.0)
    print(f"  Current position: {x_pos:.4f}mm" if x_pos is not None else "  Position query failed!")
    results['x']['position'] = x_pos
    results['x']['communicating'] = x_pos is not None

    # Check channel enable state for Y-axis (0x22)
    print("\n--- Y-Axis (Bay 2 / dest=0x22) Channel Status ---")
    print(f"  Destination address: 0x{mc.DEST_Y_AXIS:02X}")

    try:
        y_enabled = mc.get_channel_enable_state(mc.DEST_Y_AXIS, timeout=2.0)
        print(f"  Channel enabled: {y_enabled}")
        results['y']['enabled'] = y_enabled
    except Exception as e:
        print(f"  Error querying channel state: {e}")
        results['y']['enabled'] = None

    # Get Y position to verify communication
    y_pos = mc.get_position(mc.DEST_Y_AXIS, timeout=2.0)
    print(f"  Current position: {y_pos:.4f}mm" if y_pos is not None else "  Position query failed!")
    results['y']['position'] = y_pos
    results['y']['communicating'] = y_pos is not None

    return results


def test_y_axis_move_verbose(mc: MotionController) -> bool:
    """
    Test Y-axis movement using the new wait_until_settled() method.

    Sends a small Y-axis move and monitors progress.
    """
    print("\n" + "="*70)
    print("Y-AXIS VERBOSE MOVE TEST")
    print("="*70)
    print("Using wait_until_settled() for motion detection")

    # Get current Y position
    mc.poll_positions()
    time.sleep(0.1)
    start_y = mc.position_y

    if start_y is None:
        print("ERROR: Cannot get Y position")
        return False

    print(f"\nCurrent Y position: {start_y:.4f}mm")

    # Target: move 5mm in Y
    target_y = start_y + 5.0
    print(f"Target Y position:  {target_y:.4f}mm")

    # Issue move command using move_to_fast (which tracks pending moves)
    print("\n--- Issuing Y-axis move via move_to_fast() ---")
    mc.move_to_fast(y=target_y)

    # Show pending moves
    pending = mc.get_pending_targets()
    print(f"  Pending moves: {pending}")
    print(f"  is_move_pending(): {mc.is_move_pending()}")

    # Monitor using poll_until_idle
    print("\n--- Monitoring Y-axis with poll_until_idle() ---")
    print(f"{'Time':>6} {'Y_pos':>10} {'Y_err':>10} {'Pending'}")
    print("-" * 45)

    start_time = time.time()
    movement_detected = False

    while not mc.poll_until_idle(tolerance=0.005, timeout=0.05):
        elapsed = time.time() - start_time
        curr_y = mc.position_y

        if curr_y is not None:
            y_err = curr_y - target_y
            pending = "Y" if mc.is_move_pending() else "-"
            print(f"{elapsed:>5.2f}s {curr_y:>10.4f} {y_err:>+10.4f} {pending:>8}")

            if abs(curr_y - start_y) > 0.01:
                movement_detected = True

        if elapsed > 10.0:
            print("\n** TIMEOUT after 10s **")
            mc.clear_pending_moves()
            break

    # Final status
    elapsed = time.time() - start_time
    final_y = mc.position_y
    print(f"\n--- Y-Axis Move Summary ---")
    print(f"Start position:  {start_y:.4f}mm")
    print(f"Target position: {target_y:.4f}mm")
    print(f"Final position:  {final_y:.4f}mm" if final_y else "Final position: unknown")
    print(f"Move time:       {elapsed:.3f}s")

    if final_y is not None:
        distance_moved = abs(final_y - start_y)
        print(f"Distance moved:  {distance_moved:.4f}mm")

        if distance_moved < 0.01:
            print("\n** FAIL: Y-axis did not move at all! **")
            print("Possible causes:")
            print("  1. Y-axis channel is disabled")
            print("  2. Y-axis motor driver issue")
            print("  3. Command not reaching Y-axis bay")
            print("  4. Mechanical obstruction")
            return False
        elif abs(final_y - target_y) < 0.05:
            print("\n** PASS: Y-axis moved to target **")
            return True
        else:
            print(f"\n** PARTIAL: Y-axis moved but not to target **")
            return False

    return movement_detected


def test_enable_and_move_y(mc: MotionController) -> bool:
    """
    Explicitly enable Y-axis channel and attempt movement.
    """
    print("\n" + "="*70)
    print("ENABLE Y-AXIS AND MOVE TEST")
    print("="*70)

    # First, explicitly enable Y-axis channel
    print("\n--- Enabling Y-axis channel (MOD_SET_CHANENABLESTATE) ---")
    print(f"  Sending to dest=0x{mc.DEST_Y_AXIS:02X}, enable=True")

    try:
        mc.set_channel_enable_state(mc.DEST_Y_AXIS, enabled=True)
        time.sleep(0.2)
        print("  Enable command sent")
    except Exception as e:
        print(f"  Error sending enable: {e}")

    # Verify it's enabled
    try:
        y_enabled = mc.get_channel_enable_state(mc.DEST_Y_AXIS, timeout=2.0)
        print(f"  Channel enabled state: {y_enabled}")
    except Exception as e:
        print(f"  Error checking state: {e}")

    # Request fresh status
    mc.request_status_update(mc.DEST_Y_AXIS)
    time.sleep(0.2)

    y_bits = mc.status_bits_y
    print(f"  Status bits: {format_status_bits(y_bits)}")

    # Now try the verbose move test
    return test_y_axis_move_verbose(mc)


def main():
    print("="*70)
    print("BBD202 Stage Diagnostic Test")
    print("="*70)

    mc = None
    try:
        # Connect to controller
        print("\nConnecting to BBD202 controller...")
        mc = MotionController()
        mc.connect(enable_updates=True)
        print("Connected!")

        # Print diagnostic info
        print_diagnostic_info(mc)

        # Run bay and channel diagnostics
        channel_results = test_bay_and_channel_status(mc)

        # Run homing status test
        homed = test_homing_status(mc)

        if not homed:
            print("\n" + "-"*70)
            response = input("Axes not homed. Home now? (y/n): ").strip().lower()
            if response == 'y':
                print("\nHoming X-axis...")
                x_result = mc.home_x_axis(timeout=30.0)
                print(f"X-axis homing: {'SUCCESS' if x_result else 'FAILED'}")

                print("\nHoming Y-axis...")
                y_result = mc.home_y_axis(timeout=30.0)
                print(f"Y-axis homing: {'SUCCESS' if y_result else 'FAILED'}")

                if not (x_result and y_result):
                    print("\nHoming failed. Cannot continue with movement test.")
                    return 1

                # Re-check homing status
                test_homing_status(mc)
            else:
                print("\nSkipping movement tests (axes not homed)")
                return 1

        # Run Y-axis verbose test first to diagnose the issue
        y_axis_ok = test_enable_and_move_y(mc)

        # If Y-axis failed, skip the square pattern
        if not y_axis_ok:
            print("\n** Y-axis movement failed - skipping square pattern test **")
            square_ok = False
        else:
            # Run square pattern test
            square_ok = test_square_pattern(mc, side_length=15.0)

        # Final summary
        print("\n" + "="*70)
        print("DIAGNOSTIC SUMMARY")
        print("="*70)
        print(f"X-Axis Channel:  {'OK' if channel_results['x'].get('communicating') else 'FAIL'}")
        print(f"Y-Axis Channel:  {'OK' if channel_results['y'].get('communicating') else 'FAIL'}")
        print(f"Homing Status:   {'PASS' if homed else 'FAIL'}")
        print(f"Y-Axis Move:     {'PASS' if y_axis_ok else 'FAIL'}")
        print(f"Square Pattern:  {'PASS' if square_ok else 'FAIL/SKIPPED'}")

        if homed and y_axis_ok and square_ok:
            print("\nAll tests PASSED")
            return 0
        else:
            print("\nSome tests FAILED - see details above")
            return 1

    except KeyboardInterrupt:
        print("\n\nTest interrupted by user")
        return 1
    except Exception as e:
        print(f"\nError during test: {e}")
        import traceback
        traceback.print_exc()
        return 1
    finally:
        if mc is not None:
            print("\nDisconnecting from controller...")
            mc.disconnect()
            print("Disconnected")


if __name__ == "__main__":
    sys.exit(main())