ShackMaster/internal/devices/ultrabeam/ultrabeam.go

package ultrabeam

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"sync"
	"time"
)

// Protocol constants
const (
	STX byte = 0xF5 // 245 decimal
	ETX byte = 0xFA // 250 decimal
	DLE byte = 0xF6 // 246 decimal
)

// Command codes
const (
	CMD_STATUS      byte = 1  // General status query
	CMD_RETRACT     byte = 2  // Retract elements
	CMD_FREQ        byte = 3  // Change frequency
	CMD_READ_BANDS  byte = 9  // Read current band adjustments
	CMD_PROGRESS    byte = 10 // Read progress bar
	CMD_MODIFY_ELEM byte = 12 // Modify element length
)

// Reply codes
const (
	UB_OK  byte = 0 // Normal execution
	UB_BAD byte = 1 // Invalid command
	UB_PAR byte = 2 // Bad parameters
	UB_ERR byte = 3 // Error executing command
)

// Direction modes
const (
	DIR_NORMAL byte = 0
	DIR_180    byte = 1
	DIR_BIDIR  byte = 2
)

type Client struct {
	host       string
	port       int
	conn       net.Conn
	connMu     sync.Mutex
	reader     *bufio.Reader
	lastStatus *Status
	statusMu   sync.RWMutex
	stopChan   chan struct{}
	running    bool
	seqNum     byte
	seqMu      sync.Mutex
}

type Status struct {
	FirmwareMinor    int   `json:"firmware_minor"`
	FirmwareMajor    int   `json:"firmware_major"`
	CurrentOperation int   `json:"current_operation"`
	Frequency        int   `json:"frequency"` // KHz
	Band             int   `json:"band"`
	Direction        int   `json:"direction"` // 0=normal, 1=180°, 2=bi-dir
	OffState         bool  `json:"off_state"`
	MotorsMoving     int   `json:"motors_moving"`    // Bitmask
	FreqMin          int   `json:"freq_min"`         // MHz
	FreqMax          int   `json:"freq_max"`         // MHz
	ElementLengths   []int `json:"element_lengths"`  // mm
	ProgressTotal    int   `json:"progress_total"`   // mm
	ProgressCurrent  int   `json:"progress_current"` // 0-60
	Connected        bool  `json:"connected"`
}

func New(host string, port int) *Client {
	return &Client{
		host:     host,
		port:     port,
		stopChan: make(chan struct{}),
		seqNum:   0,
	}
}

func (c *Client) Start() error {
	c.running = true
	go c.pollLoop()
	return nil
}

func (c *Client) Stop() {
	if !c.running {
		return
	}
	c.running = false
	close(c.stopChan)

	c.connMu.Lock()
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}
	c.connMu.Unlock()
}

func (c *Client) pollLoop() {
	ticker := time.NewTicker(2 * time.Second) // Increased from 500ms to 2s
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:

			// Try to connect if not connected
			c.connMu.Lock()
			if c.conn == nil {
				log.Printf("Ultrabeam: Not connected, attempting connection...")
				conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", c.host, c.port), 5*time.Second)
				if err != nil {
					log.Printf("Ultrabeam: Connection failed: %v", err)
					c.connMu.Unlock()

					// Mark as disconnected
					c.statusMu.Lock()
					c.lastStatus = &Status{Connected: false}
					c.statusMu.Unlock()
					continue
				}
				c.conn = conn
				c.reader = bufio.NewReader(c.conn)
				log.Printf("Ultrabeam: Connected to %s:%d", c.host, c.port)
			}
			c.connMu.Unlock()

			// Query status
			status, err := c.queryStatus()
			if err != nil {
				log.Printf("Ultrabeam: Failed to query status: %v", err)
				// Close connection and retry
				c.connMu.Lock()
				if c.conn != nil {
					c.conn.Close()
					c.conn = nil
					c.reader = nil
				}
				c.connMu.Unlock()

				// Mark as disconnected
				c.statusMu.Lock()
				c.lastStatus = &Status{Connected: false}
				c.statusMu.Unlock()
				continue
			}

			// Mark as connected
			status.Connected = true

			// Query progress if motors moving
			if status.MotorsMoving != 0 {
				progress, err := c.queryProgress()
				if err == nil {
					status.ProgressTotal = progress[0]
					status.ProgressCurrent = progress[1]
				}
			}

			c.statusMu.Lock()
			c.lastStatus = status
			c.statusMu.Unlock()

		case <-c.stopChan:
			return
		}
	}
}

func (c *Client) GetStatus() (*Status, error) {
	c.statusMu.RLock()
	defer c.statusMu.RUnlock()

	if c.lastStatus == nil {
		return &Status{Connected: false}, nil
	}

	return c.lastStatus, nil
}

// getNextSeq returns the next sequence number
func (c *Client) getNextSeq() byte {
	c.seqMu.Lock()
	defer c.seqMu.Unlock()

	seq := c.seqNum
	c.seqNum = (c.seqNum + 1) % 128
	return seq
}

