PgTidy/pkg/lsp/server.go

// Package lsp implements a Language Server Protocol server for PgTidy.
//
// The server communicates over stdio using JSON-RPC 2.0 with Content-Length
// framing. It provides:
//   - textDocument/formatting — full-document formatting via pkg/format
//   - textDocument/publishDiagnostics — lint findings via pkg/lint, sent on
//     every didOpen/didChange notification
package lsp

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"strings"

	"git.warky.dev/wdevs/pgtidy/pkg/config"
	"git.warky.dev/wdevs/pgtidy/pkg/diagnostics"
	"git.warky.dev/wdevs/pgtidy/pkg/format"
	"git.warky.dev/wdevs/pgtidy/pkg/lint"
	"git.warky.dev/wdevs/pgtidy/pkg/parser"
)

// Serve runs the LSP server, reading JSON-RPC messages from r and writing to w.
// startDir is used for .pgtidy.yaml discovery. Returns when the client sends
// "exit" or r reaches EOF.
func Serve(ctx context.Context, r io.Reader, w io.Writer, startDir string) error {
	cfg, _ := config.Load(startDir)
	srv := &server{
		docs:  make(map[string]string),
		fixes: make(map[string][]diagnostics.Diagnostic),
		cfg:   cfg,
		w:     w,
	}
	return srv.loop(ctx, r)
}

type server struct {
	docs  map[string]string                   // URI → current text
	fixes map[string][]diagnostics.Diagnostic // URI → diagnostics that have fixes
	cfg   config.Style
	w     io.Writer
}

func (s *server) loop(ctx context.Context, r io.Reader) error {
	br := bufio.NewReader(r)
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		raw, err := readMsg(br)
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
		if s.handle(raw) {
			return nil
		}
	}
}

// handle dispatches one JSON-RPC message. Returns true when the server should exit.
func (s *server) handle(raw []byte) bool {
	var req rpcMsg
	if err := json.Unmarshal(raw, &req); err != nil {
		return false
	}

	switch req.Method {
	case "initialize":
		s.reply(req.ID, initResult{
			Capabilities: serverCaps{
				TextDocumentSync:                1, // full sync
				DocumentFormattingProvider:      true,
				DocumentRangeFormattingProvider: true,
				CodeActionProvider:              true,
			},
		})
	case "initialized": // no-op notification
	case "shutdown":
		s.reply(req.ID, nil)
	case "exit":
		return true
	case "textDocument/didOpen":
		var p didOpenParams
		_ = json.Unmarshal(req.Params, &p)
		s.docs[p.TextDocument.URI] = p.TextDocument.Text
		s.pushDiagnostics(p.TextDocument.URI, p.TextDocument.Text)
	case "textDocument/didChange":
		var p didChangeParams
		_ = json.Unmarshal(req.Params, &p)
		if len(p.ContentChanges) > 0 {
			text := p.ContentChanges[len(p.ContentChanges)-1].Text
			s.docs[p.TextDocument.URI] = text
			s.pushDiagnostics(p.TextDocument.URI, text)
		}
	case "textDocument/didClose":
		var p struct {
			TextDocument textDocID `json:"textDocument"`
		}
		_ = json.Unmarshal(req.Params, &p)
		delete(s.docs, p.TextDocument.URI)
		delete(s.fixes, p.TextDocument.URI)
		s.notify("textDocument/publishDiagnostics", publishDiagnosticsParams{
			URI:         p.TextDocument.URI,
			Diagnostics: []lspDiagnostic{},
		})
	case "textDocument/formatting":
		var p formattingParams
		_ = json.Unmarshal(req.Params, &p)
		text, ok := s.docs[p.TextDocument.URI]
		if !ok {
			s.reply(req.ID, []textEdit{})
			return false
		}
		formatted := format.File(parser.Parse(text), s.cfg)
		if formatted == text {
			s.reply(req.ID, []textEdit{})
			return false
		}
		if err := format.VerifySafe(text, formatted, s.cfg); err != nil {
			s.reply(req.ID, []textEdit{})
			return false
		}
		s.reply(req.ID, []textEdit{fullReplace(text, formatted)})
	case "textDocument/rangeFormatting":
		var p rangeFormattingParams
		_ = json.Unmarshal(req.Params, &p)
		text, ok := s.docs[p.TextDocument.URI]
		if !ok {
			s.reply(req.ID, []textEdit{})
			return false
		}
		s.reply(req.ID, s.rangeFormat(text, p.Range))
	case "textDocument/codeAction":
		var p codeActionParams
		_ = json.Unmarshal(req.Params, &p)
		s.handleCodeAction(req.ID, p)
	case "$/cancelRequest": // ignore
	default:
		if req.ID != nil {
			s.replyErr(req.ID, -32601, "method not found: "+req.Method)
		}
	}
	return false
}

