package diag

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

// fakeServer is a language server that lives in the test process and speaks
// the wire protocol over pipes, so every ugly path — silence, a half-written
// header, an early exit — can be produced on demand without a real server
// installed. serve is called for each message the client sends.
type fakeServer struct {
	in    *bufio.Reader
	out   io.WriteCloser
	outMu sync.Mutex

	mu       sync.Mutex
	received []incoming
	replies  []incoming // responses the client sent to our requests
	closed   bool
}

type serveFunc func(f *fakeServer, msg incoming)

// startFake wires a session to a fakeServer. idle is the session's idle
// timeout, 0 for none.
func startFake(t *testing.T, serve serveFunc, idle time.Duration, onExit func(*session)) (*session, *fakeServer) {
	t.Helper()
	toServer, clientStdin := io.Pipe()
	clientStdout, fromServer := io.Pipe()
	f := &fakeServer{in: bufio.NewReader(toServer), out: fromServer}
	go f.loop(serve)
	kill := func() {
		f.close()
		_ = toServer.Close()
	}
	spec := serverSpec{Bin: "fakels", Languages: map[string]string{".go": "go"}}
	s := newSession(spec, t.TempDir(), clientStdin, clientStdout, kill, nil, idle, onExit)
	t.Cleanup(func() {
		done := make(chan struct{})
		go func() { s.Close(); close(done) }()
		select {
		case <-done:
		case <-time.After(3 * time.Second):
			t.Error("session.Close hung")
		}
	})
	return s, f
}

func (f *fakeServer) loop(serve serveFunc) {
	for {
		body, err := readMessage(f.in)
		if err != nil {
			return
		}
		var msg incoming
		if err := json.Unmarshal(body, &msg); err != nil {
			return
		}
		f.mu.Lock()
		if msg.Method == "" {
			f.replies = append(f.replies, msg)
		} else {
			f.received = append(f.received, msg)
		}
		f.mu.Unlock()
		if msg.Method != "" && serve != nil {
			serve(f, msg)
		}
	}
}

func (f *fakeServer) send(v any) {
	body, _ := json.Marshal(v)
	f.outMu.Lock()
	defer f.outMu.Unlock()
	if f.closed {
		return
	}
	_ = writeMessage(f.out, body)
}

// raw writes bytes with no framing, for producing broken streams.
func (f *fakeServer) raw(s string) {
	f.outMu.Lock()
	defer f.outMu.Unlock()
	if !f.closed {
		_, _ = io.WriteString(f.out, s)
	}
}

func (f *fakeServer) close() {
	f.outMu.Lock()
	defer f.outMu.Unlock()
	if !f.closed {
		f.closed = true
		_ = f.out.Close()
	}
}

func (f *fakeServer) respond(id json.RawMessage, result any) {
	f.send(map[string]any{"jsonrpc": "2.0", "id": id, "result": result})
}

func (f *fakeServer) notify(method string, params any) {
	f.send(map[string]any{"jsonrpc": "2.0", "method": method, "params": params})
}

func (f *fakeServer) request(id int, method string, params any) {
	f.send(map[string]any{"jsonrpc": "2.0", "id": id, "method": method, "params": params})
}

func (f *fakeServer) methods() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]string, len(f.received))
	for i, m := range f.received {
		out[i] = m.Method
	}
	return out
}

func (f *fakeServer) publish(uri string, diags []map[string]any) {
	if diags == nil {
		diags = []map[string]any{}
	}
	f.notify("textDocument/publishDiagnostics", map[string]any{"uri": uri, "diagnostics": diags})
}

// lspDiag builds one LSP diagnostic at 0-based line and character.
func lspDiag(line, char, severity int, msg string) map[string]any {
	return map[string]any{
		"range": map[string]any{
			"start": map[string]any{"line": line, "character": char},
			"end":   map[string]any{"line": line, "character": char + 1},
		},
		"severity": severity,
		"message":  msg,
	}
}

