package linetest

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

// A minimal iperf3 CLIENT for one TCP upload test (client → server), speaking the iperf3 control protocol (3.x) to the
// stock `iperf3 --server` on ep0 (operator ruling 2026-10-24: the agent speaks iperf3 itself, nothing is installed on
// the host). The number that counts is the SERVER's: the bytes it received, from its results JSON — bytes the client
// handed to its socket include what still sat in the send buffer when the timer ended.
//
// The protocol, as the server drives it (one state byte at a time on the control connection):
//
//	client → 37-byte cookie
//	server → PARAM_EXCHANGE(9)      client → params JSON (4-byte big-endian length + bytes)
//	server → CREATE_STREAMS(10)     client opens N data connections, each starts with the same cookie
//	server → TEST_START(1), TEST_RUNNING(2)   client sends for the test time, then → TEST_END(4)
//	server → EXCHANGE_RESULTS(13)   client → its results JSON; server → its results JSON
//	server → DISPLAY_RESULTS(14)    client → IPERF_DONE(16)
//
// A server already running a test answers ACCESS_DENIED(-1) and closes: that is ErrBusy (another box is testing; try
// again later the same day). SERVER_ERROR(-3) / SERVER_TERMINATE(-2) end the test with an error.

const (
	stParamExchange   = 9
	stCreateStreams   = 10
	stTestStart       = 1
	stTestRunning     = 2
	stTestEnd         = 4
	stExchangeResults = 13
	stDisplayResults  = 14
	stIperfDone       = 16
	stAccessDenied    = -1
	stServerTerminate = -2
	stServerError     = -3

	cookieSize   = 37 // 36 characters + NUL
	blockSize    = 128 * 1024
	maxJSONBytes = 1 << 20
)

// ErrBusy: the server is running another client's test.
var ErrBusy = errors.New("iperf3 server busy")

// Result is one upload test.
type Result struct {
	Mbps          float64 // the server's received bits per second / 1e6
	ServerBytes   int64   // bytes the server received, all streams
	ClientBytes   int64   // bytes the client wrote (diagnostics only)
	ServerSeconds float64 // the server's measured duration
}

// Dialer opens a TCP connection (net.Dialer in production; a seam for tests).
type Dialer interface {
	DialContext(ctx context.Context, network, address string) (net.Conn, error)
}

// Client runs one test. The whole run is bounded by ctx; Run sets every connection's deadline from it.
type Client struct {
	Dialer  Dialer
	Streams int           // parallel streams (-P)
	Time    time.Duration // test time (-t)
}

func (c Client) dialer() Dialer {
	if c.Dialer != nil {
		return c.Dialer
	}
	return &net.Dialer{Timeout: 5 * time.Second}
}

func newCookie() ([]byte, error) {
	const alphabet = "abcdefghijklmnopqrstuvwxyz234567"
	b := make([]byte, cookieSize)
	if _, err := rand.Read(b[:cookieSize-1]); err != nil {
		return nil, err
	}
	for i := 0; i < cookieSize-1; i++ {
		b[i] = alphabet[int(b[i])%len(alphabet)]
	}
	b[cookieSize-1] = 0
	return b, nil
}

func writeJSON(w io.Writer, v any) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	var hdr [4]byte
	binary.BigEndian.PutUint32(hdr[:], uint32(len(data)))
	if _, err := w.Write(append(hdr[:], data...)); err != nil {
		return err
	}
	return nil
}

func readJSON(r io.Reader, v any) error {
	var hdr [4]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return err
	}
	n := binary.BigEndian.Uint32(hdr[:])
	if n == 0 || n > maxJSONBytes {
		return fmt.Errorf("iperf3: results JSON of %d bytes refused", n)
	}
	buf := make([]byte, n)
	if _, err := io.ReadFull(r, buf); err != nil {
		return err
	}
	return json.Unmarshal(buf, v)
}

type streamResult struct {
	ID             int     `json:"id"`
	Bytes          int64   `json:"bytes"`
	Retransmits    int64   `json:"retransmits"`
	Jitter         float64 `json:"jitter"`
	Errors         int64   `json:"errors"`
	OmittedErrors  int64   `json:"omitted_errors"`
	Packets        int64   `json:"packets"`
	OmittedPackets int64   `json:"omitted_packets"`
	StartTime      float64 `json:"start_time"`
	EndTime        float64 `json:"end_time"`
}

type results struct {
	CPUTotal             float64        `json:"cpu_util_total"`
	CPUUser              float64        `json:"cpu_util_user"`
	CPUSystem            float64        `json:"cpu_util_system"`
	SenderHasRetransmits int            `json:"sender_has_retransmits"`
	CongestionUsed       string         `json:"congestion_used,omitempty"`
	Streams              []streamResult `json:"streams"`
}

// streamIDs are the ids iperf3 gives its streams: the first is 1, every later one is (streams so far) + 2 — so 1, 3,
// 4, 5 … (iperf_add_stream). The server matches our results by these ids.
func streamIDs(n int) []int {
	ids := make([]int, n)
	for i := range ids {
		if i == 0 {
			ids[i] = 1
		} else {
			ids[i] = i + 2
		}
	}
	return ids
}