func (s *server) pushDiagnostics(uri, text string) {
	eng := lint.New()
	diags, _ := eng.Check(text, uri)

	var fixable []diagnostics.Diagnostic
	out := make([]lspDiagnostic, 0, len(diags))
	for _, d := range diags {
		line := uint32(d.Line)
		if line > 0 {
			line--
		}
		col := uint32(d.Col)
		if col > 0 {
			col--
		}
		lspD := lspDiagnostic{
			Range:    lspRange{Start: position{line, col}, End: position{line, col + 1}},
			Severity: severityCode(d.Severity),
			Code:     d.RuleID,
			Source:   "pgtidy",
			Message:  d.Message,
		}
		out = append(out, lspD)
		if d.Fix != nil {
			fixable = append(fixable, d)
		}
	}
	s.fixes[uri] = fixable
	s.notify("textDocument/publishDiagnostics", publishDiagnosticsParams{
		URI:         uri,
		Diagnostics: out,
	})
}

func (s *server) handleCodeAction(id json.RawMessage, p codeActionParams) {
	uri := p.TextDocument.URI
	text, ok := s.docs[uri]
	if !ok {
		s.reply(id, []codeAction{})
		return
	}
	var actions []codeAction
	for _, d := range s.fixes[uri] {
		if d.Fix == nil {
			continue
		}
		start := offsetToPosition(text, d.Fix.Offset)
		end := offsetToPosition(text, d.Fix.End)
		// Only include if the fix range overlaps the requested range.
		if !rangesOverlap(start, end, p.Range.Start, p.Range.End) {
			continue
		}
		actions = append(actions, codeAction{
			Title: d.Fix.Title,
			Kind:  "quickfix",
			Edit: &workspaceEdit{
				Changes: map[string][]textEdit{
					uri: {{
						Range:   lspRange{Start: start, End: end},
						NewText: d.Fix.New,
					}},
				},
			},
		})
	}
	s.reply(id, actions)
}

// rangeFormat formats the region of text that overlaps the given LSP range.
// It formats the complete document, then returns an edit covering only the
// minimal set of changed lines that intersects the selection. Returns nil when
// the document is already formatted or no change falls within the selection.
func (s *server) rangeFormat(text string, r lspRange) []textEdit {
	formatted := format.File(parser.Parse(text), s.cfg)
	if formatted == text {
		return nil
	}
	if err := format.VerifySafe(text, formatted, s.cfg); err != nil {
		return nil
	}

	// Split both versions into lines, keeping the trailing newline attached to
	// each element so that joining them reconstructs the original string.
	origLines := strings.SplitAfter(text, "\n")
	fmtLines := strings.SplitAfter(formatted, "\n")

	// Longest common prefix of unchanged lines.
	pfx := 0
	for pfx < len(origLines) && pfx < len(fmtLines) && origLines[pfx] == fmtLines[pfx] {
		pfx++
	}
	// Longest common suffix (must not overlap the prefix).
	sfx := 0
	for sfx < len(origLines)-pfx && sfx < len(fmtLines)-pfx &&
		origLines[len(origLines)-1-sfx] == fmtLines[len(fmtLines)-1-sfx] {
		sfx++
	}

	// The changed region in the original spans lines [pfx, changedEnd).
	changedStart := pfx
	changedEnd := len(origLines) - sfx

	// Only return an edit when the changed region overlaps the selection.
	selStart := int(r.Start.Line)
	selEnd := int(r.End.Line)
	if changedEnd <= selStart || changedStart > selEnd {
		return nil
	}

	// Compute byte offsets for the changed region boundary.
	startByte := 0
	for i := 0; i < changedStart; i++ {
		startByte += len(origLines[i])
	}
	endByte := startByte
	for i := changedStart; i < changedEnd; i++ {
		endByte += len(origLines[i])
	}

	newText := strings.Join(fmtLines[pfx:len(fmtLines)-sfx], "")
	return []textEdit{{
		Range:   lspRange{Start: offsetToPosition(text, startByte), End: offsetToPosition(text, endByte)},
		NewText: newText,
	}}
}

// offsetToPosition converts a byte offset in text to an LSP position.
func offsetToPosition(text string, byteOff int) position {
	if byteOff > len(text) {
		byteOff = len(text)
	}
	line := uint32(0)
	lastNL := -1
	for i := 0; i < byteOff; i++ {
		if text[i] == '\n' {
			line++
			lastNL = i
		}
	}
	return position{line, uint32(byteOff - lastNL - 1)}
}

// rangesOverlap returns true when [s1,e1) and [s2,e2) share any point.
// For insertions (s==e) we check if the point falls within the other range.
func rangesOverlap(s1, e1, s2, e2 position) bool {
	cmp := func(a, b position) int {
		if a.Line != b.Line {
			if a.Line < b.Line {
				return -1
			}
			return 1
		}
		if a.Character < b.Character {
			return -1
		}
		if a.Character > b.Character {
			return 1
		}
		return 0
	}
	return cmp(s1, e2) <= 0 && cmp(s2, e1) <= 0
}

func severityCode(sev diagnostics.Severity) int {
	switch sev {
	case diagnostics.SeverityError:
		return 1
	case diagnostics.SeverityWarning:
		return 2
	case diagnostics.SeverityHint:
		return 4
	}
	return 3 // info
}