// docText pulls the file content out of a didOpen or didChange.
func docText(msg incoming) (uri, text string) {
	var open struct {
		TextDocument struct {
			URI  string `json:"uri"`
			Text string `json:"text"`
		} `json:"textDocument"`
		ContentChanges []struct {
			Text string `json:"text"`
		} `json:"contentChanges"`
	}
	_ = json.Unmarshal(msg.Params, &open)
	text = open.TextDocument.Text
	if len(open.ContentChanges) > 0 {
		text = open.ContentChanges[0].Text
	}
	return open.TextDocument.URI, text
}

// wellBehaved answers like a real server: initialize gets capabilities, an
// opened or changed document gets one error per line containing BUG, one
// warning per WARN and one hint per HINT, shutdown gets null and exit closes
// the connection.
func wellBehaved(f *fakeServer, msg incoming) {
	switch msg.Method {
	case "initialize":
		f.respond(msg.ID, map[string]any{"capabilities": map[string]any{}})
	case "textDocument/didOpen", "textDocument/didChange":
		uri, text := docText(msg)
		var diags []map[string]any
		for i, line := range strings.Split(text, "\n") {
			switch {
			case strings.Contains(line, "BUG"):
				diags = append(diags, lspDiag(i, strings.Index(line, "BUG"), 1, "this is a bug"))
			case strings.Contains(line, "WARN"):
				diags = append(diags, lspDiag(i, 0, 2, "this is a warning"))
			case strings.Contains(line, "HINT"):
				diags = append(diags, lspDiag(i, 0, 4, "this is a hint"))
			}
		}
		f.publish(uri, diags)
	case "shutdown":
		f.respond(msg.ID, nil)
	case "exit":
		f.close()
	}
}

func writeTemp(t *testing.T, name, content string) string {
	t.Helper()
	p := filepath.Join(t.TempDir(), name)
	if err := os.WriteFile(p, []byte(content), 0o644); err != nil {
		t.Fatal(err)
	}
	return p
}

func shortCtx(t *testing.T, d time.Duration) context.Context {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), d)
	t.Cleanup(cancel)
	return ctx
}

func TestSessionReportsDiagnosticsForTheOpenedFile(t *testing.T) {
	s, f := startFake(t, wellBehaved, 0, nil)
	path := writeTemp(t, "main.go", "package main\n\nfunc main() {\n\tx := BUG\n\tWARN\n\tHINT\n}\n")

	ds, err := s.Check(shortCtx(t, 2*time.Second), path)
	if err != nil {
		t.Fatalf("Check: %v", err)
	}
	if len(ds) != 3 {
		t.Fatalf("got %d diagnostics, want 3: %+v", len(ds), ds)
	}
	sortForReport(ds, path)
	bug := ds[0]
	if bug.File != path || bug.Line != 4 || bug.Col != 7 || bug.Severity != SeverityError || bug.Message != "this is a bug" || bug.Source != "fakels" {
		t.Errorf("bug diagnostic wrong: %+v", bug)
	}
	if ds[1].Severity != SeverityWarning || ds[1].Line != 5 {
		t.Errorf("warning wrong: %+v", ds[1])
	}
	if ds[2].Severity != SeverityHint {
		t.Errorf("hint wrong: %+v", ds[2])
	}

	got := f.methods()
	want := []string{"initialize", "initialized", "textDocument/didOpen", "textDocument/didSave"}
	if strings.Join(got, ",") != strings.Join(want, ",") {
		t.Errorf("server saw %v, want %v", got, want)
	}
}