// calculateChecksum calculates the checksum for a packet
func calculateChecksum(data []byte) byte {
	chk := byte(0x55)
	for _, b := range data {
		chk ^= b
		chk++
	}
	return chk
}

// quoteByte handles DLE escaping
func quoteByte(b byte) []byte {
	if b == STX || b == ETX || b == DLE {
		return []byte{DLE, b & 0x7F} // Clear MSB
	}
	return []byte{b}
}

// buildPacket creates a complete packet with checksum and escaping
func (c *Client) buildPacket(cmd byte, data []byte) []byte {
	seq := c.getNextSeq()

	// Calculate checksum on unquoted data
	payload := append([]byte{seq, cmd}, data...)
	chk := calculateChecksum(payload)

	// Build packet with quoting
	packet := []byte{STX}

	// Add quoted SEQ
	packet = append(packet, quoteByte(seq)...)

	// Add quoted CMD
	packet = append(packet, quoteByte(cmd)...)

	// Add quoted data
	for _, b := range data {
		packet = append(packet, quoteByte(b)...)
	}

	// Add quoted checksum
	packet = append(packet, quoteByte(chk)...)

	// Add ETX
	packet = append(packet, ETX)

	return packet
}

// parsePacket parses a received packet, handling DLE unescaping
func parsePacket(data []byte) (seq byte, cmd byte, payload []byte, err error) {
	if len(data) < 5 { // STX + SEQ + CMD + CHK + ETX
		return 0, 0, nil, fmt.Errorf("packet too short")
	}

	if data[0] != STX {
		return 0, 0, nil, fmt.Errorf("missing STX")
	}

	if data[len(data)-1] != ETX {
		return 0, 0, nil, fmt.Errorf("missing ETX")
	}

	// Unquote the data
	var unquoted []byte
	dle := false
	for i := 1; i < len(data)-1; i++ {
		b := data[i]
		if b == DLE {
			dle = true
			continue
		}
		if dle {
			b |= 0x80 // Set MSB
			dle = false
		}
		unquoted = append(unquoted, b)
	}

	if len(unquoted) < 3 {
		return 0, 0, nil, fmt.Errorf("unquoted packet too short")
	}

	seq = unquoted[0]
	cmd = unquoted[1]
	chk := unquoted[len(unquoted)-1]
	payload = unquoted[2 : len(unquoted)-1]

	// Verify checksum
	calcChk := calculateChecksum(unquoted[:len(unquoted)-1])
	if calcChk != chk {
		return 0, 0, nil, fmt.Errorf("checksum mismatch: got %02X, expected %02X", chk, calcChk)
	}

	return seq, cmd, payload, nil
}

// sendCommand sends a command and waits for reply
func (c *Client) sendCommand(cmd byte, data []byte) ([]byte, error) {
	c.connMu.Lock()
	defer c.connMu.Unlock()

	if c.conn == nil || c.reader == nil {
		return nil, fmt.Errorf("not connected")
	}

	// Build and send packet
	packet := c.buildPacket(cmd, data)

	_, err := c.conn.Write(packet)
	if err != nil {
		return nil, fmt.Errorf("failed to write: %w", err)
	}

	// Read reply with timeout
	c.conn.SetReadDeadline(time.Now().Add(1 * time.Second)) // Reduced from 2s to 1s

	// Read until we get a complete packet
	var buffer []byte
	for {
		b, err := c.reader.ReadByte()
		if err != nil {
			return nil, fmt.Errorf("failed to read: %w", err)
		}

		buffer = append(buffer, b)

		// Check if we have a complete packet
		if b == ETX && len(buffer) > 0 && buffer[0] == STX {
			break
		}

		// Prevent infinite loop
		if len(buffer) > 256 {
			return nil, fmt.Errorf("packet too long")
		}
	}

	// Parse reply
	_, replyCmd, payload, err := parsePacket(buffer)
	if err != nil {
		return nil, fmt.Errorf("failed to parse reply: %w", err)
	}

	// Log for debugging unknown codes
	if replyCmd != UB_OK && replyCmd != UB_BAD && replyCmd != UB_PAR && replyCmd != UB_ERR {
		log.Printf("Ultrabeam: Unknown reply code %d (0x%02X), raw packet: %v", replyCmd, replyCmd, buffer)
	}

	// Check for errors
	switch replyCmd {
	case UB_BAD:
		return nil, fmt.Errorf("invalid command")
	case UB_PAR:
		return nil, fmt.Errorf("bad parameters")
	case UB_ERR:
		return nil, fmt.Errorf("execution error")
	case UB_OK:
		return payload, nil
	default:
		// Unknown codes might indicate "busy" or "in progress"
		// Treat as non-fatal, return empty payload
		log.Printf("Ultrabeam: Unusual reply code %d, treating as busy/in-progress", replyCmd)
		return []byte{}, nil
	}
}

