package network

// The speed test. It measures against Cloudflare's public speed endpoints, the same ones
// speed.cloudflare.com uses: a plain GET of N bytes down, a plain POST of N bytes up, and a
// zero-byte GET for the round trip. No account, no API key, nothing to install, and an edge
// close to every one of these machines.
//
// The method, since a number without one means little — it is stated in the UI too:
//
//   - Latency is the round trip of a zero-byte response on a connection that is already open,
//     less the time the far end says it spent on the request (its Server-Timing header). That
//     subtraction is the difference between measuring the line and measuring Cloudflare's
//     thinking time: a zero-byte response is almost entirely the latter — tens of
//     milliseconds of worker, against a few of actual wire. The figure reported is the best of
//     several probes; jitter is the average change from one to the next.
//   - Download and upload each run several streams at once for a fixed few seconds. One
//     stream measures one TCP connection's congestion window as much as the line, which is
//     why every speed test ever made uses more than one.
//   - The first second of each phase is thrown away. That is the ramp: counting it reports a
//     number nobody's connection ever actually did.
//   - Each phase is capped in bytes as well as seconds, so a very fast line cannot run away
//     with the meter.
//
// The one thing to know about the far end: it rate-limits by volume, and a test moves a few
// hundred megabytes. Several in a row will start earning 429s, which is why there is a
// cooldown between tests and why a phase that gets refused says so instead of reporting the
// decayed figure it would otherwise compute.
//
// HTTP/2 is deliberately off: over one HTTP/2 connection the streams would share a single
// congestion window, which measures the window, not the line.

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// The endpoints, and the shape of a phase. Variables rather than constants only so the tests
// can point them at a local server and shrink the windows; nothing else rebinds them.
// Source is somewhere to pull bytes from. Cloudflare's endpoint is the first choice because
// it is anycast — near wherever Ham Lab is running — and because the same host answers for
// upload and for the trace. The fallback exists because that endpoint rate-limits by volume
// and will refuse a connection that has been tested hard: a plain 100 MB file on a public
// mirror has no such opinion. It is in Seattle, so on a machine far from there it will read
// low, which is exactly why every result names the source it used.
type Source struct {
	Name string
	// URL takes the number of bytes wanted, where the source can be asked; a fixed file
	// ignores it and is fetched again each time round.
	URL func(bytes int) string
}

var downSources = []Source{
	{"cloudflare", func(b int) string { return fmt.Sprintf(cfDownURL, b) }},
	{"linode-seattle", func(int) string { return "https://speedtest.seattle.linode.com/100MB-seattle.bin" }},
}

var (
	cfDownURL = "https://speed.cloudflare.com/__down?bytes=%d"
	upURL     = "https://speed.cloudflare.com/__up"
	// traceURL is Cloudflare's standard trace: a few lines of key=value naming the address it
	// sees, the edge that answered and the country. A few hundred bytes.
	traceURL = "https://speed.cloudflare.com/cdn-cgi/trace"

	warmUp  = time.Second
	downFor = 4 * time.Second
	upFor   = 4 * time.Second
	// measureBytes ends a phase early once this much has moved inside the window. A fast line
	// says what it is in well under the time allowed, and stopping there is what keeps a test
	// down to a couple of hundred megabytes — which is the difference between the far end
	// serving this happily and rate-limiting it. A slow line never reaches it and is timed out
	// by downFor/upFor instead, which is the honest way round.
	measureBytes = int64(64) << 20
	// minWindow is how long the window must last whatever the byte cap says. Without it a very
	// fast line fills the cap in a few hundredths of a second and the phase would be over
	// before there was anything to divide by.
	minWindow = time.Second
)

const (
	latencyProbes = 8
	downStreams   = 4
	upStreams     = 3
	downChunk     = 25 << 20 // bytes per download request; the far end refuses much more
	upChunk       = 25 << 20 // and per upload request
	// phaseBudget is the runaway guard, counting the ramp too: nothing measures forever.
	phaseBudget = int64(512) << 20
	httpTimeout = 90 * time.Second
	// Cooldown is how long the line is left alone between tests. The far end rate-limits by
	// volume, and two back-to-back tests are the surest way to be refused by it.
	Cooldown = 20 * time.Second

	// Method is what the UI tells the reader about the numbers it is showing.
	Method = "Cloudflare's speed endpoints · 4 streams down, 3 up · a 1 s ramp, then up to 4 s or 64 MiB each way"
)

