controller.go

// Copyright (c) 2020–2024 The prologix developers. All rights reserved.
// Project site: https://github.com/gotmc/prologix
// Use of this source code is governed by a MIT-style license that
// can be found in the LICENSE.txt file for the project.

package prologix

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"strings"
)

// ContextReader is an optional interface that an io.ReadWriter can implement
// to support context-aware reads.
type ContextReader interface {
	ReadContext(ctx context.Context, p []byte) (n int, err error)
}

// ContextWriter is an optional interface that an io.ReadWriter can implement
// to support context-aware writes.
type ContextWriter interface {
	WriteContext(ctx context.Context, p []byte) (n int, err error)
}

// Controller models a GPIB controller-in-charge.
type Controller struct {
	rw               io.ReadWriter
	reader           *bufio.Reader
	primaryAddr      int
	hasSecondaryAddr bool
	secondaryAddr    int
	auto             bool
	eoi              bool
	usbTerm          byte
	eotChar          byte
	debug            bool // if true, print controller commands before sending. Set via WithDebug().
	ar488            bool // compatibility with Arduino AR488 - see WithAR488 documentation for details.
}

// ControllerOption applies an option to the controller.
type ControllerOption func(*Controller)

// NewController creates a GPIB controller-in-charge at the given address using
// the given Prologix driver, which can either be a Virtual COM Port (VCP), USB
// direct, or Ethernet. Enable clear to send the Selected Device Clear (SDC)
// message to the GPIB address. Optionally controller configuration can be
// included using a ControllerOption.
func NewController(
	rw io.ReadWriter,
	addr int,
	clear bool,
	opts ...ControllerOption,
) (*Controller, error) {
	c := Controller{
		rw:               rw,
		reader:           bufio.NewReader(rw),
		primaryAddr:      addr,
		hasSecondaryAddr: false,
		auto:             false,
		eoi:              true,
		usbTerm:          '\n',
		eotChar:          '\n',
	}

	// Apply options using the functional option pattern.
	for _, opt := range opts {
		opt(&c)
	}

	// Verify validate primary address.
	if !isPrimaryAddressValid(c.primaryAddr) {
		return nil, fmt.Errorf("invalid primary address %d (must by 0-30)", c.primaryAddr)
	}

	// Configure the Prologix GPIB controller.
	addrCmd := fmt.Sprintf("addr %d", c.primaryAddr)
	if c.hasSecondaryAddr {
		if !isSecondaryAddressValid(c.secondaryAddr) {
			return nil, fmt.Errorf("invalid secondary address %d (must be 96-126)", c.secondaryAddr)
		}
		addrCmd = fmt.Sprintf("addr %d %d", c.primaryAddr, c.secondaryAddr)
	}
	eotCharCmd := fmt.Sprintf("eot_char %d", c.eotChar)
	cmds := []string{}
	if !c.ar488 {
		cmds = append(cmds,
			"verbose 0", // turn off verbosity if on
			"savecfg 0", // Disable saving of configuration parameters in EPROM
		)
	}
	cmds = append(cmds,
		addrCmd,           // Set the primary address.
		"mode 1",          // Switch to controller mode.
		"auto 0",          // Turn off read-after-write and address instrument to listen.
		"eoi 1",           // Enable EOI assertion with last character.
		"eos 0",           // Set GPIB termination.
		"read_tmo_ms 500", // Set the read timeout to 500 ms.
		eotCharCmd,        // Set the EOT char
		"eot_enable 1",    // Append character when EOI detected?
	)
	if !c.ar488 {
		cmds = append(cmds,
			"savecfg 1", // Enable saving of configuration parameters in EPROM
		)
	}
	if clear {
		cmds = append(cmds, "clr")
	}
	for _, cmd := range cmds {
		if err := c.CommandController(cmd); err != nil {
			return nil, err
		}
	}

	return &c, nil
}

// WithSecondaryAddress sets a secondary address, which must be in the range of
// 96 and 126, inclusive.
func WithSecondaryAddress(addr int) ControllerOption {
	return func(c *Controller) {
		c.hasSecondaryAddr = true
		c.secondaryAddr = addr
	}
}

// WithDebug causes commands and responses to be logged.
func WithDebug() ControllerOption { return func(c *Controller) { c.debug = true } }

// WithAR488 slightly alters the init commands, for compatibility with the
// Arduino-based AR488. Specifically, we do not emit 'verbose 0', nor do
// we toggle savecfg.
func WithAR488() ControllerOption { return func(c *Controller) { c.ar488 = true } }

// Write writes the given data to the instrument at the currently assigned GPIB
// address.
func (c *Controller) Write(p []byte) (n int, err error) {
	return c.rw.Write(p)
}

// WriteBinary writes binary data to the instrument at the currently assigned
// GPIB address without adding a terminator. If the underlying writer implements
// ContextWriter, the call is delegated directly. Otherwise, the write is
// performed in a goroutine so the context cancellation is respected.
func (c *Controller) WriteBinary(ctx context.Context, p []byte) (n int, err error) {
	if cw, ok := c.rw.(ContextWriter); ok {
		return cw.WriteContext(ctx, p)
	}
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	type result struct {
		n   int
		err error
	}
	ch := make(chan result, 1)
	go func() {
		n, err := c.rw.Write(p)
		ch <- result{n, err}
	}()
	select {
	case <-ctx.Done():
		return 0, ctx.Err()
	case r := <-ch:
		return r.n, r.err
	}
}

// Read reads from the instrument at the currently assigned GPIB address into
// the given byte slice.
func (c *Controller) Read(p []byte) (n int, err error) {
	return c.rw.Read(p)
}

