MonocoEditor-With-Lsp-Backend/backend/internal/lsp/client.go

package lsp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"monica-go-completion-backend/internal/completion"
)

var errClientClosed = errors.New("lsp client closed")

type Client struct {
	cmd   *exec.Cmd
	stdin io.WriteCloser

	writeMu sync.Mutex

	nextID atomic.Int64

	workspaceDir string

	pendingMu sync.Mutex
	pending   map[string]chan rpcResponse

	exitCh    chan error
	closeOnce sync.Once
	closed    atomic.Bool
}

type rpcResponse struct {
	result json.RawMessage
	err    error
}

type rpcError struct {
	Code    int             `json:"code"`
	Message string          `json:"message"`
	Data    json.RawMessage `json:"data,omitempty"`
}

type incomingEnvelope struct {
	ID     *json.RawMessage `json:"id,omitempty"`
	Result json.RawMessage  `json:"result,omitempty"`
	Error  *rpcError        `json:"error,omitempty"`
}

type lspCompletionItem struct {
	Label         string          `json:"label"`
	Kind          int             `json:"kind,omitempty"`
	Detail        string          `json:"detail,omitempty"`
	InsertText    string          `json:"insertText,omitempty"`
	SortText      string          `json:"sortText,omitempty"`
	FilterText    string          `json:"filterText,omitempty"`
	Documentation json.RawMessage `json:"documentation,omitempty"`
}

type lspCompletionList struct {
	IsIncomplete bool                `json:"isIncomplete"`
	Items        []lspCompletionItem `json:"items"`
}

func NewClient(parent context.Context, goplsPath, rootPath string) (*Client, error) {
	if goplsPath == "" {
		goplsPath = "gopls"
	}
	if rootPath == "" {
		cwd, err := os.Getwd()
		if err != nil {
			return nil, fmt.Errorf("get working directory: %w", err)
		}
		rootPath = cwd
	}

	cmd := exec.Command(goplsPath)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("create stdin pipe: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("create stdout pipe: %w", err)
	}
	cmd.Stderr = os.Stderr

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start gopls: %w", err)
	}

	client := &Client{
		cmd:          cmd,
		stdin:        stdin,
		workspaceDir: rootPath,
		pending:      make(map[string]chan rpcResponse),
		exitCh:       make(chan error, 1),
	}

	go func() {
		client.exitCh <- cmd.Wait()
	}()
	go client.readLoop(stdout)

	initCtx, cancel := context.WithTimeout(parent, 10*time.Second)
	defer cancel()

	if err := client.initialize(initCtx, rootPath); err != nil {
		_ = client.Close()
		return nil, err
	}

	return client, nil
}

func (c *Client) initialize(ctx context.Context, rootPath string) error {
	rootURI, err := pathToURI(rootPath)
	if err != nil {
		return fmt.Errorf("build root uri: %w", err)
	}

	params := map[string]any{
		"processId": os.Getpid(),
		"rootUri":   rootURI,
		"capabilities": map[string]any{
			"textDocument": map[string]any{
				"completion": map[string]any{
					"completionItem": map[string]any{
						"snippetSupport": true,
					},
				},
			},
		},
		"workspaceFolders": []map[string]string{
			{
				"uri":  rootURI,
				"name": filepath.Base(rootPath),
			},
		},
		"clientInfo": map[string]string{
			"name":    "monica-go-completion-backend",
			"version": "0.1.0",
		},
	}

	var initResult json.RawMessage
	if err := c.request(ctx, "initialize", params, &initResult); err != nil {
		return fmt.Errorf("initialize request failed: %w", err)
	}
	if err := c.notify("initialized", map[string]any{}); err != nil {
		return fmt.Errorf("initialized notification failed: %w", err)
	}
	return nil
}

func (c *Client) DidOpen(ctx context.Context, uri, text string, version int) error {
	normalizedURI, err := c.normalizeURI(uri)
	if err != nil {
		return err
	}

	params := map[string]any{
		"textDocument": map[string]any{
			"uri":        normalizedURI,
			"languageId": "go",
			"version":    version,
			"text":       text,
		},
	}
	return c.notifyWithContext(ctx, "textDocument/didOpen", params)
}

