package modbus

import (
	"context"
	"encoding/binary"
	"errors"
	"io"
	"net"
	"testing"
)

// serve runs a one-connection fake Modbus server. handle receives the request
// PDU and returns the reply PDU; returning nil closes the connection instead,
// which is how a mute or dying peer is simulated. The request frames the server
// saw are sent back on the returned channel once the connection ends.
func serve(t *testing.T, handle func(pdu []byte) []byte) (addr string, seen <-chan [][]byte) {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	t.Cleanup(func() { _ = ln.Close() })

	frames := make(chan [][]byte, 1)
	go func() {
		var got [][]byte
		defer func() { frames <- got }()
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer conn.Close()
		for {
			var header [mbapLen]byte
			if _, err := io.ReadFull(conn, header[:]); err != nil {
				return
			}
			length := int(binary.BigEndian.Uint16(header[4:]))
			body := make([]byte, length-1)
			if _, err := io.ReadFull(conn, body); err != nil {
				return
			}
			got = append(got, append(append([]byte{}, header[:]...), body...))

			reply := handle(body)
			if reply == nil {
				return
			}
			out := make([]byte, mbapLen+len(reply))
			copy(out[0:], header[0:2]) // echo the transaction id
			binary.BigEndian.PutUint16(out[2:], protocolID)
			binary.BigEndian.PutUint16(out[4:], uint16(len(reply)+1))
			out[6] = header[6]
			copy(out[mbapLen:], reply)
			if _, err := conn.Write(out); err != nil {
				return
			}
		}
	}()
	return ln.Addr().String(), frames
}

func dial(t *testing.T, addr string) *Client {
	t.Helper()
	c, err := Connect(context.Background(), Options{Address: addr})
	if err != nil {
		t.Fatalf("connect: %v", err)
	}
	t.Cleanup(func() { _ = c.Close() })
	return c
}

// readReply builds a well-formed FC03 reply carrying the given register values.
func readReply(values ...uint16) []byte { return replyFor(fcReadHold, values...) }

// inputReply is the same for FC04, which frames its reply identically.
func inputReply(values ...uint16) []byte { return replyFor(fcReadInput, values...) }

func replyFor(fc byte, values ...uint16) []byte {
	out := []byte{fc, byte(len(values) * 2)}
	for _, v := range values {
		out = binary.BigEndian.AppendUint16(out, v)
	}
	return out
}

func TestReadHoldingDecodesRegisters(t *testing.T) {
	addr, seen := serve(t, func(pdu []byte) []byte {
		return readReply(0x0102, 0x0304, 0xFFFF)
	})
	c := dial(t, addr)

	got, err := c.ReadHolding(context.Background(), 20053, 3)
	if err != nil {
		t.Fatalf("ReadHolding: %v", err)
	}
	want := []uint16{0x0102, 0x0304, 0xFFFF}
	if len(got) != len(want) {
		t.Fatalf("got %d registers, want %d", len(got), len(want))
	}
	for i := range want {
		if got[i] != want[i] {
			t.Errorf("register %d = 0x%04x, want 0x%04x", i, got[i], want[i])
		}
	}

	_ = c.Close()
	frames := <-seen
	if len(frames) != 1 {
		t.Fatalf("server saw %d frames, want 1", len(frames))
	}
	f := frames[0]
	if pid := binary.BigEndian.Uint16(f[2:]); pid != protocolID {
		t.Errorf("protocol id = %d, want 0", pid)
	}
	// Length counts the unit id plus the five-byte PDU.
	if l := binary.BigEndian.Uint16(f[4:]); l != 6 {
		t.Errorf("length field = %d, want 6", l)
	}
	if f[6] != defaultUnitID {
		t.Errorf("unit id = %d, want %d", f[6], defaultUnitID)
	}
	if f[mbapLen] != fcReadHold {
		t.Errorf("function code = 0x%02x, want 0x%02x", f[mbapLen], fcReadHold)
	}
	if a := binary.BigEndian.Uint16(f[mbapLen+1:]); a != 20053 {
		t.Errorf("address = %d, want 20053", a)
	}
	if n := binary.BigEndian.Uint16(f[mbapLen+3:]); n != 3 {
		t.Errorf("quantity = %d, want 3", n)
	}
}

func TestReadInputAsksTheInputTable(t *testing.T) {
	addr, seen := serve(t, func(pdu []byte) []byte {
		// A charger that keeps its measurements in input registers answers FC04
		// and refuses FC03, so replying to the wrong code would hide a mix-up.
		if pdu[0] != fcReadInput {
			return []byte{pdu[0] | excMask, 0x02}
		}
		return inputReply(0x0102, 0x0304)
	})
	c := dial(t, addr)

	got, err := c.ReadInput(context.Background(), 20041, 2)
	if err != nil {
		t.Fatalf("ReadInput: %v", err)
	}
	if len(got) != 2 || got[0] != 0x0102 || got[1] != 0x0304 {
		t.Errorf("registers = %v, want [258 772]", got)
	}

	_ = c.Close()
	frames := <-seen
	if len(frames) != 1 {
		t.Fatalf("server saw %d frames, want 1", len(frames))
	}
	if fc := frames[0][mbapLen]; fc != fcReadInput {
		t.Errorf("function code = 0x%02x, want 0x%02x", fc, fcReadInput)
	}
	if a := binary.BigEndian.Uint16(frames[0][mbapLen+1:]); a != 20041 {
		t.Errorf("address = %d, want 20041", a)
	}
}

func TestReadInputRejectsHoldingReply(t *testing.T) {
	addr, _ := serve(t, func(pdu []byte) []byte { return readReply(0x0102) })
	c := dial(t, addr)

	if _, err := c.ReadInput(context.Background(), 20041, 1); err == nil {
		t.Fatal("expected an error when the reply echoes the other function code")
	}
}