func TestSessionSecondCheckIsADidChange(t *testing.T) {
	s, f := startFake(t, wellBehaved, 0, nil)
	path := writeTemp(t, "main.go", "BUG\n")
	ctx := shortCtx(t, 4*time.Second)

	if ds, err := s.Check(ctx, path); err != nil || len(ds) != 1 {
		t.Fatalf("first check: %v, %d diagnostics", err, len(ds))
	}
	if err := os.WriteFile(path, []byte("clean now\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	ds, err := s.Check(ctx, path)
	if err != nil {
		t.Fatalf("second check: %v", err)
	}
	if len(ds) != 0 {
		t.Errorf("second check returned stale diagnostics: %+v", ds)
	}
	methods := strings.Join(f.methods(), ",")
	if strings.Count(methods, "textDocument/didOpen") != 1 || !strings.Contains(methods, "textDocument/didChange") {
		t.Errorf("expected one didOpen then a didChange, got %s", methods)
	}
}

func TestSessionCleanFileReturnsPromptly(t *testing.T) {
	s, _ := startFake(t, wellBehaved, 0, nil)
	path := writeTemp(t, "main.go", "package main\n")
	start := time.Now()
	ds, err := s.Check(shortCtx(t, 2*time.Second), path)
	if err != nil {
		t.Fatal(err)
	}
	if len(ds) != 0 {
		t.Errorf("clean file got %+v", ds)
	}
	// An empty batch is an answer. The check should cost the settle window,
	// not the deadline.
	if took := time.Since(start); took > settleWindow+time.Second {
		t.Errorf("clean check took %s", took)
	}
}

func TestSessionCollectsNeighboursAndKeepsTheLatestBatch(t *testing.T) {
	var neighbour string
	serve := func(f *fakeServer, msg incoming) {
		switch msg.Method {
		case "textDocument/didOpen":
			uri, _ := docText(msg)
			// A clearing batch, then the real one, then a neighbour — the
			// order gopls produces.
			f.publish(uri, nil)
			f.publish(uri, []map[string]any{lspDiag(0, 0, 1, "real")})
			f.publish(neighbour, []map[string]any{lspDiag(2, 0, 1, "neighbour broke")})
		default:
			wellBehaved(f, msg)
		}
	}
	s, _ := startFake(t, serve, 0, nil)
	path := writeTemp(t, "main.go", "x\n")
	neighbour = fileURI(filepath.Join(filepath.Dir(path), "other.go"))

	ds, err := s.Check(shortCtx(t, 2*time.Second), path)
	if err != nil {
		t.Fatal(err)
	}
	if len(ds) != 2 {
		t.Fatalf("got %+v", ds)
	}
	sortForReport(ds, path)
	if ds[0].Message != "real" || ds[1].Message != "neighbour broke" || filepath.Base(ds[1].File) != "other.go" {
		t.Errorf("got %+v", ds)
	}
}

func TestSessionServerThatNeverAnswersDoesNotHang(t *testing.T) {
	silent := func(*fakeServer, incoming) {}
	s, _ := startFake(t, silent, 0, nil)
	path := writeTemp(t, "main.go", "x\n")

	start := time.Now()
	_, err := s.Check(shortCtx(t, 150*time.Millisecond), path)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("err = %v, want deadline exceeded", err)
	}
	if took := time.Since(start); took > time.Second {
		t.Errorf("check took %s against a silent server", took)
	}
	// Close must also not wait on it.
	start = time.Now()
	s.Close()
	if took := time.Since(start); took > 3*shutdownGrace {
		t.Errorf("Close took %s against a silent server", took)
	}
	if !s.Dead() {
		t.Error("session not dead after Close")
	}
}

func TestSessionServerThatInitializesButNeverPublishes(t *testing.T) {
	serve := func(f *fakeServer, msg incoming) {
		if msg.Method == "initialize" {
			f.respond(msg.ID, map[string]any{"capabilities": map[string]any{}})
		}
	}
	s, _ := startFake(t, serve, 0, nil)
	path := writeTemp(t, "main.go", "x\n")
	_, err := s.Check(shortCtx(t, 150*time.Millisecond), path)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("err = %v, want deadline exceeded", err)
	}
	// The server is fine, just slow; it must not have been declared dead.
	if s.Dead() {
		t.Error("a slow server was killed")
	}
}

