package diag

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
)

// maxMessageBytes is the largest JSON-RPC message the client will read. A
// language server that announces a body larger than this is either confused
// or hostile, and either way the right answer is to drop the connection
// rather than to allocate what it asked for. Sixteen megabytes is far past
// any diagnostics batch and comfortably under the point where a reader
// goroutine's allocation shows up in the harness's memory.
const maxMessageBytes = 16 << 20

// readMessage reads one JSON-RPC message in the LSP's base-protocol framing:
// header lines terminated by CRLF, a blank line, then exactly Content-Length
// bytes of body.
//
// The body is read with io.ReadFull rather than a single Read, because a pipe
// hands over whatever the server has written so far and a 40-kilobyte
// diagnostics batch arrives in several pieces. The header is read line by
// line off a bufio.Reader for the same reason: "Content-Le" followed later by
// "ngth: 5" is one header, and a reader that gave up between them would take
// the rest of the stream for garbage.
func readMessage(r *bufio.Reader) ([]byte, error) {
	length := -1
	for {
		line, err := r.ReadString('\n')
		if err != nil {
			if err == io.EOF && line == "" && length < 0 {
				return nil, io.EOF
			}
			return nil, fmt.Errorf("reading header: %w", err)
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" {
			break
		}
		name, value, ok := strings.Cut(line, ":")
		if !ok {
			return nil, fmt.Errorf("malformed header line %q", line)
		}
		if strings.EqualFold(strings.TrimSpace(name), "Content-Length") {
			n, err := strconv.Atoi(strings.TrimSpace(value))
			if err != nil || n < 0 {
				return nil, fmt.Errorf("bad Content-Length %q", strings.TrimSpace(value))
			}
			length = n
		}
		// Content-Type is the only other header the spec defines, and its
		// only permitted charset is utf-8, so there is nothing to do with it.
	}
	if length < 0 {
		return nil, fmt.Errorf("message has no Content-Length header")
	}
	if length > maxMessageBytes {
		return nil, fmt.Errorf("message of %d bytes exceeds the %d-byte limit", length, maxMessageBytes)
	}
	body := make([]byte, length)
	if _, err := io.ReadFull(r, body); err != nil {
		return nil, fmt.Errorf("reading %d-byte body: %w", length, err)
	}
	return body, nil
}

// writeMessage frames body and writes it in one call, so two goroutines
// sharing the writer under a mutex cannot interleave a header with a body.
func writeMessage(w io.Writer, body []byte) error {
	var buf bytes.Buffer
	buf.Grow(len(body) + 32)
	fmt.Fprintf(&buf, "Content-Length: %d\r\n\r\n", len(body))
	buf.Write(body)
	_, err := w.Write(buf.Bytes())
	return err
}