// Result is one finished speed test.
type Result struct {
	At        time.Time `json:"at"`
	DownMbps  float64   `json:"down_mbps"`
	UpMbps    float64   `json:"up_mbps"`
	LatencyMs float64   `json:"latency_ms"`
	JitterMs  float64   `json:"jitter_ms"`
	DownBytes int64     `json:"down_bytes"`
	UpBytes   int64     `json:"up_bytes"`
	TookSec   float64   `json:"took_sec"`
	// Colo is the Cloudflare edge that answered — "SEA" — and IP is this connection's public
	// address as that edge saw it.
	Colo string `json:"colo,omitempty"`
	IP   string `json:"ip,omitempty"`
	// Source is where the download came from, which matters because the fallback is not in the
	// same place as the anycast one. Empty means the usual.
	Source string `json:"source,omitempty"`
	// Note is set when the figures are real but something went wrong along the way — a stream
	// that was refused part-way, most often. The number is what was actually seen; the note is
	// why it may be low.
	Note  string `json:"note,omitempty"`
	Error string `json:"error,omitempty"`
}

// Run is a test in progress, or the last one to have finished.
type Run struct {
	Started time.Time `json:"started"`
	Running bool      `json:"running"`
	// Phase is where it has got to: latency, download, upload, done.
	Phase string `json:"phase"`
	// Done is how far through the whole test, 0–1, so one bar can show all of it.
	Done float64 `json:"done"`
	// Mbps is the live figure for the phase under way — what the needle shows.
	Mbps   float64 `json:"mbps"`
	Bytes  int64   `json:"bytes"`
	Result *Result `json:"result,omitempty"`
}

// Tester runs one speed test at a time and remembers what the last ones said.
type Tester struct {
	mu   sync.RWMutex
	run  Run
	hist []Result
	// lastTry is when a test was last *started*, which is what the cooldown counts from — a
	// test that was refused still used the line, and hammering retries at an endpoint that has
	// just rate-limited us is the one thing guaranteed not to help.
	lastTry time.Time
	cancel  context.CancelFunc
	onDone  func(Result)
	maxHist int
}

func NewTester(onDone func(Result)) *Tester {
	return &Tester{onDone: onDone, maxHist: 50}
}

// Load takes the results kept from before a restart, oldest first.
func (t *Tester) Load(rs []Result) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.hist = append(t.hist[:0], rs...)
	if n := len(t.hist); n > 0 {
		last := t.hist[n-1]
		t.run = Run{Started: last.At, Phase: "done", Done: 1, Result: &last}
		t.lastTry = last.At
	}
}

func (t *Tester) State() Run {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.run
}

func (t *Tester) History() []Result {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return append([]Result{}, t.hist...)
}

func (t *Tester) Running() bool {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.run.Running
}

// Cancel stops a test in progress. Reports whether there was one.
func (t *Tester) Cancel() bool {
	t.mu.Lock()
	c := t.cancel
	t.mu.Unlock()
	if c == nil {
		return false
	}
	c()
	return true
}

var ErrBusy = errors.New("a speed test is already running")