func TestSessionServerThatExitsImmediately(t *testing.T) {
	exited := make(chan *session, 1)
	serve := func(f *fakeServer, _ incoming) { f.close() }
	s, _ := startFake(t, serve, 0, func(s *session) { exited <- s })
	path := writeTemp(t, "main.go", "x\n")

	_, err := s.Check(shortCtx(t, time.Second), path)
	if err == nil {
		t.Fatal("check against a dead server succeeded")
	}
	if errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("waited for the deadline rather than noticing the exit: %v", err)
	}
	select {
	case got := <-exited:
		if got != s {
			t.Error("onExit called with a different session")
		}
	case <-time.After(time.Second):
		t.Fatal("onExit never called")
	}
	if !s.failedToStart() {
		t.Error("a server that exited before initialize should count as failed to start")
	}
}

func TestSessionServerThatDiesMidCheck(t *testing.T) {
	serve := func(f *fakeServer, msg incoming) {
		switch msg.Method {
		case "initialize":
			f.respond(msg.ID, map[string]any{"capabilities": map[string]any{}})
		case "textDocument/didOpen":
			f.close()
		}
	}
	s, _ := startFake(t, serve, 0, nil)
	path := writeTemp(t, "main.go", "x\n")
	_, err := s.Check(shortCtx(t, time.Second), path)
	if !errors.Is(err, errServerDead) {
		t.Errorf("err = %v, want errServerDead", err)
	}
	if s.failedToStart() {
		t.Error("a server that died after initializing is a crash, not a start failure")
	}
}

func TestSessionGarbageOnTheWireEndsTheSession(t *testing.T) {
	serve := func(f *fakeServer, msg incoming) {
		if msg.Method == "initialize" {
			f.raw("this is not a language server\n")
		}
	}
	s, _ := startFake(t, serve, 0, nil)
	_, err := s.Check(shortCtx(t, time.Second), writeTemp(t, "main.go", "x\n"))
	if err == nil || errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("err = %v, want a prompt failure", err)
	}
	if !s.Dead() {
		t.Error("session survived garbage")
	}
}

func TestSessionFloodIsCapped(t *testing.T) {
	serve := func(f *fakeServer, msg incoming) {
		switch msg.Method {
		case "initialize":
			f.respond(msg.ID, map[string]any{"capabilities": map[string]any{}})
		case "textDocument/didOpen":
			uri, _ := docText(msg)
			dir := filepath.Dir(uriPath(uri))
			// Far more files than anyone can use, and far more problems
			// per file.
			for i := 0; i < 3*maxDiagnosticFiles; i++ {
				f.publish(fileURI(filepath.Join(dir, fmt.Sprintf("n%d.go", i))), []map[string]any{lspDiag(0, 0, 1, "x")})
			}
			many := make([]map[string]any, 5*maxDiagnosticsPerURI)
			for i := range many {
				many[i] = lspDiag(i, 0, 1, "again")
			}
			f.publish(uri, many)
		}
	}
	s, _ := startFake(t, serve, 0, nil)
	path := writeTemp(t, "main.go", "x\n")
	ds, err := s.Check(shortCtx(t, 3*time.Second), path)
	if err != nil {
		t.Fatal(err)
	}
	files := map[string]int{}
	for _, d := range ds {
		files[d.File]++
	}
	// The cap is on neighbours; the checked file is kept on top of it.
	if len(files) > maxDiagnosticFiles+1 {
		t.Errorf("%d files collected, cap is %d", len(files), maxDiagnosticFiles)
	}
	if files[path] != maxDiagnosticsPerURI {
		t.Errorf("%d diagnostics for the target, cap is %d", files[path], maxDiagnosticsPerURI)
	}
	if files[path] == 0 {
		t.Error("the target file was pushed out by the flood")
	}
}

