// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

package lsp

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"strings"
	"testing"
)

const cleanDoc = "#include \"textflag.h\"\n" +
	"TEXT ·foo(SB), NOSPLIT, $0\n" +
	"\tMOVQ AX, CX\n" +
	"loop:\n" +
	"\tJMP loop\n" +
	"\tRET\n"

const badDoc = "#include \"textflag.h\"\n" +
	"TEXT ·foo(SB), NOSPLIT, $0\n" +
	"\tNOSUCHINSTR AX, BX\n" +
	"\tJMP missing\n"

// frame renders one Content-Length framed JSON-RPC message.
func frame(id any, method string, params any) string {
	msg := map[string]any{"jsonrpc": "2.0"}
	if id != nil {
		msg["id"] = id
	}
	if method != "" {
		msg["method"] = method
	}
	if params != nil {
		msg["params"] = params
	}
	body, _ := json.Marshal(msg)
	return fmt.Sprintf("Content-Length: %d\r\n\r\n%s", len(body), body)
}

// run feeds input to a server and returns every parsed output message.
func run(t *testing.T, input string) []rpcMessage {
	t.Helper()
	var out bytes.Buffer
	srv := New(strings.NewReader(input), &out)
	if err := srv.Run(); err != nil {
		t.Fatalf("server run: %v", err)
	}
	return readFrames(t, &out)
}

// readFrames parses all framed messages from a buffer.
func readFrames(t *testing.T, r io.Reader) []rpcMessage {
	t.Helper()
	br := bufio.NewReader(r)
	var msgs []rpcMessage
	for {
		length := -1
		for {
			line, err := br.ReadString('\n')
			if err == io.EOF {
				return msgs
			}
			if err != nil {
				t.Fatalf("read header: %v", err)
			}
			line = strings.TrimRight(line, "\r\n")
			if line == "" {
				break
			}
			if k, v, ok := strings.Cut(line, ":"); ok && strings.EqualFold(strings.TrimSpace(k), "Content-Length") {
				length, _ = strconv.Atoi(strings.TrimSpace(v))
			}
		}
		if length < 0 {
			return msgs
		}
		body := make([]byte, length)
		if _, err := io.ReadFull(br, body); err != nil {
			t.Fatalf("read body: %v", err)
		}
		var m rpcMessage
		if err := json.Unmarshal(body, &m); err != nil {
			t.Fatalf("unmarshal: %v", err)
		}
		msgs = append(msgs, m)
	}
}

// session builds a standard scripting of messages around a document.
func session(uri, text string) string {
	var b strings.Builder
	b.WriteString(frame(1, "initialize", map[string]any{"rootUri": ""}))
	b.WriteString(frame(nil, "initialized", map[string]any{}))
	b.WriteString(frame(nil, "textDocument/didOpen", map[string]any{
		"textDocument": map[string]any{
			"uri": uri, "languageId": "gasm", "version": 1, "text": text,
		},
	}))
	return b.String()
}

func findByID(msgs []rpcMessage, n int) *rpcMessage {
	for i := range msgs {
		if msgs[i].ID != nil {
			var id int
			if json.Unmarshal(*msgs[i].ID, &id) == nil && id == n {
				return &msgs[i]
			}
		}
	}
	return nil
}

func findMethod(msgs []rpcMessage, method string) *rpcMessage {
	for i := range msgs {
		if msgs[i].Method == method {
			return &msgs[i]
		}
	}
	return nil
}

func TestInitialize(t *testing.T) {
	msgs := run(t, session("file:///f_amd64.s", cleanDoc)+frame(nil, "exit", nil))
	resp := findByID(msgs, 1)
	if resp == nil {
		t.Fatal("no initialize response")
	}
	var res initializeResult
	if err := json.Unmarshal(mustResult(t, resp), &res); err != nil {
		t.Fatal(err)
	}
	if !res.Capabilities.HoverProvider || res.Capabilities.SemanticTokensProvider == nil {
		t.Fatalf("unexpected capabilities: %+v", res.Capabilities)
	}
}

func TestDiagnosticsClean(t *testing.T) {
	msgs := run(t, session("file:///f_amd64.s", cleanDoc)+frame(nil, "exit", nil))
	pub := findMethod(msgs, "textDocument/publishDiagnostics")
	if pub == nil {
		t.Fatal("no publishDiagnostics notification")
	}
	var p publishDiagnosticsParams
	json.Unmarshal(pub.Params, &p)
	if len(p.Diagnostics) != 0 {
		t.Fatalf("clean doc should have no diagnostics, got %+v", p.Diagnostics)
	}
}

