package modbus

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

// Modbus function codes.
const (
	fcReadCoils         = 0x01
	fcReadDiscrete      = 0x02
	fcReadHolding       = 0x03
	fcReadInput         = 0x04
	fcWriteSingleCoil   = 0x05
	fcWriteSingleReg    = 0x06
	fcWriteMultipleRegs = 0x10
)

// client is a minimal Modbus TCP master for a single device address. Requests
// are serialised by mu (Modbus TCP is request/response and the connection is
// shared by all of a device's polled registers). The connection is dialled
// lazily and dropped on any I/O error so the next request reconnects.
type client struct {
	addr    string
	timeout time.Duration

	mu   sync.Mutex
	conn net.Conn
	txID uint16
}

func newClient(addr string, timeout time.Duration) *client {
	if timeout <= 0 {
		timeout = 3 * time.Second
	}
	return &client{addr: addr, timeout: timeout}
}

func (c *client) close() {
	c.mu.Lock()
	c.closeLocked()
	c.mu.Unlock()
}

func (c *client) closeLocked() {
	if c.conn != nil {
		_ = c.conn.Close()
		c.conn = nil
	}
}

// request sends a PDU to unitID and returns the response PDU (function code +
// data). It dials on demand and tears the connection down on error.
func (c *client) request(unitID byte, pdu []byte) ([]byte, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn == nil {
		conn, err := net.DialTimeout("tcp", c.addr, c.timeout)
		if err != nil {
			return nil, fmt.Errorf("modbus: dial %s: %w", c.addr, err)
		}
		c.conn = conn
	}

	resp, err := c.transact(unitID, pdu)
	if err != nil {
		c.closeLocked()
		return nil, err
	}
	return resp, nil
}

// transact performs one MBAP-framed exchange. Caller holds mu.
func (c *client) transact(unitID byte, pdu []byte) ([]byte, error) {
	c.txID++
	tx := c.txID

	frame := make([]byte, 7+len(pdu))
	binary.BigEndian.PutUint16(frame[0:], tx)                 // transaction id
	binary.BigEndian.PutUint16(frame[2:], 0)                  // protocol id (0 = Modbus)
	binary.BigEndian.PutUint16(frame[4:], uint16(1+len(pdu))) // length: unit id + PDU
	frame[6] = unitID
	copy(frame[7:], pdu)

	_ = c.conn.SetDeadline(time.Now().Add(c.timeout))
	if _, err := c.conn.Write(frame); err != nil {
		return nil, fmt.Errorf("modbus: write: %w", err)
	}

	head := make([]byte, 7)
	if _, err := io.ReadFull(c.conn, head); err != nil {
		return nil, fmt.Errorf("modbus: read header: %w", err)
	}
	if binary.BigEndian.Uint16(head[0:]) != tx {
		return nil, fmt.Errorf("modbus: transaction id mismatch")
	}
	length := binary.BigEndian.Uint16(head[4:])
	if length < 2 { // unit id + at least a function code
		return nil, fmt.Errorf("modbus: short frame length %d", length)
	}
	body := make([]byte, length-1) // header already consumed the unit id
	if _, err := io.ReadFull(c.conn, body); err != nil {
		return nil, fmt.Errorf("modbus: read body: %w", err)
	}

	fc := body[0]
	if fc&0x80 != 0 { // exception response
		var ex byte
		if len(body) >= 2 {
			ex = body[1]
		}
		return nil, fmt.Errorf("modbus: exception 0x%02x (%s)", ex, exceptionText(ex))
	}
	return body, nil
}

// readRegisters reads `quantity` 16-bit registers via fc (holding or input).
func (c *client) readRegisters(unitID, fc byte, addr, quantity uint16) ([]uint16, error) {
	pdu := []byte{fc, byte(addr >> 8), byte(addr), byte(quantity >> 8), byte(quantity)}
	resp, err := c.request(unitID, pdu)
	if err != nil {
		return nil, err
	}
	if len(resp) < 2 {
		return nil, fmt.Errorf("modbus: short register response")
	}
	byteCount := int(resp[1])
	if byteCount != int(quantity)*2 || len(resp) < 2+byteCount {
		return nil, fmt.Errorf("modbus: register byte count %d (want %d)", byteCount, quantity*2)
	}
	regs := make([]uint16, quantity)
	for i := range regs {
		regs[i] = binary.BigEndian.Uint16(resp[2+i*2:])
	}
	return regs, nil
}

// readBits reads `quantity` bits via fc (coils or discrete inputs).
func (c *client) readBits(unitID, fc byte, addr, quantity uint16) ([]bool, error) {
	pdu := []byte{fc, byte(addr >> 8), byte(addr), byte(quantity >> 8), byte(quantity)}
	resp, err := c.request(unitID, pdu)
	if err != nil {
		return nil, err
	}
	if len(resp) < 2 {
		return nil, fmt.Errorf("modbus: short bit response")
	}
	byteCount := int(resp[1])
	if len(resp) < 2+byteCount {
		return nil, fmt.Errorf("modbus: bit byte count %d exceeds frame", byteCount)
	}
	bits := make([]bool, quantity)
	for i := range bits {
		idx := 2 + i/8
		if idx >= len(resp) {
			break
		}
		bits[i] = resp[idx]&(1<<(uint(i)%8)) != 0
	}
	return bits, nil
}

func (c *client) writeSingleRegister(unitID byte, addr, value uint16) error {
	pdu := []byte{fcWriteSingleReg, byte(addr >> 8), byte(addr), byte(value >> 8), byte(value)}
	_, err := c.request(unitID, pdu)
	return err
}

func (c *client) writeMultipleRegisters(unitID byte, addr uint16, values []uint16) error {
	pdu := make([]byte, 6+len(values)*2)
	pdu[0] = fcWriteMultipleRegs
	binary.BigEndian.PutUint16(pdu[1:], addr)
	binary.BigEndian.PutUint16(pdu[3:], uint16(len(values)))
	pdu[5] = byte(len(values) * 2)
	for i, v := range values {
		binary.BigEndian.PutUint16(pdu[6+i*2:], v)
	}
	_, err := c.request(unitID, pdu)
	return err
}

func (c *client) writeSingleCoil(unitID byte, addr uint16, on bool) error {
	var v uint16
	if on {
		v = 0xFF00
	}
	pdu := []byte{fcWriteSingleCoil, byte(addr >> 8), byte(addr), byte(v >> 8), byte(v)}
	_, err := c.request(unitID, pdu)
	return err
}

func exceptionText(code byte) string {
	switch code {
	case 0x01:
		return "illegal function"
	case 0x02:
		return "illegal data address"
	case 0x03:
		return "illegal data value"
	case 0x04:
		return "server device failure"
	case 0x05:
		return "acknowledge"
	case 0x06:
		return "server device busy"
	case 0x0B:
		return "gateway target failed to respond"
	default:
		return "unknown"
	}
}