path: root/internal/lsp/handlers.go

package lsp

import (
	"context"
	"encoding/json"
	"fmt"
	"hexai/internal"
	"hexai/internal/llm"
	"hexai/internal/logging"
	"os"
	"strings"
	"time"
)

func (s *Server) handle(req Request) {
    switch req.Method {
    case "initialize":
        s.handleInitialize(req)
    case "initialized":
        s.handleInitialized()
	case "shutdown":
		s.handleShutdown(req)
	case "exit":
		s.handleExit()
	case "textDocument/didOpen":
		s.handleDidOpen(req)
	case "textDocument/didChange":
		s.handleDidChange(req)
	case "textDocument/didClose":
		s.handleDidClose(req)
    case "textDocument/completion":
        s.handleCompletion(req)
    case "textDocument/codeAction":
        s.handleCodeAction(req)
    default:
        if len(req.ID) != 0 {
            s.reply(req.ID, nil, &RespError{Code: -32601, Message: fmt.Sprintf("method not found: %s", req.Method)})
        }
    }
}

func (s *Server) handleInitialize(req Request) {
	version := internal.Version
	if s.llmClient != nil {
		version = version + " [" + s.llmClient.Name() + ":" + s.llmClient.DefaultModel() + "]"
	}
    res := InitializeResult{
        Capabilities: ServerCapabilities{
            TextDocumentSync: 1, // 1 = TextDocumentSyncKindFull
            CompletionProvider: &CompletionOptions{
                ResolveProvider: false,
                // TODO: Make the trigger characters configurable
                TriggerCharacters: []string{".", ":", "/", "_"},
            },
            CodeActionProvider: true,
        },
        ServerInfo: &ServerInfo{Name: "hexai", Version: version},
    }
    s.reply(req.ID, res, nil)
}

func (s *Server) handleCodeAction(req Request) {
    var p CodeActionParams
    if err := json.Unmarshal(req.Params, &p); err != nil {
        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
        return
    }
    if s.llmClient == nil {
        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
        return
    }
    // Extract selected text
    d := s.getDocument(p.TextDocument.URI)
    if d == nil || len(d.lines) == 0 {
        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
        return
    }
    sel := extractRangeText(d, p.Range)
    if strings.TrimSpace(sel) == "" {
        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
        return
    }
    // Derive instruction from selection comments (prefer first), including ;text; marker
    instr, cleaned := instructionFromSelection(sel)
    if strings.TrimSpace(instr) == "" {
        // No instruction; do not offer an action
        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
        return
    }
    // Build prompt for rewrite of cleaned selection according to instruction
    sys := "You are a precise code refactoring engine. Rewrite the given code strictly according to the instruction. Return only the updated code with no prose or backticks. Preserve formatting where reasonable."
    user := fmt.Sprintf("Instruction: %s\n\nSelected code to transform:\n%s", instr, cleaned)
    ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
    defer cancel()
    messages := []llm.Message{{Role: "system", Content: sys}, {Role: "user", Content: user}}
    text, err := s.llmClient.Chat(ctx, messages, llm.WithMaxTokens(s.maxTokens), llm.WithTemperature(0.1))
    if err != nil {
        logging.Logf("lsp ", "codeAction llm error: %v", err)
        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
        return
    }
    out := strings.TrimSpace(text)
    if out == "" {
        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
        return
    }
    edit := WorkspaceEdit{Changes: map[string][]TextEdit{p.TextDocument.URI: {{Range: p.Range, NewText: out}}}}
    action := CodeAction{Title: "Hexai: rewrite selection", Kind: "refactor.rewrite", Edit: &edit}
    if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{action}, nil) }
}

// instructionFromSelection extracts the first instruction from selection text.
// Preference order on each line: strict ;text; marker (no inner spaces), then
// a line comment (//, #, --). Returns the instruction string and the selection
// text cleaned of the matched instruction marker or comment.
func instructionFromSelection(sel string) (string, string) {
    lines := splitLines(sel)
    for idx, line := range lines {
        if instr, cleaned, ok := findFirstInstructionInLine(line); ok && strings.TrimSpace(instr) != "" {
            lines[idx] = cleaned
            return instr, strings.Join(lines, "\n")
        }
    }
    return "", sel
}