// queryStatus queries general status (command 1)
func (c *Client) queryStatus() (*Status, error) {
	reply, err := c.sendCommand(CMD_STATUS, nil)
	if err != nil {
		return nil, err
	}

	if len(reply) < 12 {
		return nil, fmt.Errorf("status reply too short: %d bytes", len(reply))
	}

	status := &Status{
		FirmwareMinor:    int(reply[0]),
		FirmwareMajor:    int(reply[1]),
		CurrentOperation: int(reply[2]),
		Frequency:        int(reply[3]) | (int(reply[4]) << 8),
		Band:             int(reply[5]),
		Direction:        int(reply[6] & 0x0F),
		OffState:         (reply[7] & 0x02) != 0,
		MotorsMoving:     int(reply[9]),
		FreqMin:          int(reply[10]),
		FreqMax:          int(reply[11]),
	}

	return status, nil
}

// queryElementLengths queries element lengths (command 9)
func (c *Client) queryElementLengths() ([]int, error) {
	reply, err := c.sendCommand(CMD_READ_BANDS, nil)
	if err != nil {
		return nil, err
	}

	// Debug: log raw bytes
	log.Printf("Ultrabeam element lengths raw reply (%d bytes): %v", len(reply), reply)

	// Try to extract 6 words - the protocol says 6 words (12 bytes)
	// But we're receiving 14 bytes, so there might be padding

	if len(reply) < 12 {
		return nil, fmt.Errorf("element lengths reply too short: %d bytes", len(reply))
	}

	lengths := make([]int, 6)

	// Try different interpretations
	log.Printf("=== Attempting different parsings ===")

	// Method 1: Standard little-endian from byte 0
	log.Printf("Method 1 (little-endian from 0):")
	for i := 0; i < 6 && i*2+1 < len(reply); i++ {
		lo := int(reply[i*2])
		hi := int(reply[i*2+1])
		val := lo | (hi << 8)
		log.Printf("  Element %d: bytes[%d,%d] = [%d,%d] => %d mm", i, i*2, i*2+1, lo, hi, val)
	}

	// Method 2: Big-endian from byte 0
	log.Printf("Method 2 (big-endian from 0):")
	for i := 0; i < 6 && i*2+1 < len(reply); i++ {
		hi := int(reply[i*2])
		lo := int(reply[i*2+1])
		val := lo | (hi << 8)
		log.Printf("  Element %d: bytes[%d,%d] = [%d,%d] => %d mm", i, i*2, i*2+1, hi, lo, val)
	}

	// Method 3: Skip first 2 bytes, then little-endian
	log.Printf("Method 3 (skip 2 bytes, little-endian):")
	for i := 0; i < 6 && i*2+3 < len(reply); i++ {
		lo := int(reply[i*2+2])
		hi := int(reply[i*2+3])
		val := lo | (hi << 8)
		log.Printf("  Element %d: bytes[%d,%d] = [%d,%d] => %d mm", i, i*2+2, i*2+3, lo, hi, val)
	}

	// For now, use method 1 (original)
	for i := 0; i < 6; i++ {
		if i*2+1 >= len(reply) {
			break
		}
		lo := int(reply[i*2])
		hi := int(reply[i*2+1])
		lengths[i] = lo | (hi << 8)
	}

	log.Printf("Final lengths: %v", lengths)
	return lengths, nil
}

// queryProgress queries motor progress (command 10)
func (c *Client) queryProgress() ([]int, error) {
	reply, err := c.sendCommand(CMD_PROGRESS, nil)
	if err != nil {
		return nil, err
	}

	if len(reply) < 4 {
		return nil, fmt.Errorf("progress reply too short")
	}

	total := int(reply[0]) | (int(reply[1]) << 8)
	current := int(reply[2]) | (int(reply[3]) << 8)

	return []int{total, current}, nil
}

// SetFrequency changes frequency and optional direction (command 3)
func (c *Client) SetFrequency(freqKhz int, direction int) error {
	data := []byte{
		byte(freqKhz & 0xFF),
		byte((freqKhz >> 8) & 0xFF),
		byte(direction),
	}

	_, err := c.sendCommand(CMD_FREQ, data)
	return err
}

// Retract retracts all elements (command 2)
func (c *Client) Retract() error {
	_, err := c.sendCommand(CMD_RETRACT, nil)
	return err
}

// ModifyElement modifies element length (command 12)
func (c *Client) ModifyElement(elementNum int, lengthMm int) error {
	if elementNum < 0 || elementNum > 5 {
		return fmt.Errorf("invalid element number: %d", elementNum)
	}

	data := []byte{
		byte(elementNum),
		0, // Reserved
		byte(lengthMm & 0xFF),
		byte((lengthMm >> 8) & 0xFF),
	}

	_, err := c.sendCommand(CMD_MODIFY_ELEM, data)
	return err
}