add initial lsp support

1 files changed, 429 insertions, 0 deletions

diff --git a/internal/lsp/client.go b/internal/lsp/client.go
new file mode 100644
index 000000000..ddf28ca2a
--- /dev/null
+++ b/internal/lsp/client.go
@@ -0,0 +1,429 @@
+package lsp
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"os/exec"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/kujtimiihoxha/termai/internal/lsp/protocol"
+)
+
+type Client struct {
+	Cmd    *exec.Cmd
+	stdin  io.WriteCloser
+	stdout *bufio.Reader
+	stderr io.ReadCloser
+
+	// Request ID counter
+	nextID atomic.Int32
+
+	// Response handlers
+	handlers   map[int32]chan *Message
+	handlersMu sync.RWMutex
+
+	// Server request handlers
+	serverRequestHandlers map[string]ServerRequestHandler
+	serverHandlersMu      sync.RWMutex
+
+	// Notification handlers
+	notificationHandlers map[string]NotificationHandler
+	notificationMu       sync.RWMutex
+
+	// Diagnostic cache
+	diagnostics   map[protocol.DocumentUri][]protocol.Diagnostic
+	diagnosticsMu sync.RWMutex
+
+	// Files are currently opened by the LSP
+	openFiles   map[string]*OpenFileInfo
+	openFilesMu sync.RWMutex
+}
+
+func NewClient(command string, args ...string) (*Client, error) {
+	cmd := exec.Command(command, args...)
+	// Copy env
+	cmd.Env = os.Environ()
+
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create stdin pipe: %w", err)
+	}
+
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create stdout pipe: %w", err)
+	}
+
+	stderr, err := cmd.StderrPipe()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create stderr pipe: %w", err)
+	}
+
+	client := &Client{
+		Cmd:                   cmd,
+		stdin:                 stdin,
+		stdout:                bufio.NewReader(stdout),
+		stderr:                stderr,
+		handlers:              make(map[int32]chan *Message),
+		notificationHandlers:  make(map[string]NotificationHandler),
+		serverRequestHandlers: make(map[string]ServerRequestHandler),
+		diagnostics:           make(map[protocol.DocumentUri][]protocol.Diagnostic),
+		openFiles:             make(map[string]*OpenFileInfo),
+	}
+
+	// Start the LSP server process
+	if err := cmd.Start(); err != nil {
+		return nil, fmt.Errorf("failed to start LSP server: %w", err)
+	}
+
+	// Handle stderr in a separate goroutine
+	go func() {
+		scanner := bufio.NewScanner(stderr)
+		for scanner.Scan() {
+			fmt.Fprintf(os.Stderr, "LSP Server: %s\n", scanner.Text())
+		}
+		if err := scanner.Err(); err != nil {
+			fmt.Fprintf(os.Stderr, "Error reading stderr: %v\n", err)
+		}
+	}()
+
+	// Start message handling loop
+	go client.handleMessages()
+
+	return client, nil
+}
+
+func (c *Client) RegisterNotificationHandler(method string, handler NotificationHandler) {
+	c.notificationMu.Lock()
+	defer c.notificationMu.Unlock()
+	c.notificationHandlers[method] = handler
+}
+
+func (c *Client) RegisterServerRequestHandler(method string, handler ServerRequestHandler) {
+	c.serverHandlersMu.Lock()
+	defer c.serverHandlersMu.Unlock()
+	c.serverRequestHandlers[method] = handler
+}
+
+func (c *Client) InitializeLSPClient(ctx context.Context, workspaceDir string) (*protocol.InitializeResult, error) {
+	initParams := &protocol.InitializeParams{
+		WorkspaceFoldersInitializeParams: protocol.WorkspaceFoldersInitializeParams{
+			WorkspaceFolders: []protocol.WorkspaceFolder{
+				{
+					URI:  protocol.URI("file://" + workspaceDir),
+					Name: workspaceDir,
+				},
+			},
+		},
+
+		XInitializeParams: protocol.XInitializeParams{
+			ProcessID: int32(os.Getpid()),
+			ClientInfo: &protocol.ClientInfo{
+				Name:    "mcp-language-server",
+				Version: "0.1.0",
+			},
+			RootPath: workspaceDir,
+			RootURI:  protocol.DocumentUri("file://" + workspaceDir),
+			Capabilities: protocol.ClientCapabilities{
+				Workspace: protocol.WorkspaceClientCapabilities{
+					Configuration: true,
+					DidChangeConfiguration: protocol.DidChangeConfigurationClientCapabilities{
+						DynamicRegistration: true,
+					},
+					DidChangeWatchedFiles: protocol.DidChangeWatchedFilesClientCapabilities{
+						DynamicRegistration:    true,
+						RelativePatternSupport: true,
+					},
+				},
+				TextDocument: protocol.TextDocumentClientCapabilities{
+					Synchronization: &protocol.TextDocumentSyncClientCapabilities{
+						DynamicRegistration: true,
+						DidSave:             true,
+					},
+					Completion: protocol.CompletionClientCapabilities{
+						CompletionItem: protocol.ClientCompletionItemOptions{},
+					},
+					CodeLens: &protocol.CodeLensClientCapabilities{
+						DynamicRegistration: true,
+					},
+					DocumentSymbol: protocol.DocumentSymbolClientCapabilities{},
+					CodeAction: protocol.CodeActionClientCapabilities{
+						CodeActionLiteralSupport: protocol.ClientCodeActionLiteralOptions{
+							CodeActionKind: protocol.ClientCodeActionKindOptions{
+								ValueSet: []protocol.CodeActionKind{},
+							},
+						},
+					},
+					PublishDiagnostics: protocol.PublishDiagnosticsClientCapabilities{
+						VersionSupport: true,
+					},
+					SemanticTokens: protocol.SemanticTokensClientCapabilities{
+						Requests: protocol.ClientSemanticTokensRequestOptions{
+							Range: &protocol.Or_ClientSemanticTokensRequestOptions_range{},
+							Full:  &protocol.Or_ClientSemanticTokensRequestOptions_full{},
+						},
+						TokenTypes:     []string{},
+						TokenModifiers: []string{},
+						Formats:        []protocol.TokenFormat{},
+					},
+				},
+				Window: protocol.WindowClientCapabilities{},
+			},
+			InitializationOptions: map[string]any{
+				"codelenses": map[string]bool{
+					"generate":           true,
+					"regenerate_cgo":     true,
+					"test":               true,
+					"tidy":               true,
+					"upgrade_dependency": true,
+					"vendor":             true,
+					"vulncheck":          false,
+				},
+			},
+		},
+	}
+
+	var result protocol.InitializeResult
+	if err := c.Call(ctx, "initialize", initParams, &result); err != nil {
+		return nil, fmt.Errorf("initialize failed: %w", err)
+	}
+
+	if err := c.Notify(ctx, "initialized", struct{}{}); err != nil {
+		return nil, fmt.Errorf("initialized notification failed: %w", err)
+	}
+
+	// Register handlers
+	c.RegisterServerRequestHandler("workspace/applyEdit", HandleApplyEdit)
+	c.RegisterServerRequestHandler("workspace/configuration", HandleWorkspaceConfiguration)
+	c.RegisterServerRequestHandler("client/registerCapability", HandleRegisterCapability)
+	c.RegisterNotificationHandler("window/showMessage", HandleServerMessage)
+	c.RegisterNotificationHandler("textDocument/publishDiagnostics",
+		func(params json.RawMessage) { HandleDiagnostics(c, params) })
+
+	// Notify the LSP server
+	err := c.Initialized(ctx, protocol.InitializedParams{})
+	if err != nil {
+		return nil, fmt.Errorf("initialization failed: %w", err)
+	}
+
+	// LSP sepecific Initialization
+	path := strings.ToLower(c.Cmd.Path)
+	switch {
+	case strings.Contains(path, "typescript-language-server"):
+		// err := initializeTypescriptLanguageServer(ctx, c, workspaceDir)
+		// if err != nil {
+		// 	return nil, err
+		// }
+	}
+
+	return &result, nil
+}
+
+func (c *Client) Close() error {
+	// Try to close all open files first
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	// Attempt to close files but continue shutdown regardless
+	c.CloseAllFiles(ctx)
+
+	// Close stdin to signal the server
+	if err := c.stdin.Close(); err != nil {
+		return fmt.Errorf("failed to close stdin: %w", err)
+	}
+
+	// Use a channel to handle the Wait with timeout
+	done := make(chan error, 1)
+	go func() {
+		done <- c.Cmd.Wait()
+	}()
+
+	// Wait for process to exit with timeout
+	select {
+	case err := <-done:
+		return err
+	case <-time.After(2 * time.Second):
+		// If we timeout, try to kill the process
+		if err := c.Cmd.Process.Kill(); err != nil {
+			return fmt.Errorf("failed to kill process: %w", err)
+		}
+		return fmt.Errorf("process killed after timeout")
+	}
+}
+
+type ServerState int
+
+const (
+	StateStarting ServerState = iota
+	StateReady
+	StateError
+)
+
+func (c *Client) WaitForServerReady(ctx context.Context) error {
+	// TODO: wait for specific messages or poll workspace/symbol
+	time.Sleep(time.Second * 1)
+	return nil
+}
+
+type OpenFileInfo struct {
+	Version int32
+	URI     protocol.DocumentUri
+}
+
+func (c *Client) OpenFile(ctx context.Context, filepath string) error {
+	uri := fmt.Sprintf("file://%s", filepath)
+
+	c.openFilesMu.Lock()
+	if _, exists := c.openFiles[uri]; exists {
+		c.openFilesMu.Unlock()
+		return nil // Already open
+	}
+	c.openFilesMu.Unlock()
+
+	// Skip files that do not exist or cannot be read
+	content, err := os.ReadFile(filepath)
+	if err != nil {
+		return fmt.Errorf("error reading file: %w", err)
+	}
+
+	params := protocol.DidOpenTextDocumentParams{
+		TextDocument: protocol.TextDocumentItem{
+			URI:        protocol.DocumentUri(uri),
+			LanguageID: DetectLanguageID(uri),
+			Version:    1,
+			Text:       string(content),
+		},
+	}
+
+	if err := c.Notify(ctx, "textDocument/didOpen", params); err != nil {
+		return err
+	}
+
+	c.openFilesMu.Lock()
+	c.openFiles[uri] = &OpenFileInfo{
+		Version: 1,
+		URI:     protocol.DocumentUri(uri),
+	}
+	c.openFilesMu.Unlock()
+
+	return nil
+}
+
+func (c *Client) NotifyChange(ctx context.Context, filepath string) error {
+	uri := fmt.Sprintf("file://%s", filepath)
+
+	content, err := os.ReadFile(filepath)
+	if err != nil {
+		return fmt.Errorf("error reading file: %w", err)
+	}
+
+	c.openFilesMu.Lock()
+	fileInfo, isOpen := c.openFiles[uri]
+	if !isOpen {
+		c.openFilesMu.Unlock()
+		return fmt.Errorf("cannot notify change for unopened file: %s", filepath)
+	}
+
+	// Increment version
+	fileInfo.Version++
+	version := fileInfo.Version
+	c.openFilesMu.Unlock()
+
+	params := protocol.DidChangeTextDocumentParams{
+		TextDocument: protocol.VersionedTextDocumentIdentifier{
+			TextDocumentIdentifier: protocol.TextDocumentIdentifier{
+				URI: protocol.DocumentUri(uri),
+			},
+			Version: version,
+		},
+		ContentChanges: []protocol.TextDocumentContentChangeEvent{
+			{
+				Value: protocol.TextDocumentContentChangeWholeDocument{
+					Text: string(content),
+				},
+			},
+		},
+	}
+
+	return c.Notify(ctx, "textDocument/didChange", params)
+}
+
+func (c *Client) CloseFile(ctx context.Context, filepath string) error {
+	uri := fmt.Sprintf("file://%s", filepath)
+
+	c.openFilesMu.Lock()
+	if _, exists := c.openFiles[uri]; !exists {
+		c.openFilesMu.Unlock()
+		return nil // Already closed
+	}
+	c.openFilesMu.Unlock()
+
+	params := protocol.DidCloseTextDocumentParams{
+		TextDocument: protocol.TextDocumentIdentifier{
+			URI: protocol.DocumentUri(uri),
+		},
+	}
+	log.Println("Closing", params.TextDocument.URI.Dir())
+	if err := c.Notify(ctx, "textDocument/didClose", params); err != nil {
+		return err
+	}
+
+	c.openFilesMu.Lock()
+	delete(c.openFiles, uri)
+	c.openFilesMu.Unlock()
+
+	return nil
+}
+
+func (c *Client) IsFileOpen(filepath string) bool {
+	uri := fmt.Sprintf("file://%s", filepath)
+	c.openFilesMu.RLock()
+	defer c.openFilesMu.RUnlock()
+	_, exists := c.openFiles[uri]
+	return exists
+}
+
+// CloseAllFiles closes all currently open files
+func (c *Client) CloseAllFiles(ctx context.Context) {
+	c.openFilesMu.Lock()
+	filesToClose := make([]string, 0, len(c.openFiles))
+
+	// First collect all URIs that need to be closed
+	for uri := range c.openFiles {
+		// Convert URI back to file path by trimming "file://" prefix
+		filePath := strings.TrimPrefix(uri, "file://")
+		filesToClose = append(filesToClose, filePath)
+	}
+	c.openFilesMu.Unlock()
+
+	// Then close them all
+	for _, filePath := range filesToClose {
+		err := c.CloseFile(ctx, filePath)
+		if err != nil && debug {
+			log.Printf("Error closing file %s: %v", filePath, err)
+		}
+	}
+
+	if debug {
+		log.Printf("Closed %d files", len(filesToClose))
+	}
+}
+
+func (c *Client) GetFileDiagnostics(uri protocol.DocumentUri) []protocol.Diagnostic {
+	c.diagnosticsMu.RLock()
+	defer c.diagnosticsMu.RUnlock()
+
+	return c.diagnostics[uri]
+}
+
+func (c *Client) GetDiagnostics() map[protocol.DocumentUri][]protocol.Diagnostic {
+	return c.diagnostics
+}