func (c *Client) DidChange(ctx context.Context, uri, text string, version int) error {
	normalizedURI, err := c.normalizeURI(uri)
	if err != nil {
		return err
	}

	params := map[string]any{
		"textDocument": map[string]any{
			"uri":     normalizedURI,
			"version": version,
		},
		"contentChanges": []map[string]string{
			{
				"text": text,
			},
		},
	}
	return c.notifyWithContext(ctx, "textDocument/didChange", params)
}

func (c *Client) Completion(ctx context.Context, uri string, line, character int) (completion.Response, error) {
	normalizedURI, err := c.normalizeURI(uri)
	if err != nil {
		return completion.Response{}, err
	}

	params := map[string]any{
		"textDocument": map[string]string{
			"uri": normalizedURI,
		},
		"position": map[string]int{
			"line":      line,
			"character": character,
		},
	}

	var raw json.RawMessage
	if err := c.request(ctx, "textDocument/completion", params, &raw); err != nil {
		return completion.Response{}, err
	}

	body := bytes.TrimSpace(raw)
	if len(body) == 0 || bytes.Equal(body, []byte("null")) {
		return completion.Response{}, nil
	}

	if body[0] == '[' {
		var items []lspCompletionItem
		if err := json.Unmarshal(body, &items); err != nil {
			return completion.Response{}, fmt.Errorf("decode completion items: %w", err)
		}
		return completion.Response{Items: mapCompletionItems(items)}, nil
	}

	var list lspCompletionList
	if err := json.Unmarshal(body, &list); err != nil {
		return completion.Response{}, fmt.Errorf("decode completion list: %w", err)
	}
	return completion.Response{
		Items:        mapCompletionItems(list.Items),
		IsIncomplete: list.IsIncomplete,
	}, nil
}

var windowsDrivePattern = regexp.MustCompile(`^[A-Za-z]:`)

func (c *Client) normalizeURI(rawURI string) (string, error) {
	if rawURI == "" {
		return "", errors.New("empty uri")
	}

	u, err := url.Parse(rawURI)
	if err != nil {
		return "", fmt.Errorf("invalid uri: %w", err)
	}
	if u.Scheme != "file" {
		return rawURI, nil
	}

	path := u.Path
	if path == "" {
		return rawURI, nil
	}

	var localPath string
	trimmed := strings.TrimPrefix(path, "/")
	if windowsDrivePattern.MatchString(trimmed) {
		localPath = filepath.FromSlash(trimmed)
	} else {
		rel := strings.TrimLeft(trimmed, "/\\")
		localPath = filepath.Join(c.workspaceDir, filepath.FromSlash(rel))
	}

	return pathToURI(localPath)
}

func (c *Client) Close() error {
	c.closeOnce.Do(func() {
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()

		_ = c.request(shutdownCtx, "shutdown", map[string]any{}, nil)
		_ = c.notify("exit", map[string]any{})
		_ = c.stdin.Close()

		select {
		case <-c.exitCh:
		case <-time.After(2 * time.Second):
			if c.cmd.Process != nil {
				_ = c.cmd.Process.Kill()
			}
			<-c.exitCh
		}

		c.closed.Store(true)
		c.failPending(errClientClosed)
	})
	return nil
}

func (c *Client) request(ctx context.Context, method string, params any, out any) error {
	if c.closed.Load() {
		return errClientClosed
	}

	id := c.nextID.Add(1)
	key := strconv.FormatInt(id, 10)
	wait := make(chan rpcResponse, 1)

	c.pendingMu.Lock()
	c.pending[key] = wait
	c.pendingMu.Unlock()

	msg := map[string]any{
		"jsonrpc": "2.0",
		"id":      id,
		"method":  method,
		"params":  params,
	}
	if err := c.writeMessage(msg); err != nil {
		c.removePending(key)
		return err
	}

	select {
	case <-ctx.Done():
		c.removePending(key)
		return ctx.Err()
	case resp := <-wait:
		if resp.err != nil {
			return resp.err
		}
		if out != nil && len(resp.result) > 0 && !bytes.Equal(bytes.TrimSpace(resp.result), []byte("null")) {
			if err := json.Unmarshal(resp.result, out); err != nil {
				return fmt.Errorf("decode response for %s: %w", method, err)
			}
		}
		return nil
	}
}

func (c *Client) notify(method string, params any) error {
	return c.notifyWithContext(context.Background(), method, params)
}

