// 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": 0, "character": 0}, "end": map[string]any{"line": 0, "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" {
		t.Errorf("action title = %q, want Add RET", 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)
	}
}

// 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)
	}
}