import serial
import argparse
import sys
import time

def handle_comm(ser, cmd):
    ser.reset_input_buffer()
    ser.write((cmd + "\n").encode())
    time.sleep(0.1)
    start = time.time()
    while time.time() - start < 1.0:
        line = ser.readline().decode(errors='replace').strip()
        if line and any(x in line for x in ["ACK", "PONG", "ERR", "READY"]):
            return line
    return "[No Response]"

def main():
    parser = argparse.ArgumentParser(description='ESP32-P4 Waveform Gen CLI')
    parser.add_argument('--port', default='/dev/ttyACM0', help='Serial port')
    parser.add_argument('--interactive', action='store_true', help='Interactive mode')
    
    # State Commands
    parser.add_argument('--enable', action='store_true', help='Enable power stage')
    parser.add_argument('--disable', action='store_true', help='Force both outputs LOW')
    parser.add_argument('--on', action='store_true', help='Manual ON (A=0, B=1)')
    parser.add_argument('--off', action='store_true', help='Manual OFF (A=1, B=0)')
    parser.add_argument('--mod-on', action='store_true', help='Start modulation')
    parser.add_argument('--mod-off', action='store_true', help='Stop modulation (return to OFF)')
    parser.add_argument('--freq', type=int, help='Set frequency (10-5000 Hz)')
    parser.add_argument('--ping', action='store_true')

    args = parser.parse_args()

    try:
        ser = serial.Serial()
        ser.port = args.port
        ser.baudrate = 115200
        ser.timeout = 0.5
        ser.dtr = False
        ser.rts = False
        ser.open()
        
        if args.interactive:
            print(f"Connected to {args.port}.")
            print("Commands: ENABLE, DISABLE, ON, OFF, MOD_ON, MOD_OFF, F <hz>, PING, exit")
            while True:
                try:
                    cmd = input(">> ").strip()
                    if cmd.lower() in ['exit', 'quit']: break
                    if not cmd: continue
                    print(f"Response: {handle_comm(ser, cmd)}")
                except KeyboardInterrupt:
                    break
            ser.close()
            return

        # Determine which command to send
        cmd = None
        if args.ping: cmd = "PING"
        elif args.enable: cmd = "ENABLE"
        elif args.disable: cmd = "DISABLE"
        elif args.on: cmd = "ON"
        elif args.off: cmd = "OFF"
        elif args.mod_on: cmd = "MOD_ON"
        elif args.mod_off: cmd = "MOD_OFF"
        elif args.freq is not None: cmd = f"F {args.freq}"

        if cmd:
            print(f"Sending: {cmd}")
            print(f"Response: {handle_comm(ser, cmd)}")
        else:
            print("Error: No action specified. Use --help for options or --interactive.")

        ser.close()

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

if __name__ == '__main__':
    main()