func (c *Client) notifyWithContext(ctx context.Context, method string, params any) error {
	if c.closed.Load() {
		return errClientClosed
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

	msg := map[string]any{
		"jsonrpc": "2.0",
		"method":  method,
		"params":  params,
	}
	return c.writeMessage(msg)
}

func (c *Client) writeMessage(msg any) error {
	payload, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("marshal json-rpc message: %w", err)
	}

	header := fmt.Sprintf("Content-Length: %d\r\n\r\n", len(payload))

	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	if _, err := io.WriteString(c.stdin, header); err != nil {
		return fmt.Errorf("write lsp header: %w", err)
	}
	if _, err := c.stdin.Write(payload); err != nil {
		return fmt.Errorf("write lsp payload: %w", err)
	}
	return nil
}

func (c *Client) readLoop(stdout io.Reader) {
	reader := bufio.NewReader(stdout)
	for {
		body, err := readMessage(reader)
		if err != nil {
			c.failPending(fmt.Errorf("read lsp message: %w", err))
			return
		}
		c.handleIncoming(body)
	}
}

func readMessage(reader *bufio.Reader) ([]byte, error) {
	contentLength := 0
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			return nil, err
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" {
			break
		}

		parts := strings.SplitN(line, ":", 2)
		if len(parts) != 2 {
			continue
		}
		if strings.EqualFold(strings.TrimSpace(parts[0]), "Content-Length") {
			n, err := strconv.Atoi(strings.TrimSpace(parts[1]))
			if err != nil {
				return nil, fmt.Errorf("invalid content-length %q: %w", parts[1], err)
			}
			contentLength = n
		}
	}

	if contentLength <= 0 {
		return nil, errors.New("missing content-length")
	}

	body := make([]byte, contentLength)
	if _, err := io.ReadFull(reader, body); err != nil {
		return nil, err
	}
	return body, nil
}

func (c *Client) handleIncoming(body []byte) {
	var envelope incomingEnvelope
	if err := json.Unmarshal(body, &envelope); err != nil {
		return
	}

	if envelope.ID == nil {
		return
	}

	key := normalizeID(*envelope.ID)
	if key == "" {
		return
	}

	c.pendingMu.Lock()
	wait := c.pending[key]
	delete(c.pending, key)
	c.pendingMu.Unlock()

	if wait == nil {
		return
	}

	if envelope.Error != nil {
		wait <- rpcResponse{
			err: fmt.Errorf("lsp error code=%d message=%s", envelope.Error.Code, envelope.Error.Message),
		}
		return
	}

	wait <- rpcResponse{result: envelope.Result}
}

func normalizeID(raw json.RawMessage) string {
	raw = bytes.TrimSpace(raw)
	if len(raw) == 0 {
		return ""
	}

	if raw[0] == '"' {
		var s string
		if err := json.Unmarshal(raw, &s); err == nil {
			return s
		}
	}
	return string(raw)
}

func (c *Client) removePending(key string) {
	c.pendingMu.Lock()
	delete(c.pending, key)
	c.pendingMu.Unlock()
}

func (c *Client) failPending(err error) {
	c.pendingMu.Lock()
	defer c.pendingMu.Unlock()

	for key, wait := range c.pending {
		delete(c.pending, key)
		wait <- rpcResponse{err: err}
	}
}

func mapCompletionItems(items []lspCompletionItem) []completion.Item {
	out := make([]completion.Item, 0, len(items))
	for _, it := range items {
		out = append(out, completion.Item{
			Label:         it.Label,
			Kind:          it.Kind,
			Detail:        it.Detail,
			Documentation: decodeDocumentation(it.Documentation),
			InsertText:    it.InsertText,
			SortText:      it.SortText,
			FilterText:    it.FilterText,
		})
	}
	return out
}

func decodeDocumentation(raw json.RawMessage) string {
	raw = bytes.TrimSpace(raw)
	if len(raw) == 0 || bytes.Equal(raw, []byte("null")) {
		return ""
	}

	if raw[0] == '"' {
		var s string
		if err := json.Unmarshal(raw, &s); err == nil {
			return s
		}
		return ""
	}

	var markup struct {
		Value string `json:"value"`
	}
	if err := json.Unmarshal(raw, &markup); err == nil && markup.Value != "" {
		return markup.Value
	}

	return ""
}

func pathToURI(path string) (string, error) {
	abs, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}
	slashed := filepath.ToSlash(abs)
	if !strings.HasPrefix(slashed, "/") {
		slashed = "/" + slashed
	}
	u := url.URL{Scheme: "file", Path: slashed}
	return u.String(), nil
}