// findFirstInstructionInLine returns the earliest instruction marker on the
// line and the line with that marker removed. Supported markers, ordered by
// earliest byte offset in the line:
// - ;text; (strict, no space after first ';' or before last ';')
// - /* text */ (single-line only)
// - <!-- text --> (single-line only)
// - // text
// - # text
// - -- text
func findFirstInstructionInLine(line string) (instr string, cleaned string, ok bool) {
    type cand struct{ start, end int; text string }
    cands := []cand{}
    if t, l, r, ok := findStrictSemicolonTag(line); ok {
        cands = append(cands, cand{start: l, end: r, text: t})
    }
    if i := strings.Index(line, "/*"); i >= 0 {
        if j := strings.Index(line[i+2:], "*/"); j >= 0 {
            start := i
            end := i + 2 + j + 2
            text := strings.TrimSpace(line[i+2 : i+2+j])
            cands = append(cands, cand{start: start, end: end, text: text})
        }
    }
    if i := strings.Index(line, "<!--"); i >= 0 {
        if j := strings.Index(line[i+4:], "-->"); j >= 0 {
            start := i
            end := i + 4 + j + 3
            text := strings.TrimSpace(line[i+4 : i+4+j])
            cands = append(cands, cand{start: start, end: end, text: text})
        }
    }
    if i := strings.Index(line, "//"); i >= 0 { cands = append(cands, cand{start: i, end: len(line), text: strings.TrimSpace(line[i+2:])}) }
    if i := strings.Index(line, "#"); i >= 0 { cands = append(cands, cand{start: i, end: len(line), text: strings.TrimSpace(line[i+1:])}) }
    if i := strings.Index(line, "--"); i >= 0 { cands = append(cands, cand{start: i, end: len(line), text: strings.TrimSpace(line[i+2:])}) }
    if len(cands) == 0 { return "", line, false }
    // pick earliest start index
    best := cands[0]
    for _, c := range cands[1:] {
        if c.start >= 0 && (best.start < 0 || c.start < best.start) {
            best = c
        }
    }
    cleaned = strings.TrimRight(line[:best.start]+line[best.end:], " \t")
    return best.text, cleaned, true
}

// findStrictSemicolonTag finds ;text; with no space after first ';' and no space
// before the last ';' on the given line. Returns the text between semicolons,
// the start index of the opening ';', the end index just after the closing ';',
// and whether it was found.
func findStrictSemicolonTag(line string) (string, int, int, bool) {
    pos := 0
    for pos < len(line) {
        j := strings.Index(line[pos:], ";")
        if j < 0 { return "", 0, 0, false }
        j += pos
        // ensure single ';' (not ';;') and non-space after
        if j+1 >= len(line) || line[j+1] == ';' || line[j+1] == ' ' { pos = j + 1; continue }
        k := strings.Index(line[j+1:], ";")
        if k < 0 { return "", 0, 0, false }
        closeIdx := j + 1 + k
        if closeIdx-1 < 0 || line[closeIdx-1] == ' ' { pos = closeIdx + 1; continue }
        inner := strings.TrimSpace(line[j+1 : closeIdx])
        if inner == "" { pos = closeIdx + 1; continue }
        end := closeIdx + 1
        return inner, j, end, true
    }
    return "", 0, 0, false
}

// extractRangeText returns the exact text within the given document range.
func extractRangeText(d *document, r Range) string {
    if r.Start.Line == r.End.Line {
        line := d.lines[r.Start.Line]
        if r.Start.Character < 0 { r.Start.Character = 0 }
        if r.End.Character > len(line) { r.End.Character = len(line) }
        if r.Start.Character > r.End.Character { return "" }
        return line[r.Start.Character:r.End.Character]
    }
    var b strings.Builder
    // first line
    first := d.lines[r.Start.Line]
    if r.Start.Character < 0 { r.Start.Character = 0 }
    if r.Start.Character > len(first) { r.Start.Character = len(first) }
    b.WriteString(first[r.Start.Character:])
    b.WriteString("\n")
    // middle lines
    for i := r.Start.Line + 1; i < r.End.Line; i++ {
        b.WriteString(d.lines[i])
        if i+1 <= r.End.Line { b.WriteString("\n") }
    }
    // last line
    last := d.lines[r.End.Line]
    if r.End.Character < 0 { r.End.Character = 0 }
    if r.End.Character > len(last) { r.End.Character = len(last) }
    b.WriteString(last[:r.End.Character])
    return b.String()
}

func (s *Server) handleInitialized() {
	logging.Logf("lsp ", "client initialized")
}

func (s *Server) handleShutdown(req Request) {
	s.reply(req.ID, nil, nil)
}

func (s *Server) handleExit() {
	s.exited = true
	os.Exit(0)
}