// Start begins a test in the background and returns its opening state.
func (t *Tester) Start() (Run, error) {
	t.mu.Lock()
	if t.run.Running {
		r := t.run
		t.mu.Unlock()
		return r, ErrBusy
	}
	if since := time.Since(t.lastTry); since < Cooldown && !t.lastTry.IsZero() {
		r := t.run
		t.mu.Unlock()
		return r, fmt.Errorf("give the line %d more seconds — testing back to back is what gets this connection rate-limited", int((Cooldown-since).Seconds())+1)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	t.cancel = cancel
	t.lastTry = time.Now()
	t.run = Run{Started: t.lastTry, Running: true, Phase: "latency"}
	r := t.run
	t.mu.Unlock()
	go t.measure(ctx, cancel)
	return r, nil
}

func (t *Tester) set(fn func(*Run)) {
	t.mu.Lock()
	fn(&t.run)
	t.mu.Unlock()
}

func (t *Tester) measure(ctx context.Context, cancel context.CancelFunc) {
	defer cancel()
	started := time.Now()
	res := Result{At: started}

	tr := newTransport(downStreams)
	defer tr.CloseIdleConnections()
	client := &http.Client{Transport: tr, Timeout: httpTimeout}

	// latency, and the meta the same responses carry
	lat, meta, err := t.latency(ctx, client)
	if err == nil {
		res.LatencyMs, res.JitterMs = best(lat), jitter(lat)
		res.Colo, res.IP = meta.Colo, meta.IP
	}
	var note string
	if err == nil {
		t.set(func(r *Run) { r.Phase, r.Done, r.Mbps = "download", 0.1, 0 })
		// Each source in turn: a refusal from the first is exactly what the second is for, and
		// anything else (no route, no DNS) would fail the same way twice, so it stops there.
		for i, src := range downSources {
			res.Source = src.Name
			res.DownMbps, res.DownBytes, note, err = t.phase(ctx, client, "download", downStreams, downFor, 0.1, 0.6, src)
			if err == nil || !refused(err) || i == len(downSources)-1 {
				break
			}
			log.Printf("speed test: %s refused the download, trying %s", src.Name, downSources[i+1].Name)
		}
		res.Note = join(res.Note, note)
	}
	if err == nil {
		t.set(func(r *Run) { r.Phase, r.Done, r.Mbps = "upload", 0.7, 0 })
		res.UpMbps, res.UpBytes, note, err = t.phase(ctx, client, "upload", upStreams, upFor, 0.7, 1.0, Source{Name: "cloudflare"})
		res.Note = join(res.Note, note)
	}
	res.TookSec = round(time.Since(started).Seconds(), 1)
	if err != nil {
		if ctx.Err() != nil {
			res.Error = "cancelled"
		} else {
			res.Error = err.Error()
		}
	}

	t.mu.Lock()
	t.cancel = nil
	t.run = Run{Started: started, Phase: "done", Done: 1, Result: &res, Bytes: res.DownBytes + res.UpBytes}
	if res.Error == "" || res.DownMbps > 0 { // a cancelled test that got a number is still a reading
		t.hist = append(t.hist, res)
		if len(t.hist) > t.maxHist {
			t.hist = t.hist[len(t.hist)-t.maxHist:]
		}
	}
	done, out := t.onDone, res
	t.mu.Unlock()
	if done != nil {
		done(out)
	}
}

// meta is what the edge says about the connection. The address comes from a header of the
// speed endpoint's own; the edge's name is the tail of the ray ID, which every Cloudflare
// response carries ("a41f6720f966a36c-SEA").
type meta struct{ IP, Colo string }

func readMeta(h http.Header) meta {
	m := meta{IP: h.Get("cf-meta-ip")}
	if ray := h.Get("cf-ray"); ray != "" {
		if i := strings.LastIndexByte(ray, '-'); i >= 0 {
			m.Colo = ray[i+1:]
		}
	}
	return m
}

// serverMs is how long the far end says it spent on this request, from Server-Timing. The
// cfL4 entry carries no dur= and contributes nothing, which is what we want — it is a
// description of the connection, not time spent.
func serverMs(h http.Header) float64 {
	var total float64
	for _, v := range h.Values("Server-Timing") {
		for _, part := range strings.Split(v, ",") {
			_, after, ok := strings.Cut(part, "dur=")
			if !ok {
				continue
			}
			end := 0
			for end < len(after) && (after[end] == '.' || (after[end] >= '0' && after[end] <= '9')) {
				end++
			}
			if f, err := strconv.ParseFloat(after[:end], 64); err == nil {
				total += f
			}
		}
	}
	return total
}

// latency times several zero-byte responses on an already-open connection.
func (t *Tester) latency(ctx context.Context, client *http.Client) ([]float64, meta, error) {
	var m meta
	var out []float64
	for i := 0; i <= latencyProbes; i++ { // the first is the handshake, and is not counted
		start := time.Now()
		req, err := http.NewRequestWithContext(ctx, "GET", fmt.Sprintf(cfDownURL, 0), nil)
		if err != nil {
			return nil, m, err
		}
		resp, err := client.Do(req)
		if err != nil {
			return nil, m, err
		}
		_, _ = io.Copy(io.Discard, resp.Body)
		resp.Body.Close()
		if resp.StatusCode != 200 {
			return nil, m, fmt.Errorf("speed endpoint answered %s", resp.Status)
		}
		if i == 0 {
			m = readMeta(resp.Header)
			continue
		}
		rtt := float64(time.Since(start).Microseconds())/1000 - serverMs(resp.Header)
		if rtt < 0.01 { // the far end claimed more time than the whole request took
			rtt = 0.01
		}
		out = append(out, rtt)
		t.set(func(r *Run) { r.Done = 0.02 + 0.08*float64(i)/latencyProbes })
	}
	if len(out) == 0 {
		return nil, m, errors.New("no reply from the speed endpoint")
	}
	return out, m, nil
}

// phase pushes traffic with several streams for a fixed spell and reports the rate over the
// window after the ramp. from and to are where this phase sits in the whole test's progress.
//
// The window ends when the traffic does, not when the clock runs out: if the far end starts
// refusing part-way, the figure is what was actually moved while it was moving, and the note
// says what happened. Coasting to the end of the clock with no bytes arriving would report a
// number a third of the truth.
func (t *Tester) phase(parent context.Context, client *http.Client, dir string, streams int, forHow time.Duration, from, to float64, src Source) (float64, int64, string, error) {
	// parent is the test's own context, which is cancelled when somebody presses Stop. ctx is
	// this phase's, which *this function* cancels to shut the streams down — so the two are
	// kept apart: asking ctx whether the test was cancelled would always say yes by then, and
	// every failure would be reported as "context canceled" whatever really happened.
	ctx, stop := context.WithCancel(parent)
	defer stop()

	var moved atomic.Int64
	var firstErr error
	var errOnce sync.Once
	var alive atomic.Int32
	alive.Store(int32(streams))
	dead := make(chan struct{})
	var deadOnce sync.Once

	for i := 0; i < streams; i++ {
		go func() {
			defer func() {
				if alive.Add(-1) == 0 {
					deadOnce.Do(func() { close(dead) })
				}
			}()
			tries := 0
			for ctx.Err() == nil && moved.Load() < phaseBudget {
				var err error
				if dir == "download" {
					err = stream(ctx, client, src, &moved)
				} else {
					err = upstream(ctx, client, &moved)
				}
				if err == nil {
					tries = 0
					continue
				}
				if ctx.Err() != nil {
					return
				}
				errOnce.Do(func() { firstErr = err })
				if tries++; tries > 2 { // refused three times: this stream is finished
					return
				}
				select {
				case <-time.After(400 * time.Millisecond):
				case <-ctx.Done():
					return
				}
			}
		}()
	}

	// the ramp, measured and thrown away
	select {
	case <-time.After(warmUp):
	case <-dead:
	case <-ctx.Done():
	}
	mark, at := moved.Load(), time.Now()
	last, lastAt := mark, at // the last time bytes were still arriving
	deadline := at.Add(forHow)
	tick := time.NewTicker(200 * time.Millisecond)
	defer tick.Stop()

	for time.Now().Before(deadline) {
		select {
		case <-tick.C:
		case <-dead:
		case <-ctx.Done():
		}
		if parent.Err() != nil {
			break
		}
		now := moved.Load()
		if now > last {
			last, lastAt = now, time.Now()
		}
		// enough has moved to know, and it has been long enough to mean it
		if (now-mark >= measureBytes && time.Since(at) >= minWindow) || now >= phaseBudget {
			break
		}
		if secs := time.Since(at).Seconds(); secs > 0 {
			live := mbps(now-mark, secs)
			part := math.Min(1, secs/forHow.Seconds())
			t.set(func(r *Run) { r.Mbps, r.Bytes, r.Done = live, now, from+(to-from)*part })
		}
		// every stream has given up and nothing more is coming
		if alive.Load() == 0 {
			break
		}
	}
	total, secs := last-mark, lastAt.Sub(at).Seconds()
	stop()

	if total <= 0 {
		switch {
		case firstErr != nil: // what the far end said explains it better than anything else
			return 0, moved.Load(), "", firstErr
		case parent.Err() != nil:
			return 0, moved.Load(), "", parent.Err()
		}
		return 0, moved.Load(), "", fmt.Errorf("nothing moved in the %s phase", dir)
	}
	if secs < 0.02 { // too short a window to divide by; the guard is against dividing by zero,
		// not a judgement about how long a good measurement takes
		return 0, moved.Load(), "", fmt.Errorf("the %s phase was over before it could be measured", dir)
	}
	note := ""
	if firstErr != nil {
		note = fmt.Sprintf("the %s phase was cut short: %s", dir, firstErr)
	}
	return mbps(total, secs), moved.Load(), note, nil
}

// join puts two notes together without leaving an empty one dangling.
func join(a, b string) string {
	switch {
	case a == "":
		return b
	case b == "":
		return a
	}
	return a + "; " + b
}

// stream fetches one chunk from src, counting bytes as they arrive.
func stream(ctx context.Context, client *http.Client, src Source, n *atomic.Int64) error {
	req, err := http.NewRequestWithContext(ctx, "GET", src.URL(downChunk), nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return answered(resp.StatusCode, resp.Status)
	}
	buf := make([]byte, 64<<10)
	for {
		got, err := resp.Body.Read(buf)
		if got > 0 {
			n.Add(int64(got))
		}
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
	}
}

// upstream posts one chunk of nothing in particular, counting bytes as they go out.
func upstream(ctx context.Context, client *http.Client, n *atomic.Int64) error {
	body := &zeros{left: upChunk, n: n}
	req, err := http.NewRequestWithContext(ctx, "POST", upURL, body)
	if err != nil {
		return err
	}
	req.ContentLength = upChunk
	req.Header.Set("Content-Type", "application/octet-stream")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	_, _ = io.Copy(io.Discard, resp.Body)
	resp.Body.Close()
	if resp.StatusCode != 200 && resp.StatusCode != 204 {
		return answered(resp.StatusCode, resp.Status)
	}
	return nil
}

// answered turns the far end's refusal into something worth reading. It rate-limits by
// volume: enough testing in a short while and it stops serving the big transfers, which is
// the only failure of this test anyone here is likely to meet.
func answered(code int, status string) error {
	if code == 429 || code == 403 {
		return errRefused
	}
	return fmt.Errorf("the speed endpoint answered %s", status)
}

// errRefused is the far end rate-limiting this connection, which is the only failure of this
// test anyone here is likely to meet: it limits by volume, and a test moves a few hundred
// megabytes. It is recognisable so the download phase can try its other source.
var errRefused = errors.New("the speed endpoint is rate-limiting this connection — it does that after a lot of testing; leave it a few minutes")

func refused(err error) bool { return errors.Is(err, errRefused) }

// zeros is a request body of a fixed size that costs nothing to produce, counting itself as
// it is read out onto the wire.
type zeros struct {
	left int
	n    *atomic.Int64
}

func (z *zeros) Read(p []byte) (int, error) {
	if z.left <= 0 {
		return 0, io.EOF
	}
	if len(p) > z.left {
		p = p[:z.left]
	}
	for i := range p {
		p[i] = 0
	}
	z.left -= len(p)
	z.n.Add(int64(len(p)))
	return len(p), nil
}

func newTransport(streams int) *http.Transport {
	return &http.Transport{
		Proxy:               http.ProxyFromEnvironment,
		DialContext:         (&net.Dialer{Timeout: 10 * time.Second}).DialContext,
		TLSHandshakeTimeout: 10 * time.Second,
		MaxIdleConns:        streams * 2,
		MaxIdleConnsPerHost: streams * 2,
		// An empty map turns HTTP/2 off: see the note at the top of this file.
		TLSNextProto: map[string]func(string, *tls.Conn) http.RoundTripper{},
	}
}

// probeOutside asks the far end who it thinks we are. One small request.
func probeOutside() Outside {
	tr := newTransport(1)
	defer tr.CloseIdleConnections()
	client := &http.Client{Transport: tr, Timeout: 15 * time.Second}
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, "GET", traceURL, nil)
	if err != nil {
		return Outside{At: time.Now(), Error: err.Error()}
	}
	resp, err := client.Do(req)
	if err != nil {
		return Outside{At: time.Now(), Error: shortErr(err)}
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<10))
	if resp.StatusCode != 200 {
		return Outside{At: time.Now(), Error: "the trace endpoint answered " + resp.Status}
	}
	o := Outside{At: time.Now()}
	for _, line := range strings.Split(string(body), "\n") {
		k, v, ok := strings.Cut(line, "=")
		if !ok {
			continue
		}
		switch k {
		case "ip":
			o.IP = v
		case "colo":
			o.Colo = v
		case "loc":
			o.Country = v
		}
	}
	if o.IP == "" {
		o.Error = "the trace endpoint said nothing about this connection"
	}
	return o
}

// shortErr keeps the part of a network error a person can act on.
func shortErr(err error) string {
	s := err.Error()
	if i := strings.LastIndex(s, ": "); i > 0 && len(s)-i < 60 {
		return s[i+2:]
	}
	return s
}

func mbps(bytes int64, secs float64) float64 {
	return round(float64(bytes)*8/secs/1e6, 1)
}

func best(v []float64) float64 {
	m := v[0]
	for _, x := range v {
		if x < m {
			m = x
		}
	}
	return round(m, 1)
}

// jitter is how much the line wobbles: the average change from one round trip to the next,
// which is the usual definition and, unlike a standard deviation, is not run away with by a
// single slow probe.
func jitter(v []float64) float64 {
	if len(v) < 2 {
		return 0
	}
	var sum float64
	for i := 1; i < len(v); i++ {
		sum += math.Abs(v[i] - v[i-1])
	}
	return round(sum/float64(len(v)-1), 1)
}

func round(v float64, places int) float64 {
	p := math.Pow(10, float64(places))
	return math.Round(v*p) / p
}

// Mbits renders a rate the way the line is sold, for a log line.
func Mbits(v float64) string { return strconv.FormatFloat(v, 'f', 1, 64) + " Mbps" }