package diag

import (
	"bufio"
	"bytes"
	"io"
	"strings"
	"testing"
	"time"
)

// trickle hands out at most n bytes per Read, so a message arrives in as many
// pieces as a pipe might deliver it in.
type trickle struct {
	r io.Reader
	n int
}

func (t trickle) Read(p []byte) (int, error) {
	if len(p) > t.n {
		p = p[:t.n]
	}
	return t.r.Read(p)
}

func TestReadMessageBodyArrivingInPieces(t *testing.T) {
	body := `{"jsonrpc":"2.0","id":1,"result":{"capabilities":{}}}`
	var buf bytes.Buffer
	if err := writeMessage(&buf, []byte(body)); err != nil {
		t.Fatal(err)
	}
	// One byte at a time is the worst case a pipe can do; bufio's own
	// buffering must not be what makes this pass, so its size is tiny too.
	r := bufio.NewReaderSize(trickle{&buf, 1}, 16)
	got, err := readMessage(r)
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != body {
		t.Errorf("got %q, want %q", got, body)
	}
}

func TestReadMessageHeaderSplitAcrossWrites(t *testing.T) {
	pr, pw := io.Pipe()
	go func() {
		defer pw.Close()
		for _, piece := range []string{"Content-Le", "ngth: 2", "\r\n", "Content-Type: application/vscode-jsonrpc; charset=utf-8\r\n", "\r", "\n", "{", "}"} {
			_, _ = io.WriteString(pw, piece)
			time.Sleep(2 * time.Millisecond)
		}
	}()
	got, err := readMessage(bufio.NewReader(pr))
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != "{}" {
		t.Errorf("got %q", got)
	}
}

func TestReadMessageTwoMessagesBackToBack(t *testing.T) {
	var buf bytes.Buffer
	_ = writeMessage(&buf, []byte(`{"a":1}`))
	_ = writeMessage(&buf, []byte(`{"b":2}`))
	r := bufio.NewReader(&buf)
	for _, want := range []string{`{"a":1}`, `{"b":2}`} {
		got, err := readMessage(r)
		if err != nil {
			t.Fatal(err)
		}
		if string(got) != want {
			t.Errorf("got %q, want %q", got, want)
		}
	}
	if _, err := readMessage(r); err != io.EOF {
		t.Errorf("after the last message: %v, want io.EOF", err)
	}
}

func TestReadMessageRejectsBadFraming(t *testing.T) {
	cases := map[string]string{
		"no content-length": "Content-Type: x\r\n\r\n{}",
		"garbage header":    "hello\r\n\r\n",
		"negative length":   "Content-Length: -1\r\n\r\n",
		"non-numeric":       "Content-Length: abc\r\n\r\n",
		"over the limit":    "Content-Length: 99999999999\r\n\r\n",
		"truncated body":    "Content-Length: 10\r\n\r\n{}",
		"truncated header":  "Content-Length: 2",
	}
	for name, in := range cases {
		t.Run(name, func(t *testing.T) {
			done := make(chan error, 1)
			go func() {
				_, err := readMessage(bufio.NewReader(strings.NewReader(in)))
				done <- err
			}()
			select {
			case err := <-done:
				if err == nil {
					t.Errorf("%q was accepted", in)
				}
				if err == io.EOF {
					t.Errorf("%q reported clean EOF rather than an error", in)
				}
			case <-time.After(time.Second):
				t.Fatal("readMessage hung")
			}
		})
	}
}

func TestReadMessageOversizeIsNotAllocated(t *testing.T) {
	// The limit check has to happen before the make, or a hostile header
	// would be a multi-gigabyte allocation. The over-the-limit case above
	// proves the error; this one proves the body is not read either, by
	// handing over a reader that would block forever if asked.
	pr, _ := io.Pipe()
	header := strings.NewReader("Content-Length: 99999999\r\n\r\n")
	done := make(chan error, 1)
	go func() {
		_, err := readMessage(bufio.NewReader(io.MultiReader(header, pr)))
		done <- err
	}()
	select {
	case err := <-done:
		if err == nil {
			t.Error("oversize message accepted")
		}
	case <-time.After(time.Second):
		t.Fatal("readMessage tried to read an oversize body")
	}
}