func TestReadHoldingRejectsBadCounts(t *testing.T) {
	addr, _ := serve(t, func(pdu []byte) []byte { return readReply(0) })
	c := dial(t, addr)

	for _, count := range []uint16{0, maxReadRegs + 1} {
		if _, err := c.ReadHolding(context.Background(), 20000, count); err == nil {
			t.Errorf("count %d was accepted, want a refusal before the wire", count)
		}
	}
}

func TestReadHoldingRejectsShortPayload(t *testing.T) {
	// The server claims three registers but sends two.
	addr, _ := serve(t, func(pdu []byte) []byte {
		return []byte{fcReadHold, 6, 0x00, 0x01, 0x00, 0x02}
	})
	c := dial(t, addr)

	if _, err := c.ReadHolding(context.Background(), 20000, 3); err == nil {
		t.Fatal("a truncated payload was accepted")
	}
}

func TestWriteSingleEchoesAddressAndValue(t *testing.T) {
	addr, seen := serve(t, func(pdu []byte) []byte {
		return append([]byte{fcWriteReg}, pdu[1:5]...) // echo address + value
	})
	c := dial(t, addr)

	if err := c.WriteSingle(context.Background(), 21000, 1); err != nil {
		t.Fatalf("WriteSingle: %v", err)
	}

	_ = c.Close()
	frames := <-seen
	if len(frames) != 1 {
		t.Fatalf("server saw %d frames, want 1", len(frames))
	}
	f := frames[0]
	if f[mbapLen] != fcWriteReg {
		t.Errorf("function code = 0x%02x, want 0x%02x", f[mbapLen], fcWriteReg)
	}
	if a := binary.BigEndian.Uint16(f[mbapLen+1:]); a != 21000 {
		t.Errorf("address = %d, want 21000", a)
	}
	if v := binary.BigEndian.Uint16(f[mbapLen+3:]); v != 1 {
		t.Errorf("value = %d, want 1", v)
	}
}

// A charger that clamps a written value reports the clamp. That must surface as
// an error rather than be mistaken for the value having been applied.
func TestWriteSingleRejectsClampedValue(t *testing.T) {
	addr, _ := serve(t, func(pdu []byte) []byte {
		reply := append([]byte{fcWriteReg}, pdu[1:5]...)
		binary.BigEndian.PutUint16(reply[3:], 160) // clamped to 16.0 A
		return reply
	})
	c := dial(t, addr)

	err := c.WriteSingle(context.Background(), 21001, 320)
	if err == nil {
		t.Fatal("a clamped write was reported as successful")
	}
}

func TestWriteSingleRejectsWrongAddressEcho(t *testing.T) {
	addr, _ := serve(t, func(pdu []byte) []byte {
		reply := append([]byte{fcWriteReg}, pdu[1:5]...)
		binary.BigEndian.PutUint16(reply[1:], 29999)
		return reply
	})
	c := dial(t, addr)

	if err := c.WriteSingle(context.Background(), 21000, 1); err == nil {
		t.Fatal("a write echoing the wrong register was accepted")
	}
}

func TestExceptionReplyBecomesError(t *testing.T) {
	addr, _ := serve(t, func(pdu []byte) []byte {
		return []byte{fcReadHold | excMask, 0x02} // illegal data address
	})
	c := dial(t, addr)

	_, err := c.ReadHolding(context.Background(), 29999, 1)
	var mbErr *Error
	if !errors.As(err, &mbErr) {
		t.Fatalf("error = %v, want a *modbus.Error", err)
	}
	if mbErr.Function != fcReadHold {
		t.Errorf("Function = 0x%02x, want 0x%02x", mbErr.Function, fcReadHold)
	}
	if mbErr.Code != 0x02 {
		t.Errorf("Code = 0x%02x, want 0x02", mbErr.Code)
	}
}

// A reply carrying someone else's transaction id means the stream is out of
// step, so it must not be handed back as this request's answer.
func TestMismatchedTransactionIDIsRejected(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer ln.Close()
	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer conn.Close()
		var header [mbapLen]byte
		if _, err := io.ReadFull(conn, header[:]); err != nil {
			return
		}
		body := make([]byte, int(binary.BigEndian.Uint16(header[4:]))-1)
		if _, err := io.ReadFull(conn, body); err != nil {
			return
		}
		reply := readReply(1)
		out := make([]byte, mbapLen+len(reply))
		binary.BigEndian.PutUint16(out[0:], binary.BigEndian.Uint16(header[0:])+7) // wrong
		binary.BigEndian.PutUint16(out[2:], protocolID)
		binary.BigEndian.PutUint16(out[4:], uint16(len(reply)+1))
		out[6] = header[6]
		copy(out[mbapLen:], reply)
		_, _ = conn.Write(out)
	}()

	c := dial(t, ln.Addr().String())
	if _, err := c.ReadHolding(context.Background(), 20000, 1); err == nil {
		t.Fatal("a reply for another transaction was accepted")
	}
}

func TestClosedClientRefusesRequests(t *testing.T) {
	addr, _ := serve(t, func(pdu []byte) []byte { return readReply(0) })
	c := dial(t, addr)
	if err := c.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
	if _, err := c.ReadHolding(context.Background(), 20000, 1); !errors.Is(err, ErrClosed) {
		t.Fatalf("error = %v, want ErrClosed", err)
	}
	// Close is deferred by callers and may also run on an error path.
	if err := c.Close(); err != nil {
		t.Fatalf("second Close: %v", err)
	}
}

func TestConnectRequiresAddress(t *testing.T) {
	if _, err := Connect(context.Background(), Options{}); err == nil {
		t.Fatal("an empty address was accepted")
	}
}