// Run performs one upload test to addr ("10.77.0.1:5201").
func (c Client) Run(ctx context.Context, addr string) (Result, error) {
	var res Result
	streams, dur := c.Streams, c.Time
	if streams < 1 {
		streams = 1
	}
	if dur <= 0 {
		dur = 10 * time.Second
	}
	deadline, ok := ctx.Deadline()
	if !ok {
		return res, errors.New("iperf3: Run needs a context with a deadline")
	}
	cookie, err := newCookie()
	if err != nil {
		return res, err
	}
	ctrl, err := c.dialer().DialContext(ctx, "tcp", addr)
	if err != nil {
		return res, fmt.Errorf("iperf3: connect: %w", err)
	}
	defer ctrl.Close()
	ctrl.SetDeadline(deadline)
	if _, err := ctrl.Write(cookie); err != nil {
		return res, fmt.Errorf("iperf3: cookie: %w", err)
	}

	var data []net.Conn
	defer func() {
		for _, d := range data {
			d.Close()
		}
	}()
	var sent atomic.Int64
	var server results
	var wg sync.WaitGroup
	for {
		var st [1]byte
		if _, err := io.ReadFull(ctrl, st[:]); err != nil {
			return res, fmt.Errorf("iperf3: control read: %w", err)
		}
		switch int8(st[0]) {
		case stParamExchange:
			params := map[string]any{"tcp": true, "omit": 0, "time": int(dur / time.Second), "num": 0, "blockcount": 0,
				"parallel": streams, "len": blockSize, "pacing_timer": 1000, "client_version": "3.18"}
			if err := writeJSON(ctrl, params); err != nil {
				return res, fmt.Errorf("iperf3: params: %w", err)
			}
		case stCreateStreams:
			for i := 0; i < streams; i++ {
				d, err := c.dialer().DialContext(ctx, "tcp", addr)
				if err != nil {
					return res, fmt.Errorf("iperf3: data stream %d: %w", i+1, err)
				}
				d.SetDeadline(deadline)
				data = append(data, d)
				if _, err := d.Write(cookie); err != nil {
					return res, fmt.Errorf("iperf3: data stream %d cookie: %w", i+1, err)
				}
			}
		case stTestStart:
		case stTestRunning:
			stopAt := time.Now().Add(dur)
			if stopAt.After(deadline) {
				return res, errors.New("iperf3: no time left for the test inside the deadline")
			}
			buf := make([]byte, blockSize)
			for _, d := range data {
				d.SetWriteDeadline(stopAt)
				wg.Add(1)
				go func(d net.Conn) {
					defer wg.Done()
					for {
						n, err := d.Write(buf)
						sent.Add(int64(n))
						if err != nil {
							return
						}
					}
				}(d)
			}
			wg.Wait()
			if _, err := ctrl.Write([]byte{stTestEnd}); err != nil {
				return res, fmt.Errorf("iperf3: test end: %w", err)
			}
		case stExchangeResults:
			ids := streamIDs(len(data))
			mine := results{SenderHasRetransmits: -1, CongestionUsed: "cubic"}
			per := sent.Load() / int64(max(len(data), 1))
			for _, id := range ids {
				mine.Streams = append(mine.Streams, streamResult{ID: id, Bytes: per, Retransmits: -1, EndTime: dur.Seconds()})
			}
			if err := writeJSON(ctrl, mine); err != nil {
				return res, fmt.Errorf("iperf3: send results: %w", err)
			}
			if err := readJSON(ctrl, &server); err != nil {
				return res, fmt.Errorf("iperf3: read server results: %w", err)
			}
		case stDisplayResults:
			ctrl.Write([]byte{stIperfDone})
			return summarize(server, sent.Load(), dur)
		case stAccessDenied:
			return res, ErrBusy
		case stServerTerminate, stServerError:
			var codes [8]byte
			io.ReadFull(ctrl, codes[:])
			return res, fmt.Errorf("iperf3: server error (state %d, iperf errno %d, errno %d)", int8(st[0]),
				int32(binary.BigEndian.Uint32(codes[:4])), int32(binary.BigEndian.Uint32(codes[4:])))
		default:
			return res, fmt.Errorf("iperf3: unexpected state %s", strconv.Itoa(int(int8(st[0]))))
		}
	}
}

// summarize turns the server's results into a Result. No stream bytes is an error, never a zero line.
func summarize(server results, clientBytes int64, dur time.Duration) (Result, error) {
	res := Result{ClientBytes: clientBytes}
	for _, s := range server.Streams {
		res.ServerBytes += s.Bytes
		if s.EndTime > res.ServerSeconds {
			res.ServerSeconds = s.EndTime
		}
	}
	if res.ServerBytes <= 0 {
		return res, errors.New("iperf3: the server reported no received bytes")
	}
	secs := res.ServerSeconds
	if secs <= 0 {
		secs = dur.Seconds()
	}
	res.Mbps = float64(res.ServerBytes) * 8 / secs / 1e6
	return res, nil
}