func TestDiagnosticsErrors(t *testing.T) {
	msgs := run(t, session("file:///f_amd64.s", badDoc)+frame(nil, "exit", nil))
	pub := findMethod(msgs, "textDocument/publishDiagnostics")
	if pub == nil {
		t.Fatal("no publishDiagnostics notification")
	}
	var p publishDiagnosticsParams
	json.Unmarshal(pub.Params, &p)
	codes := map[string]bool{}
	for _, d := range p.Diagnostics {
		codes[d.Code] = true
	}
	if !codes["unknown-instruction"] || !codes["undefined-label"] {
		t.Fatalf("expected unknown-instruction and undefined-label, got %+v", p.Diagnostics)
	}
}

func TestCompletion(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(2, "textDocument/completion", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"position":     map[string]any{"line": 2, "character": 1},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 2)
	if resp == nil {
		t.Fatal("no completion response")
	}
	var items []CompletionItem
	if err := json.Unmarshal(mustResult(t, resp), &items); err != nil {
		t.Fatal(err)
	}
	labels := map[string]bool{}
	for _, it := range items {
		labels[it.Label] = true
	}
	for _, want := range []string{"MOVQ", "AX", "TEXT", "NOSPLIT", "loop"} {
		if !labels[want] {
			t.Errorf("completion missing %q", want)
		}
	}
}

func TestHover(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(3, "textDocument/hover", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"position":     map[string]any{"line": 2, "character": 2}, // on MOVQ
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 3)
	if resp == nil {
		t.Fatal("no hover response")
	}
	var h Hover
	if err := json.Unmarshal(mustResult(t, resp), &h); err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(h.Contents.Value, "MOVQ") {
		t.Fatalf("hover = %q, want MOVQ docs", h.Contents.Value)
	}
}

func TestDocumentSymbols(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(4, "textDocument/documentSymbol", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 4)
	if resp == nil {
		t.Fatal("no documentSymbol response")
	}
	var syms []DocumentSymbol
	if err := json.Unmarshal(mustResult(t, resp), &syms); err != nil {
		t.Fatal(err)
	}
	if len(syms) == 0 || syms[0].Name != "foo" {
		t.Fatalf("symbols = %+v, want function foo", syms)
	}
	found := false
	for _, c := range syms[0].Children {
		if c.Name == "loop" {
			found = true
		}
	}
	if !found {
		t.Errorf("function foo should contain label loop: %+v", syms[0].Children)
	}
}

func TestSemanticTokens(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(5, "textDocument/semanticTokens/full", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 5)
	if resp == nil {
		t.Fatal("no semanticTokens response")
	}
	var st SemanticTokens
	if err := json.Unmarshal(mustResult(t, resp), &st); err != nil {
		t.Fatal(err)
	}
	if len(st.Data) == 0 || len(st.Data)%5 != 0 {
		t.Fatalf("semantic tokens data invalid: len=%d", len(st.Data))
	}
	// There must be at least one "function" (mnemonic) and one "comment"-free
	// keyword token; sanity-check that a MOVQ-classified function token exists.
	seenFunction := false
	for i := 3; i < len(st.Data); i += 5 {
		if st.Data[i] == stFunction {
			seenFunction = true
		}
	}
	if !seenFunction {
		t.Error("expected at least one function (mnemonic) semantic token")
	}
}

func TestReferences(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(6, "textDocument/references", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"position":     map[string]any{"line": 4, "character": 5}, // on "loop" in JMP loop
			"context":      map[string]any{"includeDeclaration": true},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 6)
	if resp == nil {
		t.Fatal("no references response")
	}
	var locs []Location
	if err := json.Unmarshal(mustResult(t, resp), &locs); err != nil {
		t.Fatal(err)
	}
	if len(locs) < 2 {
		t.Fatalf("want at least 2 references (definition + use), got %d", len(locs))
	}
}

func TestRename(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(7, "textDocument/rename", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"position":     map[string]any{"line": 3, "character": 0}, // on "loop:" label
			"newName":      "newloop",
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 7)
	if resp == nil {
		t.Fatal("no rename response")
	}
	var edit WorkspaceEdit
	if err := json.Unmarshal(mustResult(t, resp), &edit); err != nil {
		t.Fatal(err)
	}
	edits := edit.Changes["file:///f_amd64.s"]
	if len(edits) < 2 {
		t.Fatalf("want at least 2 edits (definition + use), got %d", len(edits))
	}
	for _, e := range edits {
		if e.NewText != "newloop" {
			t.Errorf("edit newText = %q, want newloop", e.NewText)
		}
	}
}

func TestFormatting(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(8, "textDocument/formatting", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"options":      map[string]any{"tabSize": 4, "insertSpaces": false},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 8)
	if resp == nil {
		t.Fatal("no formatting response")
	}
	var edits []TextEdit
	if err := json.Unmarshal(mustResult(t, resp), &edits); err != nil {
		t.Fatal(err)
	}
	// Formatting should return at most one whole-document edit.
	if len(edits) > 1 {
		t.Fatalf("formatting returned %d edits, want 0 or 1", len(edits))
	}
}