func (s *Server) handleDidOpen(req Request) {
	var p DidOpenTextDocumentParams
	if err := json.Unmarshal(req.Params, &p); err == nil {
		s.setDocument(p.TextDocument.URI, p.TextDocument.Text)
		s.markActivity()
	}
}

func (s *Server) handleDidChange(req Request) {
	var p DidChangeTextDocumentParams
	if err := json.Unmarshal(req.Params, &p); err == nil {
		if len(p.ContentChanges) > 0 {
			s.setDocument(p.TextDocument.URI, p.ContentChanges[len(p.ContentChanges)-1].Text)
		}
		s.markActivity()
	}
}

func (s *Server) handleDidClose(req Request) {
	var p DidCloseTextDocumentParams
	if err := json.Unmarshal(req.Params, &p); err == nil {
		s.deleteDocument(p.TextDocument.URI)
		s.markActivity()
	}
}

func (s *Server) handleCompletion(req Request) {
	var p CompletionParams
	var docStr string
	if err := json.Unmarshal(req.Params, &p); err == nil {
		above, current, below, funcCtx := s.lineContext(p.TextDocument.URI, p.Position)
		docStr = s.buildDocString(p, above, current, below, funcCtx)
		if s.logContext {
			s.logCompletionContext(p, above, current, below, funcCtx)
		}
		if s.llmClient != nil {
			newFunc := s.isDefiningNewFunction(p.TextDocument.URI, p.Position)
			extra, has := s.buildAdditionalContext(newFunc, p.TextDocument.URI, p.Position)
			items, ok := s.tryLLMCompletion(p, above, current, below, funcCtx, docStr, has, extra)
			if ok {
				s.reply(req.ID, CompletionList{IsIncomplete: false, Items: items}, nil)
				return
			}
		}
	}
	items := s.fallbackCompletionItems(docStr)
	s.reply(req.ID, CompletionList{IsIncomplete: false, Items: items}, nil)
}

func (s *Server) reply(id json.RawMessage, result any, err *RespError) {
	resp := Response{JSONRPC: "2.0", ID: id, Result: result, Error: err}
	s.writeMessage(resp)
}

// --- completion helpers ---

func (s *Server) buildDocString(p CompletionParams, above, current, below, funcCtx string) string {
	return fmt.Sprintf("file: %s\nline: %d\nabove: %s\ncurrent: %s\nbelow: %s\nfunction: %s",
		p.TextDocument.URI, p.Position.Line, trimLen(above), trimLen(current), trimLen(below), trimLen(funcCtx))
}

func (s *Server) logCompletionContext(p CompletionParams, above, current, below, funcCtx string) {
	logging.Logf("lsp ", "completion ctx uri=%s line=%d char=%d above=%q current=%q below=%q function=%q",
		p.TextDocument.URI, p.Position.Line, p.Position.Character, trimLen(above), trimLen(current), trimLen(below), trimLen(funcCtx))
}