func TestSessionAnswersServerRequests(t *testing.T) {
	serve := func(f *fakeServer, msg incoming) {
		switch msg.Method {
		case "initialize":
			f.respond(msg.ID, map[string]any{"capabilities": map[string]any{}})
		case "initialized":
			f.request(7, "workspace/configuration", map[string]any{"items": []any{map[string]any{}, map[string]any{}}})
			f.request(8, "client/registerCapability", map[string]any{"registrations": []any{}})
			f.request(9, "something/unknown", nil)
		case "textDocument/didOpen":
			uri, _ := docText(msg)
			f.publish(uri, nil)
		}
	}
	s, f := startFake(t, serve, 0, nil)
	if _, err := s.Check(shortCtx(t, 2*time.Second), writeTemp(t, "main.go", "x\n")); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(time.Second)
	for {
		f.mu.Lock()
		n := len(f.replies)
		f.mu.Unlock()
		if n >= 3 || time.Now().After(deadline) {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	byID := map[string]incoming{}
	for _, r := range f.replies {
		byID[string(r.ID)] = r
	}
	if r, ok := byID["7"]; !ok || string(r.Result) != "[null,null]" {
		t.Errorf("workspace/configuration reply = %+v", r)
	}
	if r, ok := byID["8"]; !ok || r.Error != nil {
		t.Errorf("client/registerCapability reply = %+v", r)
	}
	if r, ok := byID["9"]; !ok || r.Error == nil || r.Error.Code != -32601 {
		t.Errorf("unknown method reply = %+v", r)
	}
}

func TestSessionIdleTimeoutShutsDownPolitely(t *testing.T) {
	exited := make(chan *session, 1)
	s, f := startFake(t, wellBehaved, 100*time.Millisecond, func(s *session) { exited <- s })
	if _, err := s.Check(shortCtx(t, 2*time.Second), writeTemp(t, "main.go", "x\n")); err != nil {
		t.Fatal(err)
	}
	select {
	case <-exited:
	case <-time.After(3 * time.Second):
		t.Fatal("idle server was never shut down")
	}
	methods := strings.Join(f.methods(), ",")
	if !strings.Contains(methods, "shutdown") || !strings.HasSuffix(methods, "exit") {
		t.Errorf("idle shutdown was not the polite sequence: %s", methods)
	}
	if _, err := s.Check(shortCtx(t, time.Second), writeTemp(t, "main.go", "x\n")); !errors.Is(err, errServerDead) {
		t.Errorf("check after idle shutdown: %v, want errServerDead", err)
	}
}

func TestSessionConcurrentChecksDoNotInterleave(t *testing.T) {
	s, _ := startFake(t, wellBehaved, 0, nil)
	paths := []string{
		writeTemp(t, "a.go", "BUG\n"),
		writeTemp(t, "b.go", "ok\nBUG\nBUG\n"),
		writeTemp(t, "c.go", "ok\n"),
	}
	want := []int{1, 2, 0}
	var wg sync.WaitGroup
	errs := make([]error, len(paths))
	counts := make([]int, len(paths))
	for i, p := range paths {
		wg.Add(1)
		go func(i int, p string) {
			defer wg.Done()
			ds, err := s.Check(shortCtx(t, 5*time.Second), p)
			errs[i] = err
			for _, d := range ds {
				if d.File == p {
					counts[i]++
				}
			}
		}(i, p)
	}
	wg.Wait()
	for i := range paths {
		if errs[i] != nil {
			t.Errorf("%s: %v", paths[i], errs[i])
		}
		if counts[i] != want[i] {
			t.Errorf("%s: %d diagnostics, want %d", filepath.Base(paths[i]), counts[i], want[i])
		}
	}
}

func TestFileURIRoundTrip(t *testing.T) {
	for _, p := range []string{"/tmp/a.go", "/tmp/with space/b.go", "/tmp/ünïcode/c.go", "/tmp/q?x#y.go"} {
		uri := fileURI(p)
		if !strings.HasPrefix(uri, "file:///") {
			t.Errorf("fileURI(%q) = %q", p, uri)
		}
		if got := uriPath(uri); got != p {
			t.Errorf("uriPath(fileURI(%q)) = %q", p, got)
		}
	}
	if uriPath("untitled:Untitled-1") != "" {
		t.Error("non-file URI should map to no path")
	}
	if !sameURI("file:///tmp/with%20space/b.go", fileURI("/tmp/with space/b.go")) {
		t.Error("sameURI should compare by path")
	}
}