func TestInlayHints(t *testing.T) {
	docWithFrame := "#include \"textflag.h\"\n" +
		"TEXT ·foo(SB), NOSPLIT, $16-24\n" +
		"\tMOVQ AX, CX\n" +
		"\tRET\n"
	in := session("file:///f_amd64.s", docWithFrame) +
		frame(9, "textDocument/inlayHint", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"range": map[string]any{
				"start": map[string]any{"line": 0, "character": 0},
				"end":   map[string]any{"line": 3, "character": 0},
			},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 9)
	if resp == nil {
		t.Fatal("no inlayHint response")
	}
	var hints []InlayHint
	if err := json.Unmarshal(mustResult(t, resp), &hints); err != nil {
		t.Fatal(err)
	}
	if len(hints) == 0 {
		t.Fatal("want at least 1 inlay hint for frame size")
	}
	if !strings.Contains(hints[0].Label, "frame=") {
		t.Errorf("hint label = %q, want frame hint", hints[0].Label)
	}
}

func TestCodeActions(t *testing.T) {
	badDoc := "#include \"textflag.h\"\n" +
		"TEXT ·foo(SB), NOSPLIT, $0\n" +
		"\tMOVQ AX, CX\n"
	in := session("file:///f_amd64.s", badDoc) +
		frame(10, "textDocument/codeAction", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"range":        map[string]any{"start": map[string]any{"line": 0, "character": 0}, "end": map[string]any{"line": 2, "character": 0}},
			"context":      map[string]any{"diagnostics": []map[string]any{{"code": "missing-ret", "range": map[string]any{"start": map[string]any{"line": 1, "character": 0}, "end": map[string]any{"line": 1, "character": 4}}}}},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 10)
	if resp == nil {
		t.Fatal("no codeAction response")
	}
	var actions []CodeAction
	if err := json.Unmarshal(mustResult(t, resp), &actions); err != nil {
		t.Fatal(err)
	}
	if len(actions) == 0 {
		t.Fatal("want at least 1 code action for missing-ret")
	}
	if actions[0].Title != "Add RET to foo" {
		t.Errorf("action title = %q, want Add RET to foo", actions[0].Title)
	}
}

func TestSignatureHelp(t *testing.T) {
	doc := "#include \"textflag.h\"\n" +
		"TEXT ·add(SB), NOSPLIT, $0-24\n" +
		"\tMOVQ AX, CX\n" +
		"\tRET\n" +
		"TEXT ·caller(SB), NOSPLIT, $0\n" +
		"\tCALL ·add(SB)\n" +
		"\tRET\n"
	in := session("file:///f_amd64.s", doc) +
		frame(11, "textDocument/signatureHelp", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"position":     map[string]any{"line": 5, "character": 6},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 11)
	if resp == nil {
		t.Fatal("no signatureHelp response")
	}
	var sh SignatureHelp
	if err := json.Unmarshal(mustResult(t, resp), &sh); err != nil {
		t.Fatal(err)
	}
	if len(sh.Signatures) == 0 {
		t.Fatal("want at least 1 signature")
	}
	if !strings.Contains(sh.Signatures[0].Label, "add") {
		t.Errorf("signature label = %q, want add", sh.Signatures[0].Label)
	}
}

func TestDocumentHighlights(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(12, "textDocument/documentHighlight", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"position":     map[string]any{"line": 4, "character": 5},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 12)
	if resp == nil {
		t.Fatal("no documentHighlight response")
	}
	var highlights []DocumentHighlight
	if err := json.Unmarshal(mustResult(t, resp), &highlights); err != nil {
		t.Fatal(err)
	}
	if len(highlights) < 2 {
		t.Fatalf("want at least 2 highlights (definition + use), got %d", len(highlights))
	}
}