func (s *Server) tryLLMCompletion(p CompletionParams, above, current, below, funcCtx, docStr string, hasExtra bool, extraText string) ([]CompletionItem, bool) {
	ctx, cancel := context.WithTimeout(context.Background(), 6*time.Second)
	defer cancel()

	inParams := inParamList(current, p.Position.Character)
	sysPrompt, userPrompt := buildPrompts(inParams, p, above, current, below, funcCtx)
	messages := []llm.Message{
		{Role: "system", Content: sysPrompt},
		{Role: "user", Content: userPrompt},
	}
	if hasExtra && extraText != "" {
		messages = append(messages, llm.Message{Role: "user", Content: "Additional context:\n" + extraText})
	}

	// Compute total sent context size (sum of message contents)
	var sentSize int
	for _, m := range messages {
		sentSize += len(m.Content)
	}
	// Update request counters (sent)
	s.mu.Lock()
	s.llmReqTotal++
	s.llmSentBytesTotal += int64(sentSize)
	s.mu.Unlock()

	text, err := s.llmClient.Chat(ctx, messages, llm.WithMaxTokens(s.maxTokens), llm.WithTemperature(0.2))
	if err != nil {
		logging.Logf("lsp ", "llm completion error: %v", err)
		// Log updated averages after this request (even if failed)
		s.mu.RLock()
		avgSent := int64(0)
		if s.llmReqTotal > 0 {
			avgSent = s.llmSentBytesTotal / s.llmReqTotal
		}
		avgRecv := int64(0)
		if s.llmRespTotal > 0 {
			avgRecv = s.llmRespBytesTotal / s.llmRespTotal
		}
		reqs, sentTot, recvTot := s.llmReqTotal, s.llmSentBytesTotal, s.llmRespBytesTotal
		s.mu.RUnlock()
		mins := time.Since(s.startTime).Minutes()
		if mins <= 0 {
			mins = 0.001
		}
		rpm := float64(reqs) / mins
		sentPerMin := float64(sentTot) / mins
		recvPerMin := float64(recvTot) / mins
		logging.Logf("lsp ", "llm stats reqs=%d avg_sent=%d avg_recv=%d sent_total=%d recv_total=%d rpm=%.2f sent_per_min=%.0f recv_per_min=%.0f", reqs, avgSent, avgRecv, sentTot, recvTot, rpm, sentPerMin, recvPerMin)
		return nil, false
	}
	// Update response counters (received)
	recvSize := len(text)
	s.mu.Lock()
	s.llmRespTotal++
	s.llmRespBytesTotal += int64(recvSize)
	avgSent := int64(0)
	if s.llmReqTotal > 0 {
		avgSent = s.llmSentBytesTotal / s.llmReqTotal
	}
	avgRecv := int64(0)
	if s.llmRespTotal > 0 {
		avgRecv = s.llmRespBytesTotal / s.llmRespTotal
	}
	reqs, sentTot, recvTot := s.llmReqTotal, s.llmSentBytesTotal, s.llmRespBytesTotal
	s.mu.Unlock()
	mins := time.Since(s.startTime).Minutes()
	if mins <= 0 {
		mins = 0.001
	}
	rpm := float64(reqs) / mins
	sentPerMin := float64(sentTot) / mins
	recvPerMin := float64(recvTot) / mins
	logging.Logf("lsp ", "llm stats reqs=%d avg_sent=%d avg_recv=%d sent_total=%d recv_total=%d rpm=%.2f sent_per_min=%.0f recv_per_min=%.0f", reqs, avgSent, avgRecv, sentTot, recvTot, rpm, sentPerMin, recvPerMin)
	cleaned := strings.TrimSpace(text)
	if cleaned == "" {
		return nil, false
	}

	te, filter := computeTextEditAndFilter(cleaned, inParams, current, p)
	rm := s.collectPromptRemovalEdits(p.TextDocument.URI)
	label := labelForCompletion(cleaned, filter)
	// Detail shows provider/model for visibility in client UI
	detail := "Hexai LLM completion"
	if s.llmClient != nil {
		detail = "Hexai " + s.llmClient.Name() + ":" + s.llmClient.DefaultModel()
	}
	items := []CompletionItem{{
		Label:               label,
		Kind:                1,
		Detail:              detail,
		InsertTextFormat:    1,
		FilterText:          strings.TrimLeft(filter, " \t"),
		TextEdit:            te,
		AdditionalTextEdits: rm,
		SortText:            "0000",
		Documentation:       docStr,
	}}
	return items, true
}

// collectPromptRemovalEdits returns edits to remove all inline prompt markers.
// Supported form (inclusive):
// - ";...;" where there is no space immediately after the first ';'
//   and no space immediately before the last ';'. An optional single space
//   after the trailing ';' is also removed for cleanliness.
// Multiple markers per line are supported.
func (s *Server) collectPromptRemovalEdits(uri string) []TextEdit {
    d := s.getDocument(uri)
    if d == nil || len(d.lines) == 0 {
        return nil
    }
    var edits []TextEdit
    for i, line := range d.lines {
        // If the line contains a double-semicolon trigger of the form
        // ";;text;" (no space after the ";;" and no space before the closing ';'),
        // remove the entire line.
        removeWholeLine := false
        {
            pos := 0
            for pos < len(line) {
                j := strings.Index(line[pos:], ";;")
                if j < 0 { break }
                j += pos
                // ensure there's a non-space after the two semicolons
                if j+2 >= len(line) || line[j+2] == ' ' { pos = j + 2; continue }
                // find closing ';' after the content
                k := strings.Index(line[j+2:], ";")
                if k < 0 { break }
                closeIdx := j + 2 + k
                // ensure char before closing ';' is not a space
                if closeIdx-1 < 0 || line[closeIdx-1] == ' ' { pos = closeIdx + 1; continue }
                removeWholeLine = true
                break
            }
        }
        if removeWholeLine {
            edits = append(edits, TextEdit{Range: Range{Start: Position{Line: i, Character: 0}, End: Position{Line: i, Character: len(line)}}, NewText: ""})
            continue
        }
        // Scan for ;...; markers that have no spaces directly inside the semicolons
        startSemi := 0
        for startSemi < len(line) {
            j := strings.Index(line[startSemi:], ";")
            if j < 0 {
                break
            }
            j += startSemi
            k := strings.Index(line[j+1:], ";")
            if k < 0 {
                break
            }
            // Require no space immediately after the first ';'
            if j+1 >= len(line) || line[j+1] == ' ' {
                startSemi = j + 1
                continue
            }
            // Ignore patterns that start with double semicolon here; handled above
            if line[j+1] == ';' {
                startSemi = j + 2
                continue
            }
            // Index of the closing ';'
            closeIdx := j + 1 + k
            // Require no space immediately before the closing ';'
            if closeIdx-1 < 0 || line[closeIdx-1] == ' ' {
                startSemi = closeIdx + 1
                continue
            }
            // Require at least one character between the semicolons
            if closeIdx-(j+1) < 1 {
                startSemi = closeIdx + 1
                continue
            }
            endChar := closeIdx + 1 // include trailing ';'
            if endChar < len(line) && line[endChar] == ' ' {
                endChar++
            }
            edits = append(edits, TextEdit{Range: Range{Start: Position{Line: i, Character: j}, End: Position{Line: i, Character: endChar}}, NewText: ""})
            startSemi = endChar
        }
    }
    return edits
}

