#!/usr/bin/env python3
"""
Helios Laser Serial Communication Diagnostic Tool
Helps troubleshoot communication issues with the Helios laser.
"""

import serial
import time
import sys

def test_port(port, baudrate=9600):
    """Test basic communication on a serial port."""
    print(f"\n{'='*60}")
    print(f"Testing {port} at {baudrate} baud")
    print(f"{'='*60}")

    try:
        ser = serial.Serial(
            port=port,
            baudrate=baudrate,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=1.0
        )
        print(f"✓ Port opened successfully")
        time.sleep(0.1)

        # Try to query the controller serial number
        print("\nSending: 'SN?'")
        ser.write(b'SN?\r')
        time.sleep(0.5)

        response = ser.readline().decode('ascii', errors='replace').strip()
        print(f"Response: '{response}'")

        if response:
            print(f"✓ Got response: {response}")
            return True, response
        else:
            print(f"✗ No response received")

        # Try head serial number
        print("\nSending: 'HSN?'")
        ser.write(b'HSN?\r')
        time.sleep(0.5)

        response = ser.readline().decode('ascii', errors='replace').strip()
        print(f"Response: '{response}'")

        if response:
            print(f"✓ Got response: {response}")
            ser.close()
            return True, response
        else:
            print(f"✗ No response received")

        # Try laser enable status
        print("\nSending: 'LE?'")
        ser.write(b'LE?\r')
        time.sleep(0.5)

        response = ser.readline().decode('ascii', errors='replace').strip()
        print(f"Response: '{response}'")

        if response:
            print(f"✓ Got response: {response}")
            ser.close()
            return True, response
        else:
            print(f"✗ No response received")

        ser.close()
        return False, "No response to any query"

    except Exception as e:
        print(f"✗ Error: {e}")
        return False, str(e)


def test_raw_communication(port, baudrate=9600):
    """Test raw serial communication and display hex."""
    print(f"\n{'='*60}")
    print(f"Raw Communication Test: {port} at {baudrate} baud")
    print(f"{'='*60}")

    try:
        ser = serial.Serial(
            port=port,
            baudrate=baudrate,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=2.0
        )
        print(f"✓ Port opened successfully")
        time.sleep(0.2)

        # Send a simple query
        command = b'SN?\r'
        print(f"\nSending command (hex): {command.hex()}")
        print(f"Sending command (ascii): {command}")

        ser.write(command)
        time.sleep(0.5)

        # Read response byte by byte
        response = b''
        while True:
            byte = ser.read(1)
            if not byte:
                break
            response += byte
            if byte == b'\n' or byte == b'\r':
                break

        print(f"\nRaw response (hex): {response.hex()}")
        print(f"Raw response (ascii): {response}")
        print(f"Response length: {len(response)} bytes")

        # Check for common issues
        if not response:
            print("✗ No response - device may not be responding or wrong baud rate")
        elif response == b'\r' or response == b'\n':
            print("⚠ Only got line terminator - device may be echoing but not responding to command")
        else:
            print("✓ Got a response!")

        ser.close()
        return True

    except Exception as e:
        print(f"✗ Error: {e}")
        return False


def test_echo(port, baudrate=9600):
    """Test if the device echoes commands back."""
    print(f"\n{'='*60}")
    print(f"Echo Test: {port} at {baudrate} baud")
    print(f"{'='*60}")

    try:
        ser = serial.Serial(
            port=port,
            baudrate=baudrate,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=1.0
        )

        # Send a test character
        test_char = b'T'
        print(f"Sending test character: {test_char}")
        ser.write(test_char)
        time.sleep(0.1)

        echo = ser.read(1)
        if echo == test_char:
            print(f"✓ Device echoes input")
        elif echo:
            print(f"⚠ Device sent something but not the same: {echo}")
        else:
            print(f"✗ No echo")

        ser.close()
        return True

    except Exception as e:
        print(f"✗ Error: {e}")
        return False


def main():
    """Run diagnostic tests."""
    port = "/dev/ttyUSB2"

    if len(sys.argv) > 1:
        port = sys.argv[1]

    print(f"\n{'#'*60}")
    print(f"# Helios Laser Serial Diagnostic Tool")
    print(f"# Testing port: {port}")
    print(f"{'#'*60}")

    # Test standard baud rate
    success, response = test_port(port, 9600)

    if not success:
        print("\n" + "="*60)
        print("Standard baud rate (9600) failed. Trying alternatives...")
        print("="*60)

        # Try other common baud rates
        for baudrate in [115200, 19200, 4800, 2400]:
            success, response = test_port(port, baudrate)
            if success:
                print(f"\n✓ SUCCESS! Device responds at {baudrate} baud")
                break
    else:
        print(f"\n✓ SUCCESS! Device responds at 9600 baud")

    # Run additional diagnostics
    print("\n")
    test_raw_communication(port, 9600)

    print("\n")
    test_echo(port, 9600)

    print(f"\n{'#'*60}")
    print("# Diagnostic Tests Complete")
    print(f"{'#'*60}\n")


if __name__ == "__main__":
    main()