func TestWorkspaceSymbols(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(13, "workspace/symbol", map[string]any{
			"query": "foo",
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 13)
	if resp == nil {
		t.Fatal("no workspace/symbol response")
	}
	var syms []WorkspaceSymbol
	if err := json.Unmarshal(mustResult(t, resp), &syms); err != nil {
		t.Fatal(err)
	}
	found := false
	for _, s := range syms {
		if s.Name == "foo" {
			found = true
		}
	}
	if !found {
		t.Fatalf("want symbol foo in results, got %+v", syms)
	}
}

func TestMethodNotFound(t *testing.T) {
	in := frame(9, "bogus/method", map[string]any{}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 9)
	if resp == nil || resp.Error == nil || resp.Error.Code != errMethodNotFound {
		t.Fatalf("expected method-not-found error, got %+v", resp)
	}
}

// TestMalformedContentLengthAnsweredAndRecovered pins the framing contract:
// an unparseable Content-Length is answered with a -32700 parse error against
// the null id, and the server keeps serving the next well-formed request
// instead of exiting the process.
func TestMalformedContentLengthAnsweredAndRecovered(t *testing.T) {
	in := "Content-Length: 12x\r\n\r\n" +
		frame(1, "initialize", map[string]any{}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	if n := countParseErrors(msgs); n != 1 {
		t.Fatalf("want exactly one -32700 response, got %d", n)
	}
	if findByID(msgs, 1) == nil {
		t.Fatal("server did not serve the request after the malformed header")
	}
}

// TestMalformedJSONBodyAnsweredAndRecovered pins the same contract for a
// body that does not decode as JSON-RPC.
func TestMalformedJSONBodyAnsweredAndRecovered(t *testing.T) {
	body := "{not json"
	in := fmt.Sprintf("Content-Length: %d\r\n\r\n%s", len(body), body) +
		frame(1, "initialize", map[string]any{}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	if n := countParseErrors(msgs); n != 1 {
		t.Fatalf("want exactly one -32700 response, got %d", n)
	}
	if findByID(msgs, 1) == nil {
		t.Fatal("server did not serve the request after the malformed body")
	}
}

func countParseErrors(msgs []rpcMessage) int {
	n := 0
	for i := range msgs {
		if msgs[i].Error != nil && msgs[i].Error.Code == errParse {
			n++
		}
	}
	return n
}

// TestNullResultPresent pins the response envelope: JSON-RPC 2.0 requires a
// result member even when its value is null (shutdown, hover over nothing),
// while notifications must not grow one.
func TestNullResultPresent(t *testing.T) {
	var out bytes.Buffer
	in := frame(1, "shutdown", nil) + frame(nil, "exit", nil)
	srv := New(strings.NewReader(in), &out)
	if err := srv.Run(); err != nil {
		t.Fatalf("server run: %v", err)
	}
	if !strings.Contains(out.String(), `"result":null`) {
		t.Fatalf("shutdown response must serialise an explicit null result, got %q", out.String())
	}
}

// TestHoverNullResultPresent checks the same envelope on a handler whose
// result is genuinely null: hover over a position with no symbol.
func TestHoverNullResultPresent(t *testing.T) {
	var out bytes.Buffer
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(3, "textDocument/hover", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"position":     map[string]any{"line": 0, "character": 0},
		}) + frame(nil, "exit", nil)
	srv := New(strings.NewReader(in), &out)
	if err := srv.Run(); err != nil {
		t.Fatalf("server run: %v", err)
	}
	raw := out.String()
	if !strings.Contains(raw, `"result":null`) {
		t.Fatalf("null hover result must serialise as an explicit null, got %q", raw)
	}
	// The publishDiagnostics notification must not carry a result member.
	for _, m := range readFrames(t, strings.NewReader(raw)) {
		if m.Method == "textDocument/publishDiagnostics" && m.Result != nil {
			t.Fatalf("notification grew a result member: %+v", m)
		}
	}
}

// applyEdits applies text edits whose character columns are rune columns
// (the middle dot makes byte and rune offsets diverge).
func applyEdits(text string, edits []TextEdit) string {
	lines := strings.Split(text, "\n")
	for _, e := range edits {
		s, en := e.Range.Start, e.Range.End
		if s.Line != en.Line || s.Line >= len(lines) {
			continue
		}
		runes := []rune(lines[s.Line])
		if s.Character > len(runes) || en.Character > len(runes) {
			continue
		}
		lines[s.Line] = string(runes[:s.Character]) + e.NewText + string(runes[en.Character:])
	}
	return strings.Join(lines, "\n")
}

// TestRenameMiddleDotCoversFullToken pins the reference ranges: a symbol
// reference carries the `·` prefix, so every edit must span the whole raw
// token or the rename corrupts the source (`·helpe` + `r` shapes).
func TestRenameMiddleDotCoversFullToken(t *testing.T) {
	doc := "#include \"textflag.h\"\n" +
		"TEXT ·helper(SB), NOSPLIT, $0\n" +
		"\tRET\n" +
		"TEXT ·caller(SB), NOSPLIT, $0\n" +
		"\tCALL ·helper(SB)\n" +
		"\tRET\n"
	in := session("file:///f_amd64.s", doc) +
		frame(7, "textDocument/rename", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"position":     map[string]any{"line": 4, "character": 7}, // on ·helper in CALL
			"newName":      "renamed",
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 7)
	if resp == nil {
		t.Fatal("no rename response")
	}
	var edit WorkspaceEdit
	if err := json.Unmarshal(mustResult(t, resp), &edit); err != nil {
		t.Fatal(err)
	}
	edits := edit.Changes["file:///f_amd64.s"]
	if len(edits) != 2 {
		t.Fatalf("want 2 edits (definition + call), got %d", len(edits))
	}
	for _, e := range edits {
		if got := e.Range.End.Character - e.Range.Start.Character; got != runeLen("·helper") {
			t.Errorf("edit range covers %d characters, want the full %d of ·helper", got, runeLen("·helper"))
		}
	}
	got := applyEdits(doc, edits)
	if strings.Contains(got, "·helper") || strings.Contains(got, "·renamed") {
		t.Fatalf("rename left a partial spelling behind: %q", got)
	}
	if !strings.Contains(got, "TEXT renamed(SB)") || !strings.Contains(got, "CALL renamed(SB)") {
		t.Fatalf("rename did not rewrite both sites: %q", got)
	}
}

// TestDocumentHighlightsMiddleDot pins the trim the handler was missing:
// the word under the cursor keeps its `·` prefix while symbol names are
// stored without it, so the untrimmed comparison returned nothing.
func TestDocumentHighlightsMiddleDot(t *testing.T) {
	doc := "#include \"textflag.h\"\n" +
		"TEXT ·helper(SB), NOSPLIT, $0\n" +
		"\tRET\n" +
		"TEXT ·caller(SB), NOSPLIT, $0\n" +
		"\tCALL ·helper(SB)\n" +
		"\tRET\n"
	in := session("file:///f_amd64.s", doc) +
		frame(12, "textDocument/documentHighlight", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"position":     map[string]any{"line": 4, "character": 7},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 12)
	if resp == nil {
		t.Fatal("no documentHighlight response")
	}
	var highlights []DocumentHighlight
	if err := json.Unmarshal(mustResult(t, resp), &highlights); err != nil {
		t.Fatal(err)
	}
	if len(highlights) != 2 {
		t.Fatalf("want the definition and the call highlighted, got %d", len(highlights))
	}
}

// TestSymbolsSkipMalformedTextPlaceholder checks both symbol listings
// against a mid-edit buffer whose first TEXT line has no symbol yet: the
// placeholder is skipped and the remaining symbols are still served.
func TestSymbolsSkipMalformedTextPlaceholder(t *testing.T) {
	doc := "TEXT $0\n" +
		"\tMOVQ AX, BX\n" +
		"TEXT ·ok(SB), NOSPLIT, $0\n" +
		"\tRET\n"
	in := session("file:///f_amd64.s", doc) +
		frame(4, "textDocument/documentSymbol", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
		}) +
		frame(13, "workspace/symbol", map[string]any{"query": "ok"}) +
		frame(nil, "exit", nil)
	msgs := run(t, in)

	dresp := findByID(msgs, 4)
	if dresp == nil || dresp.Error != nil {
		t.Fatalf("documentSymbol failed on a mid-edit buffer: %+v", dresp)
	}
	var syms []DocumentSymbol
	if err := json.Unmarshal(mustResult(t, dresp), &syms); err != nil {
		t.Fatal(err)
	}
	if len(syms) != 1 || syms[0].Name != "ok" {
		t.Fatalf("want only the ·ok function, got %+v", syms)
	}

	wresp := findByID(msgs, 13)
	if wresp == nil || wresp.Error != nil {
		t.Fatalf("workspace/symbol failed on a mid-edit buffer: %+v", wresp)
	}
	var wsyms []WorkspaceSymbol
	if err := json.Unmarshal(mustResult(t, wresp), &wsyms); err != nil {
		t.Fatal(err)
	}
	if len(wsyms) != 1 || wsyms[0].Name != "ok" {
		t.Fatalf("want only the ·ok symbol, got %+v", wsyms)
	}
}

// TestUTF16PositionEncoding exercises the boundary conversions directly:
// one astral rune is one rune but two UTF-16 code units.
func TestUTF16PositionEncoding(t *testing.T) {
	const line = "\U0001F389ab" // emoji, then two ASCII letters
	if got := utf16Len(line); got != 4 {
		t.Errorf("utf16Len = %d, want 4", got)
	}
	if got := utf16Column(line, 3); got != 4 {
		t.Errorf("utf16Column(col 3) = %d, want 4", got)
	}
	if got := runeColumn(line, 4); got != 3 {
		t.Errorf("runeColumn(units 4) = %d, want 3", got)
	}
	rng := clientRange("x\n"+line, Range{
		Start: Position{Line: 1, Character: 1},
		End:   Position{Line: 1, Character: 3},
	})
	if rng.Start.Character != 2 || rng.End.Character != 4 {
		t.Errorf("clientRange = %+v, want 2..4 UTF-16 units", rng)
	}
	// wordAt takes UTF-16 units from the client and lands on the right rune.
	w, wr := wordAt("loop: // \U0001F389 x\n", Position{Line: 0, Character: 12})
	if w != "x" {
		t.Fatalf("wordAt = %q, want x", w)
	}
	if wr.Start.Character != 11 || wr.End.Character != 12 {
		t.Errorf("wordAt range = %+v, want 11..12 (rune columns are returned unconverted)", wr)
	}
}

// TestSemanticTokensAstralCommentLength checks that an astral rune in a
// comment is counted as two UTF-16 code units in the token length.
func TestSemanticTokensAstralCommentLength(t *testing.T) {
	comment := "// \U0001F389 ok"
	doc := "#include \"textflag.h\"\nTEXT ·f(SB), NOSPLIT, $0\n\tRET " + comment + "\n"
	in := session("file:///f_amd64.s", doc) +
		frame(5, "textDocument/semanticTokens/full", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 5)
	if resp == nil {
		t.Fatal("no semanticTokens response")
	}
	var st SemanticTokens
	if err := json.Unmarshal(mustResult(t, resp), &st); err != nil {
		t.Fatal(err)
	}
	// Decode the delta stream: the RET line carries the comment token.
	line, char := 0, 0
	found := false
	for i := 0; i+4 < len(st.Data); i += 5 {
		if st.Data[i] != 0 {
			line += st.Data[i]
			char = st.Data[i+1]
		} else {
			char += st.Data[i+1]
		}
		if line == 2 && st.Data[i+3] == stComment {
			found = true
			if st.Data[i+2] != utf16Len(comment) {
				t.Errorf("comment token length = %d, want %d UTF-16 code units", st.Data[i+2], utf16Len(comment))
			}
			if want := utf16Len("\tRET "); char != want {
				t.Errorf("comment token char = %d, want %d UTF-16 code units", char, want)
			}
		}
	}
	if !found {
		t.Fatalf("no comment token on the RET line: %v", st.Data)
	}
}

// TestSemanticTokensPipeOperator checks that the `|` joining TEXT flags is
// classified as an operator instead of falling out of the switch unstyled.
func TestSemanticTokensPipeOperator(t *testing.T) {
	doc := "#include \"textflag.h\"\nTEXT ·f(SB), NOSPLIT|DUPOK, $0\n\tRET\n"
	in := session("file:///f_amd64.s", doc) +
		frame(5, "textDocument/semanticTokens/full", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 5)
	if resp == nil {
		t.Fatal("no semanticTokens response")
	}
	var st SemanticTokens
	if err := json.Unmarshal(mustResult(t, resp), &st); err != nil {
		t.Fatal(err)
	}
	// Decode the delta stream and look for a one-character operator token
	// on the TEXT line.
	line, char := 0, 0
	found := false
	for i := 0; i+4 < len(st.Data); i += 5 {
		if st.Data[i] != 0 {
			line += st.Data[i]
			char = st.Data[i+1]
		} else {
			char += st.Data[i+1]
		}
		if line == 1 && st.Data[i+2] == 1 && st.Data[i+3] == stOperator {
			found = true
		}
	}
	if !found {
		t.Fatalf("no operator token found for the | in the flag list: %v", st.Data)
	}
}

// mustResult re-marshals a response result into raw JSON for typed decoding.
func mustResult(t *testing.T, m *rpcMessage) []byte {
	t.Helper()
	b, err := json.Marshal(m.Result)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

func TestPullDiagnostics(t *testing.T) {
	in := session("file:///f_amd64.s", badDoc) +
		frame(20, "textDocument/diagnostic", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 20)
	if resp == nil {
		t.Fatal("no diagnostic response")
	}
	var report documentDiagnosticReport
	if err := json.Unmarshal(mustResult(t, resp), &report); err != nil {
		t.Fatal(err)
	}
	if report.Kind != "full" {
		t.Errorf("report kind = %q, want full", report.Kind)
	}
	found := false
	for _, d := range report.Items {
		if d.Code == "unknown-instruction" {
			found = true
		}
	}
	if !found {
		t.Errorf("want an unknown-instruction diagnostic, got %+v", report.Items)
	}
}

func TestDocumentLinks(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(21, "textDocument/documentLink", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 21)
	if resp == nil {
		t.Fatal("no documentLink response")
	}
	var links []DocumentLink
	if err := json.Unmarshal(mustResult(t, resp), &links); err != nil {
		t.Fatal(err)
	}
	if len(links) != 1 {
		t.Fatalf("want 1 document link, got %d", len(links))
	}
	if !strings.HasSuffix(links[0].Target, "pkg/include/textflag.h") {
		t.Errorf("link target = %q, want the GOROOT pkg/include header", links[0].Target)
	}
	if links[0].Range.Start.Line != 0 || links[0].Range.Start.Character != 9 {
		t.Errorf("link range start = %+v, want line 0 character 9", links[0].Range.Start)
	}
}

func TestFoldingRanges(t *testing.T) {
	in := session("file:///f_amd64.s", cleanDoc) +
		frame(22, "textDocument/foldingRange", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 22)
	if resp == nil {
		t.Fatal("no foldingRange response")
	}
	var ranges []FoldingRange
	if err := json.Unmarshal(mustResult(t, resp), &ranges); err != nil {
		t.Fatal(err)
	}
	if len(ranges) != 1 {
		t.Fatalf("want 1 folding range, got %d", len(ranges))
	}
	if ranges[0].StartLine != 1 || ranges[0].EndLine != 5 {
		t.Errorf("folding range = %d-%d, want 1-5", ranges[0].StartLine, ranges[0].EndLine)
	}
	if ranges[0].Kind != "region" {
		t.Errorf("folding kind = %q, want region", ranges[0].Kind)
	}
}

// TestInitializeVersion checks the version reported in the initialize
// result when the caller stamps one.
func TestInitializeVersion(t *testing.T) {
	var out bytes.Buffer
	srv := New(strings.NewReader(frame(1, "initialize", map[string]any{})+frame(nil, "exit", nil)), &out)
	srv.SetVersion("9.9.9")
	if err := srv.Run(); err != nil {
		t.Fatalf("server run: %v", err)
	}
	msgs := readFrames(t, &out)
	resp := findByID(msgs, 1)
	if resp == nil {
		t.Fatal("no initialize response")
	}
	var res initializeResult
	if err := json.Unmarshal(mustResult(t, resp), &res); err != nil {
		t.Fatal(err)
	}
	if res.ServerInfo["version"] != "9.9.9" {
		t.Errorf("serverInfo = %v, want version 9.9.9", res.ServerInfo)
	}
}

// TestSyntaxDiagnosticsPublished checks that parse errors reach the editor
// as error-severity diagnostics with the syntax code.
func TestSyntaxDiagnosticsPublished(t *testing.T) {
	msgs := run(t, session("file:///f_amd64.s", "TEXT $\n")+frame(nil, "exit", nil))
	pub := findMethod(msgs, "textDocument/publishDiagnostics")
	if pub == nil {
		t.Fatal("no publishDiagnostics notification")
	}
	var p publishDiagnosticsParams
	json.Unmarshal(pub.Params, &p)
	found := false
	for _, d := range p.Diagnostics {
		if d.Code == "syntax" && d.Severity == sevError {
			found = true
		}
	}
	if !found {
		t.Fatalf("expected a syntax error diagnostic, got %+v", p.Diagnostics)
	}
}

// TestDispatchRecoversFromPanic exercises the per-message recover: a panic
// inside a handler is answered as an internal error instead of taking the
// server down.
func TestDispatchRecoversFromPanic(t *testing.T) {
	var out bytes.Buffer
	srv := New(strings.NewReader(""), &out)
	srv.docs = nil // force a nil-map write inside didOpen
	raw := json.RawMessage(`{"textDocument":{"uri":"file:///x.s","text":"RET"}}`)
	srv.dispatch(&rpcMessage{ID: rawID(t, 7), Method: "textDocument/didOpen", Params: raw})
	msgs := readFrames(t, &out)
	resp := findByID(msgs, 7)
	if resp == nil {
		t.Fatal("no error response after panic")
	}
	if resp.Error == nil || resp.Error.Code != errInternalError {
		t.Fatalf("error = %+v, want internal error", resp.Error)
	}
}

func rawID(t *testing.T, n int) *json.RawMessage {
	t.Helper()
	b, err := json.Marshal(n)
	if err != nil {
		t.Fatal(err)
	}
	raw := json.RawMessage(b)
	return &raw
}

// TestURIDecoding pins the percent-decoding of file URIs: clients encode
// non-ASCII paths, and the decoded form is what resolves on disk.
func TestURIDecoding(t *testing.T) {
	got := uriPath("file:///home/petrbalvin/Repozit%C3%A1%C5%99e/k.s")
	if got != "/home/petrbalvin/Repozitáře/k.s" {
		t.Errorf("uriPath = %q", got)
	}
	if got := uriPath("/plain/path.s"); got != "/plain/path.s" {
		t.Errorf("uriPath plain = %q", got)
	}
}

// TestCodeActionsTargetsFlaggedFunctionOnly checks that a missing-ret
// diagnostic offers an edit for the flagged function only, even when the
// file defines several functions.
func TestCodeActionsTargetsFlaggedFunctionOnly(t *testing.T) {
	doc := "TEXT \u00b7first(SB), NOSPLIT, $0\n" +
		"\tMOVQ AX, CX\n" +
		"\tRET\n" +
		"TEXT \u00b7second(SB), NOSPLIT, $0\n" +
		"\tMOVQ AX, CX\n"
	in := session("file:///f_amd64.s", doc) +
		frame(11, "textDocument/codeAction", map[string]any{
			"textDocument": map[string]any{"uri": "file:///f_amd64.s"},
			"range":        map[string]any{"start": map[string]any{"line": 0, "character": 0}, "end": map[string]any{"line": 4, "character": 0}},
			"context":      map[string]any{"diagnostics": []map[string]any{{"code": "missing-ret", "range": map[string]any{"start": map[string]any{"line": 3, "character": 0}, "end": map[string]any{"line": 3, "character": 4}}}}},
		}) + frame(nil, "exit", nil)
	msgs := run(t, in)
	resp := findByID(msgs, 11)
	if resp == nil {
		t.Fatal("no codeAction response")
	}
	var actions []CodeAction
	if err := json.Unmarshal(mustResult(t, resp), &actions); err != nil {
		t.Fatal(err)
	}
	if len(actions) != 1 {
		t.Fatalf("actions = %d, want 1", len(actions))
	}
	if actions[0].Title != "Add RET to second" {
		t.Errorf("action title = %q, want Add RET to second", actions[0].Title)
	}
	if !strings.Contains(actions[0].Edit.Changes["file:///f_amd64.s"][0].NewText, "\tRET\nTEXT \u00b7second") {
		t.Errorf("edit does not insert RET at the end of second:\n%s", actions[0].Edit.Changes["file:///f_amd64.s"][0].NewText)
	}
}

// TestCrossFileDefinitionAndReferences opens two documents: docA calls
// ·helper(SB), docB defines it.  Definition must jump to docB and references
// must collect the call site in docA plus the definition in docB.
func TestCrossFileDefinitionAndReferences(t *testing.T) {
	docA := "TEXT \u00b7caller(SB), NOSPLIT, $0\n" +
		"\tCALL \u00b7helper(SB)\n" +
		"\tRET\n"
	docB := "TEXT \u00b7helper(SB), NOSPLIT, $0\n" +
		"\tRET\n"
	in := frame(1, "initialize", map[string]any{}) +
		frame(nil, "initialized", map[string]any{}) +
		frame(nil, "textDocument/didOpen", map[string]any{
			"textDocument": map[string]any{"uri": "file:///a_amd64.s", "languageId": "gasm", "version": 1, "text": docA},
		}) +
		frame(nil, "textDocument/didOpen", map[string]any{
			"textDocument": map[string]any{"uri": "file:///b_amd64.s", "languageId": "gasm", "version": 1, "text": docB},
		}) +
		frame(2, "textDocument/definition", map[string]any{
			"textDocument": map[string]any{"uri": "file:///a_amd64.s"},
			"position":     map[string]any{"line": 1, "character": 8}, // on helper in CALL ·helper(SB)
		}) +
		frame(3, "textDocument/references", map[string]any{
			"textDocument": map[string]any{"uri": "file:///b_amd64.s"},
			"position":     map[string]any{"line": 0, "character": 7}, // on helper in TEXT ·helper(SB)
			"context":      map[string]any{"includeDeclaration": true},
		}) +
		frame(nil, "exit", nil)
	msgs := run(t, in)

	dresp := findByID(msgs, 2)
	if dresp == nil {
		t.Fatal("no definition response")
	}
	var locs []Location
	if err := json.Unmarshal(mustResult(t, dresp), &locs); err != nil {
		t.Fatal(err)
	}
	if len(locs) != 1 || locs[0].URI != "file:///b_amd64.s" || locs[0].Range.Start.Line != 0 {
		t.Fatalf("definition = %+v, want the TEXT in b_amd64.s line 0", locs)
	}

	rresp := findByID(msgs, 3)
	if rresp == nil {
		t.Fatal("no references response")
	}
	locs = nil
	if err := json.Unmarshal(mustResult(t, rresp), &locs); err != nil {
		t.Fatal(err)
	}
	if len(locs) != 2 {
		t.Fatalf("references = %+v, want the definition in b_amd64.s and the call in a_amd64.s", locs)
	}
	byURI := map[string]int{}
	for _, l := range locs {
		byURI[l.URI]++
	}
	if byURI["file:///a_amd64.s"] != 1 || byURI["file:///b_amd64.s"] != 1 {
		t.Errorf("references by uri = %v, want one in each file", byURI)
	}
}