package m31

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"time"
)

const (
	connectTimeout = 2 * time.Second
	writeTimeout   = time.Second
	readTimeout    = 2 * time.Second
)

type tcpClient struct {
	conn        net.Conn
	unit        byte
	transaction uint16
}

// NewTCPClient supplies Modbus TCP functions 01, 02, 03, 05, 15, and 16 without
// opening a connection. Connect has a two-second timeout. Writes have a
// one-second timeout, and reads have a two-second timeout. Earlier caller
// deadlines apply to each operation. The caller owns Close.
func NewTCPClient() Client {
	return &tcpClient{}
}

func (c *tcpClient) Connect(ctx context.Context, address string, unit byte) error {
	if c.conn != nil {
		return fmt.Errorf("modbus connection is already open")
	}

	if unit == 0 || unit > 247 {
		return fmt.Errorf("modbus unit must be from one through 247")
	}

	ctx, cancel := context.WithTimeout(ctx, connectTimeout)
	defer cancel()

	conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", address)
	if err != nil {
		return fmt.Errorf("connect Modbus transport: %w", err)
	}

	c.conn, c.unit = conn, unit

	return nil
}

func (c *tcpClient) Close() error {
	if c.conn == nil {
		return nil
	}

	err := c.conn.Close()
	c.conn = nil
	if err != nil {
		return fmt.Errorf("close Modbus transport: %w", err)
	}

	return nil
}

func (c *tcpClient) ReadInputs(ctx context.Context, count uint16) ([]bool, error) {
	return c.readBits(ctx, 2, count)
}

func (c *tcpClient) ReadCoils(ctx context.Context, count uint16) ([]bool, error) {
	return c.readBits(ctx, 1, count)
}

func (c *tcpClient) readBits(ctx context.Context, function byte, count uint16) ([]bool, error) {
	if count == 0 || count > 2000 {
		return nil, fmt.Errorf("bit read count must be from one through 2000")
	}

	data, err := c.exchange(ctx, function, words(0, count))
	if err != nil {
		return nil, fmt.Errorf("read Modbus bits: %w", err)
	}

	size := (int(count) + 7) / 8
	if len(data) != size+1 || int(data[0]) != size {
		return nil, fmt.Errorf("modbus bit response has an invalid length")
	}

	values := make([]bool, count)

	for index := range values {
		values[index] = data[1+index/8]&(1<<uint(index%8)) != 0
	}

	return values, nil
}

func (c *tcpClient) ReadRegisters(ctx context.Context, address uint16, count uint16) ([]uint16, error) {
	if count == 0 || count > 125 || uint32(address)+uint32(count) > 65536 {
		return nil, fmt.Errorf("invalid register read range")
	}

	data, err := c.exchange(ctx, 3, words(address, count))
	if err != nil {
		return nil, fmt.Errorf("read Modbus registers: %w", err)
	}

	if len(data) != int(count)*2+1 || int(data[0]) != int(count)*2 {
		return nil, fmt.Errorf("modbus register response has an invalid length")
	}

	values := make([]uint16, count)

	for index := range values {
		values[index] = binary.BigEndian.Uint16(data[1+2*index:])
	}

	return values, nil
}

func (c *tcpClient) WriteCoil(ctx context.Context, address uint16, value bool) error {
	var encoded uint16

	if value {
		encoded = 0xff00
	}

	return c.write(ctx, 5, words(address, encoded), words(address, encoded))
}

func (c *tcpClient) WriteCoils(ctx context.Context, address uint16, values []bool) error {
	if len(values) == 0 || len(values) > 1968 || uint32(address)+uint32(len(values)) > 65536 {
		return fmt.Errorf("invalid coil write range")
	}

	data := append(words(address, uint16(len(values))), byte((len(values)+7)/8))
	data = append(data, make([]byte, (len(values)+7)/8)...)

	for index, value := range values {
		if value {
			data[5+index/8] |= 1 << uint(index%8)
		}
	}

	return c.write(ctx, 15, data, words(address, uint16(len(values))))
}

func (c *tcpClient) WriteRegister(ctx context.Context, address uint16, value uint16) error {
	data := append(words(address, 1), 2)
	data = append(data, words(value)...)

	return c.write(ctx, 16, data, words(address, 1))
}

func (c *tcpClient) write(ctx context.Context, function byte, request []byte, expected []byte) error {
	data, err := c.exchange(ctx, function, request)
	if err != nil {
		return fmt.Errorf("write Modbus value: %w", err)
	}

	if !bytes.Equal(data, expected) {
		return fmt.Errorf("modbus write acknowledgement differs from request")
	}

	return nil
}

func (c *tcpClient) exchange(ctx context.Context, function byte, data []byte) (response []byte, err error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	if c.conn == nil {
		return nil, fmt.Errorf("modbus transport is closed")
	}

	c.transaction++
	request := append(words(c.transaction, 0, uint16(len(data)+2)), c.unit, function)
	request = append(request, data...)
	err = c.conn.SetWriteDeadline(operationDeadline(ctx, writeTimeout))
	if err != nil {
		return nil, fmt.Errorf("set Modbus write deadline: %w", err)
	}

	err = c.conn.SetReadDeadline(operationDeadline(ctx, readTimeout))
	if err != nil {
		return nil, fmt.Errorf("set Modbus read deadline: %w", err)
	}

	interrupted := make(chan struct{})

	var interruptErr error

	stop := context.AfterFunc(ctx, func() {
		interruptErr = c.conn.SetDeadline(time.Now())

		close(interrupted)
	})
	defer func() {
		if !stop() {
			<-interrupted

			err = errors.Join(err, interruptErr, ctx.Err())
		}
	}()

	count, err := c.conn.Write(request)
	if err != nil {
		return nil, fmt.Errorf("write Modbus request: %w", err)
	}

	if count != len(request) {
		return nil, fmt.Errorf("write Modbus request: %w", io.ErrShortWrite)
	}

	var header [7]byte

	_, err = io.ReadFull(c.conn, header[:])
	if err != nil {
		return nil, fmt.Errorf("read Modbus header: %w", err)
	}

	length := binary.BigEndian.Uint16(header[4:6])
	if binary.BigEndian.Uint16(header[:2]) != c.transaction || binary.BigEndian.Uint16(header[2:4]) != 0 || header[6] != c.unit || length < 2 || length > 254 {
		return nil, fmt.Errorf("invalid Modbus response header")
	}

	body := make([]byte, length-1)
	_, err = io.ReadFull(c.conn, body)
	if err != nil {
		return nil, fmt.Errorf("read Modbus body: %w", err)
	}

	if body[0] == function|0x80 && len(body) == 2 {
		return nil, fmt.Errorf("modbus exception %d", body[1])
	}

	if body[0] != function {
		return nil, fmt.Errorf("invalid Modbus response function")
	}

	return body[1:], nil
}

func words(values ...uint16) []byte {
	data := make([]byte, 2*len(values))

	for index, value := range values {
		binary.BigEndian.PutUint16(data[2*index:], value)
	}

	return data
}

func operationDeadline(ctx context.Context, timeout time.Duration) time.Time {
	deadline := time.Now().Add(timeout)
	caller, exists := ctx.Deadline()
	if exists && caller.Before(deadline) {
		return caller
	}

	return deadline
}