// fullReplace builds a TextEdit that replaces the entire document.
func fullReplace(orig, formatted string) textEdit {
	lines := strings.Split(orig, "\n")
	lastLine := uint32(len(lines) - 1)
	lastChar := uint32(len(lines[lastLine]))
	return textEdit{
		Range:   lspRange{Start: position{0, 0}, End: position{lastLine, lastChar}},
		NewText: formatted,
	}
}

// --- JSON-RPC 2.0 transport ---

func readMsg(r *bufio.Reader) ([]byte, error) {
	length := 0
	for {
		line, err := r.ReadString('\n')
		if err != nil {
			return nil, err
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" {
			break
		}
		if after, ok := strings.CutPrefix(line, "Content-Length: "); ok {
			length, _ = strconv.Atoi(after)
		}
	}
	if length == 0 {
		return nil, fmt.Errorf("lsp: missing Content-Length")
	}
	buf := make([]byte, length)
	_, err := io.ReadFull(r, buf)
	return buf, err
}

func (s *server) send(v interface{}) {
	data, err := json.Marshal(v)
	if err != nil {
		return
	}
	_, _ = fmt.Fprintf(s.w, "Content-Length: %d\r\n\r\n", len(data))
	s.w.Write(data) //nolint:errcheck
}

func (s *server) reply(id json.RawMessage, result interface{}) {
	s.send(rpcResponse{JSONRPC: "2.0", ID: id, Result: result})
}

func (s *server) replyErr(id json.RawMessage, code int, msg string) {
	s.send(rpcResponse{JSONRPC: "2.0", ID: id, Error: &rpcError{Code: code, Message: msg}})
}

func (s *server) notify(method string, params interface{}) {
	s.send(rpcNotification{JSONRPC: "2.0", Method: method, Params: params})
}

// --- JSON-RPC 2.0 wire types ---

type rpcMsg struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

type rpcResponse struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"`
	Result  interface{}     `json:"result"`
	Error   *rpcError       `json:"error,omitempty"`
}

type rpcNotification struct {
	JSONRPC string      `json:"jsonrpc"`
	Method  string      `json:"method"`
	Params  interface{} `json:"params"`
}

type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// --- LSP protocol types (minimal subset) ---

type position struct {
	Line      uint32 `json:"line"`
	Character uint32 `json:"character"`
}

type lspRange struct {
	Start position `json:"start"`
	End   position `json:"end"`
}

type textEdit struct {
	Range   lspRange `json:"range"`
	NewText string   `json:"newText"`
}

type initResult struct {
	Capabilities serverCaps `json:"capabilities"`
}

type serverCaps struct {
	TextDocumentSync                int  `json:"textDocumentSync"`
	DocumentFormattingProvider      bool `json:"documentFormattingProvider"`
	DocumentRangeFormattingProvider bool `json:"documentRangeFormattingProvider"`
	CodeActionProvider              bool `json:"codeActionProvider"`
}

type textDocItem struct {
	URI        string `json:"uri"`
	LanguageID string `json:"languageId"`
	Version    int    `json:"version"`
	Text       string `json:"text"`
}

type textDocID struct {
	URI string `json:"uri"`
}

type contentChange struct {
	Text string `json:"text"`
}

type didOpenParams struct {
	TextDocument textDocItem `json:"textDocument"`
}

type didChangeParams struct {
	TextDocument   textDocID       `json:"textDocument"`
	ContentChanges []contentChange `json:"contentChanges"`
}

type formattingParams struct {
	TextDocument textDocID `json:"textDocument"`
	Options      struct {
		TabSize      int  `json:"tabSize"`
		InsertSpaces bool `json:"insertSpaces"`
	} `json:"options"`
}

type rangeFormattingParams struct {
	TextDocument textDocID `json:"textDocument"`
	Range        lspRange  `json:"range"`
	Options      struct {
		TabSize      int  `json:"tabSize"`
		InsertSpaces bool `json:"insertSpaces"`
	} `json:"options"`
}

type lspDiagnostic struct {
	Range    lspRange `json:"range"`
	Severity int      `json:"severity"`
	Code     string   `json:"code,omitempty"`
	Source   string   `json:"source,omitempty"`
	Message  string   `json:"message"`
}

type publishDiagnosticsParams struct {
	URI         string          `json:"uri"`
	Diagnostics []lspDiagnostic `json:"diagnostics"`
}

type codeActionParams struct {
	TextDocument textDocID `json:"textDocument"`
	Range        lspRange  `json:"range"`
	Context      struct {
		Diagnostics []lspDiagnostic `json:"diagnostics"`
	} `json:"context"`
}

type workspaceEdit struct {
	Changes map[string][]textEdit `json:"changes"`
}

type codeAction struct {
	Title string         `json:"title"`
	Kind  string         `json:"kind,omitempty"`
	Edit  *workspaceEdit `json:"edit,omitempty"`
}