func inParamList(current string, cursor int) bool {
	if !strings.Contains(current, "func ") {
		return false
	}
	open := strings.Index(current, "(")
	close := strings.Index(current, ")")
	return open >= 0 && cursor > open && (close == -1 || cursor <= close)
}

func buildPrompts(inParams bool, p CompletionParams, above, current, below, funcCtx string) (string, string) {
	if inParams {
		sys := "You are a code completion engine for function signatures. Return only the parameter list contents (without parentheses), no braces, no prose. Prefer idiomatic names and types."
		user := fmt.Sprintf("Cursor is inside the function parameter list. Suggest only the parameter list (no parentheses).\nFunction line: %s\nCurrent line (cursor at %d): %s", funcCtx, p.Position.Character, current)
		return sys, user
	}
	sys := "You are a terse code completion engine. Return only the code to insert, no surrounding prose or backticks."
	user := fmt.Sprintf("Provide the next likely code to insert at the cursor.\nFile: %s\nFunction/context: %s\nAbove line: %s\nCurrent line (cursor at character %d): %s\nBelow line: %s\nOnly return the completion snippet.", p.TextDocument.URI, funcCtx, above, p.Position.Character, current, below)
	return sys, user
}

func computeTextEditAndFilter(cleaned string, inParams bool, current string, p CompletionParams) (*TextEdit, string) {
	if inParams {
		open := strings.Index(current, "(")
		close := strings.Index(current, ")")
		if open >= 0 {
			left := open + 1
			right := len(current)
			if close >= 0 && close >= left {
				right = close
			}
			if p.Position.Character < right {
				right = p.Position.Character
			}
			te := &TextEdit{Range: Range{Start: Position{Line: p.Position.Line, Character: left}, End: Position{Line: p.Position.Line, Character: right}}, NewText: cleaned}
			var filter string
			if left >= 0 && right >= left && right <= len(current) {
				filter = strings.TrimLeft(current[left:right], " \t")
			}
			return te, filter
		}
	}
	startChar := computeWordStart(current, p.Position.Character)
	te := &TextEdit{Range: Range{Start: Position{Line: p.Position.Line, Character: startChar}, End: Position{Line: p.Position.Line, Character: p.Position.Character}}, NewText: cleaned}
	filter := strings.TrimLeft(current[startChar:p.Position.Character], " \t")
	return te, filter
}

func computeWordStart(current string, at int) int {
	if at > len(current) {
		at = len(current)
	}
	for at > 0 {
		ch := current[at-1]
		if (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch >= '0' && ch <= '9') || ch == '_' {
			at--
			continue
		}
		break
	}
	return at
}

func labelForCompletion(cleaned, filter string) string {
	label := trimLen(firstLine(cleaned))
	if filter != "" && !strings.HasPrefix(strings.ToLower(label), strings.ToLower(filter)) {
		return filter
	}
	return label
}

func (s *Server) fallbackCompletionItems(docStr string) []CompletionItem {
	return []CompletionItem{{
		Label:         "hexai-complete",
		Kind:          1,
		Detail:        "dummy completion",
		InsertText:    "hexai",
		SortText:      "9999",
		Documentation: docStr,
	}}
}