// ReadBinary reads binary data from the instrument at the currently assigned
// GPIB address into the given byte slice without terminator interpretation. If
// the underlying reader implements ContextReader, the call is delegated
// directly. Otherwise, the read is performed in a goroutine so the context
// cancellation is respected.
func (c *Controller) ReadBinary(ctx context.Context, p []byte) (n int, err error) {
	if cr, ok := c.rw.(ContextReader); ok {
		return cr.ReadContext(ctx, p)
	}
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	type result struct {
		n   int
		err error
	}
	ch := make(chan result, 1)
	go func() {
		n, err := c.rw.Read(p)
		ch <- result{n, err}
	}()
	select {
	case <-ctx.Done():
		return 0, ctx.Err()
	case r := <-ch:
		return r.n, r.err
	}
}

// WriteString writes a string to the instrument at the currently assigned GPIB
// address.
func (c *Controller) WriteString(s string) (n int, err error) {
	cmd := fmt.Sprintf("%s%c", strings.TrimSpace(s), c.usbTerm)
	log.Printf("prologix driver writing string: %s", cmd)
	return c.rw.Write([]byte(cmd))
}

// Command formats according to a format specifier if provided and sends a
// SCPI/ASCII command to the instrument at the currently assigned GPIB address.
// All leading and trailing whitespace is removed before appending the USB
// terminator to the command sent to the Prologix.
func (c *Controller) Command(ctx context.Context, cmd string, a ...any) error {
	if a != nil {
		cmd = fmt.Sprintf(cmd, a...)
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	cmd = fmt.Sprintf("%s%c", strings.TrimSpace(cmd), c.usbTerm)
	if c.debug {
		log.Printf("cmd %q (%x)", cmd, cmd)
	}
	_, err := fmt.Fprint(c.rw, cmd)
	return err
}

// Query queries the instrument at the currently assigned GPIB using the given
// SCPI/ASCII command. The cmd string does not need to include a new line
// character, since all leading and trailing whitespace is removed before
// appending the USB terminator to the command sent to the Prologix.  When data
// from host is received over USB, the Prologix controller removes all
// non-escaped LF, CR and ESC characters and appends the GPIB terminator, as
// specified by the `eos` command, before sending the data to instruments.  To
// change the GPIB terminator use the SetGPIBTermination method.
func (c *Controller) Query(ctx context.Context, cmd string) (string, error) {
	if err := ctx.Err(); err != nil {
		return "", err
	}
	cmd = fmt.Sprintf("%s%c", strings.TrimSpace(cmd), c.usbTerm)
	if c.debug {
		log.Printf("query: %q", cmd)
	}
	_, err := fmt.Fprint(c.rw, cmd)
	if err != nil {
		return "", fmt.Errorf("error writing command: %w", err)
	}
	// If read-after-write is disabled, need to tell the Prologix controller to
	// read.
	if !c.auto {
		readCmd := "++read eoi"
		_, err = fmt.Fprintf(c.rw, "%s%c", readCmd, c.usbTerm)
		if err != nil {
			return "", fmt.Errorf("error sending %q command: %w", readCmd, err)
		}
	}
	s, err := c.reader.ReadString(c.eotChar)
	if errors.Is(err, io.EOF) {
		log.Printf("found EOF")
		return s, nil
	}
	return s, err
}

// Close closes the underlying transport if it implements io.Closer.
func (c *Controller) Close() error {
	if closer, ok := c.rw.(io.Closer); ok {
		return closer.Close()
	}
	return nil
}

// QueryController sends the given command to the Prologix controller and
// returns its response as a string. To indicate this is a command for the
// Prologix controller, thereby not transmitting over GPIB, two plus signs `++`
// are prepended. Addtionally, a new line is appended to act as the USB
// termination character.
func (c *Controller) QueryController(cmd string) (string, error) {
	err := c.CommandController(cmd)
	if err != nil {
		return "", err
	}
	s, err := c.reader.ReadString(c.eotChar)
	if c.debug {
		log.Printf("read data: %q", s)
	}
	return s, err
}

// CommandController sends the given command to the Prologix controller. To
// indicate this is a command for the Prologix controller, thereby not
// transmitting to the instrument over GPIB, two plus signs `++` are prepended.
// Addtionally, a new line is appended to act as the USB termination character.
func (c *Controller) CommandController(cmd string) error {
	cmd = fmt.Sprintf("++%s%c", strings.ToLower(strings.TrimSpace(cmd)), c.usbTerm)
	if c.debug {
		log.Printf("cmd %q (%2x)", cmd, cmd)
	}
	_, err := c.rw.Write([]byte(cmd))
	return err
}

// GpibTerm provides the type for the available GPIB terminators.
type GpibTerm int

// Available GPIB terminators for the Prologix Controller.
const (
	AppendCRLF GpibTerm = iota
	AppendCR
	AppendLF
	AppendNothing
)

var gpibTermDesc = map[GpibTerm]string{
	AppendCRLF:    `Append CR+LF (\r\n) to instrument commands`,
	AppendCR:      `Append CR (\r) to instrument commands`,
	AppendLF:      `Append LF (\n) to instrument commands`,
	AppendNothing: `Do not append anything to instrument commands`,
}

func (term GpibTerm) String() string {
	return gpibTermDesc[term]
}

// isPrimaryAddressValid checks that the primary GPIB address is between 0 and
// 30, inclusive.
func isPrimaryAddressValid(addr int) bool {
	if addr < 0 || addr > 30 {
		return false
	}
	return true
}

// isSecondaryAddressValid checks that the secondary GPIB address is between 96
// and 126, inclusive.
func isSecondaryAddressValid(addr int) bool {
	if addr < 96 || addr > 126 {
		return false
	}
	return true
}