main.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
)

type Result struct {
	Source     string `json:"source"`
	Target     string `json:"target"`
	Mode       string `json:"mode"`
	LinesAdded int    `json:"lines_added"`
	Line       int    `json:"line,omitempty"`
	Summary    string `json:"summary"`
}

type Config struct {
	Source  string
	Target  string
	Mode    string
	Line    int
	CLIMode bool
}

// MCP JSON-RPC types
type JSONRPCRequest struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      any             `json:"id"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

type JSONRPCResponse struct {
	JSONRPC string `json:"jsonrpc"`
	ID      any    `json:"id"`
	Result  any    `json:"result,omitempty"`
	Error   *Error `json:"error,omitempty"`
}

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

type InitializeResult struct {
	ProtocolVersion string       `json:"protocolVersion"`
	ServerInfo      ServerInfo   `json:"serverInfo"`
	Capabilities    Capabilities `json:"capabilities"`
}

type ServerInfo struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

type Capabilities struct {
	Tools map[string]bool `json:"tools"`
}

type ToolsListResult struct {
	Tools []Tool `json:"tools"`
}

type Tool struct {
	Name        string      `json:"name"`
	Description string      `json:"description"`
	InputSchema InputSchema `json:"inputSchema"`
}

type InputSchema struct {
	Type       string              `json:"type"`
	Properties map[string]Property `json:"properties"`
	Required   []string            `json:"required"`
}

type Property struct {
	Type        string `json:"type"`
	Description string `json:"description"`
	Enum        []string `json:"enum,omitempty"`
	Default     any    `json:"default,omitempty"`
}

type ToolCallParams struct {
	Name      string         `json:"name"`
	Arguments map[string]any `json:"arguments"`
}

type ToolCallResult struct {
	Content []ContentItem `json:"content"`
}

type ContentItem struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

// Exit codes for CLI mode
const (
	ExitSuccess = 0
	ExitError   = 1
)

func main() {
	config := parseFlags()

	if config.CLIMode {
		runCLI(config)
		return
	}

	runMCPServer()
}

func parseFlags() Config {
	config := Config{}

	flag.BoolVar(&config.CLIMode, "cli", false, "Run in CLI mode (default is MCP server mode)")
	flag.StringVar(&config.Source, "source", "", "File to read content from (required)")
	flag.StringVar(&config.Target, "target", "", "File to write content to (required)")
	flag.StringVar(&config.Mode, "mode", "", "Operation mode: append|prepend|replace|insert (required)")
	flag.IntVar(&config.Line, "line", 0, "Line number for insert mode (required for insert)")

	flag.Parse()

	return config
}

func runCLI(config Config) {
	if err := validateConfig(config); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		flag.Usage()
		os.Exit(ExitError)
	}

	result, err := splice(config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(ExitError)
	}

	output, err := json.Marshal(result)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error marshaling JSON: %v\n", err)
		os.Exit(ExitError)
	}

	fmt.Println(string(output))
}

func validateConfig(config Config) error {
	if config.Source == "" {
		return fmt.Errorf("--source is required")
	}
	if config.Target == "" {
		return fmt.Errorf("--target is required")
	}
	switch config.Mode {
	case "append", "prepend", "replace", "insert":
		// valid
	case "":
		return fmt.Errorf("--mode is required (append|prepend|replace|insert)")
	default:
		return fmt.Errorf("invalid mode %q: must be append|prepend|replace|insert", config.Mode)
	}
	if config.Mode == "insert" && config.Line < 1 {
		return fmt.Errorf("--line is required for insert mode and must be >= 1")
	}
	return nil
}

func runMCPServer() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-sigChan
		fmt.Fprintln(os.Stderr, "Received shutdown signal, exiting gracefully...")
		cancel()
	}()

	scanner := bufio.NewScanner(os.Stdin)

	lineChan := make(chan string)
	errChan := make(chan error, 1)

	go func() {
		for scanner.Scan() {
			lineChan <- scanner.Text()
		}
		if err := scanner.Err(); err != nil {
			errChan <- err
		}
		close(lineChan)
	}()

	for {
		select {
		case <-ctx.Done():
			return
		case err := <-errChan:
			fmt.Fprintf(os.Stderr, "Scanner error: %v\n", err)
			return
		case line, ok := <-lineChan:
			if !ok {
				return
			}
			if line == "" {
				continue
			}

			var req JSONRPCRequest
			if err := json.Unmarshal([]byte(line), &req); err != nil {
				sendError(nil, -32700, "Parse error")
				continue
			}

			handleRequest(req)
		}
	}
}

func handleRequest(req JSONRPCRequest) {
	isNotification := req.ID == nil

	switch req.Method {
	case "initialize":
		handleInitialize(req)
	case "notifications/initialized":
		return
	case "tools/list":
		handleToolsList(req)
	case "tools/call":
		handleToolsCall(req)
	default:
		if isNotification {
			return
		}
		sendError(req.ID, -32601, "Method not found")
	}
}

func handleInitialize(req JSONRPCRequest) {
	result := InitializeResult{
		ProtocolVersion: "2024-11-05",
		ServerInfo: ServerInfo{
			Name:    "splice",
			Version: "1.0.0",
		},
		Capabilities: Capabilities{
			Tools: map[string]bool{
				"list": true,
				"call": true,
			},
		},
	}
	sendResponse(req.ID, result)
}

func handleToolsList(req JSONRPCRequest) {
	result := ToolsListResult{
		Tools: []Tool{
			{
				Name:        "splice",
				Description: "Splice file contents into a target file. Supports four modes: append (add after last line), prepend (add before first line), replace (overwrite target with source), and insert (split target at line N and insert source between halves).",
				InputSchema: InputSchema{
					Type: "object",
					Properties: map[string]Property{
						"source": {
							Type:        "string",
							Description: "Absolute path to the file to read content from.",
						},
						"target": {
							Type:        "string",
							Description: "Absolute path to the file to write content to.",
						},
						"mode": {
							Type:        "string",
							Description: "Operation mode.",
							Enum:        []string{"append", "prepend", "replace", "insert"},
						},
						"line": {
							Type:        "integer",
							Description: "Line number for insert mode. Content is inserted after this line. Must be >= 1.",
						},
					},
					Required: []string{"source", "target", "mode"},
				},
			},
		},
	}
	sendResponse(req.ID, result)
}

func handleToolsCall(req JSONRPCRequest) {
	var params ToolCallParams
	if err := json.Unmarshal(req.Params, &params); err != nil {
		sendError(req.ID, -32602, "Invalid params")
		return
	}

	if params.Name != "splice" {
		sendError(req.ID, -32602, "Unknown tool")
		return
	}

	source, ok := params.Arguments["source"].(string)
	if !ok || source == "" {
		sendError(req.ID, -32602, "Missing or invalid 'source' parameter")
		return
	}

	target, ok := params.Arguments["target"].(string)
	if !ok || target == "" {
		sendError(req.ID, -32602, "Missing or invalid 'target' parameter")
		return
	}

	mode, ok := params.Arguments["mode"].(string)
	if !ok || mode == "" {
		sendError(req.ID, -32602, "Missing or invalid 'mode' parameter")
		return
	}

	config := Config{
		Source: source,
		Target: target,
		Mode:   mode,
	}

	if lineVal, exists := params.Arguments["line"]; exists {
		switch v := lineVal.(type) {
		case float64:
			config.Line = int(v)
		case int:
			config.Line = v
		}
	}

	if err := validateConfig(config); err != nil {
		sendError(req.ID, -32602, err.Error())
		return
	}

	result, err := splice(config)
	if err != nil {
		sendError(req.ID, -32603, fmt.Sprintf("Splice failed: %v", err))
		return
	}

	jsonResult, err := json.Marshal(result)
	if err != nil {
		sendError(req.ID, -32603, "Failed to marshal result")
		return
	}

	response := ToolCallResult{
		Content: []ContentItem{
			{
				Type: "text",
				Text: string(jsonResult),
			},
		},
	}

	sendResponse(req.ID, response)
}

func sendResponse(id any, result any) {
	resp := JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      id,
		Result:  result,
	}
	data, err := json.Marshal(resp)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to marshal response: %v\n", err)
		return
	}
	fmt.Println(string(data))
}

func sendError(id any, code int, message string) {
	resp := JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      id,
		Error: &Error{
			Code:    code,
			Message: message,
		},
	}
	data, err := json.Marshal(resp)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to marshal error response: %v\n", err)
		return
	}
	fmt.Println(string(data))
}

func splice(config Config) (*Result, error) {
	sourceData, err := os.ReadFile(config.Source)
	if err != nil {
		return nil, fmt.Errorf("failed to read source file: %w", err)
	}

	sourceContent := string(sourceData)
	sourceLineCount := countLines(sourceContent)

	result := &Result{
		Source:     config.Source,
		Target:     config.Target,
		Mode:       config.Mode,
		LinesAdded: sourceLineCount,
	}

	switch config.Mode {
	case "replace":
		err = writeFileAtomic(config.Target, sourceData)
	case "append":
		err = spliceAppend(config.Target, sourceData)
	case "prepend":
		err = splicePrepend(config.Target, sourceData)
	case "insert":
		result.Line = config.Line
		err = spliceInsert(config.Target, sourceData, config.Line)
	}

	if err != nil {
		return nil, err
	}

	result.Summary = buildSummary(config, sourceLineCount)
	return result, nil
}

func spliceAppend(targetPath string, sourceData []byte) error {
	targetData, err := os.ReadFile(targetPath)
	if err != nil {
		return fmt.Errorf("failed to read target file: %w", err)
	}

	var combined []byte
	if len(targetData) > 0 && targetData[len(targetData)-1] != '\n' {
		combined = make([]byte, 0, len(targetData)+1+len(sourceData))
		combined = append(combined, targetData...)
		combined = append(combined, '\n')
		combined = append(combined, sourceData...)
	} else {
		combined = make([]byte, 0, len(targetData)+len(sourceData))
		combined = append(combined, targetData...)
		combined = append(combined, sourceData...)
	}

	return writeFileAtomic(targetPath, combined)
}

func splicePrepend(targetPath string, sourceData []byte) error {
	targetData, err := os.ReadFile(targetPath)
	if err != nil {
		return fmt.Errorf("failed to read target file: %w", err)
	}

	var combined []byte
	if len(sourceData) > 0 && sourceData[len(sourceData)-1] != '\n' {
		combined = make([]byte, 0, len(sourceData)+1+len(targetData))
		combined = append(combined, sourceData...)
		combined = append(combined, '\n')
		combined = append(combined, targetData...)
	} else {
		combined = make([]byte, 0, len(sourceData)+len(targetData))
		combined = append(combined, sourceData...)
		combined = append(combined, targetData...)
	}

	return writeFileAtomic(targetPath, combined)
}

func spliceInsert(targetPath string, sourceData []byte, lineNum int) error {
	targetData, err := os.ReadFile(targetPath)
	if err != nil {
		return fmt.Errorf("failed to read target file: %w", err)
	}

	targetContent := string(targetData)
	lines := strings.Split(targetContent, "\n")

	// Handle trailing newline: if file ends with \n, Split produces an empty last element
	hasTrailingNewline := len(targetContent) > 0 && targetContent[len(targetContent)-1] == '\n'
	if hasTrailingNewline && len(lines) > 0 && lines[len(lines)-1] == "" {
		lines = lines[:len(lines)-1]
	}

	if lineNum > len(lines) {
		return fmt.Errorf("line %d exceeds file length (%d lines)", lineNum, len(lines))
	}

	// Ensure source content doesn't have a trailing newline for clean joining
	sourceStr := string(sourceData)
	if len(sourceStr) > 0 && sourceStr[len(sourceStr)-1] == '\n' {
		sourceStr = sourceStr[:len(sourceStr)-1]
	}
	sourceLines := strings.Split(sourceStr, "\n")

	// Build new content: lines[0:lineNum] + sourceLines + lines[lineNum:]
	newLines := make([]string, 0, len(lines)+len(sourceLines))
	newLines = append(newLines, lines[:lineNum]...)
	newLines = append(newLines, sourceLines...)
	newLines = append(newLines, lines[lineNum:]...)

	output := strings.Join(newLines, "\n")
	if hasTrailingNewline {
		output += "\n"
	}

	return writeFileAtomic(targetPath, []byte(output))
}

func writeFileAtomic(path string, data []byte) error {
	resolvedPath, err := filepath.EvalSymlinks(path)
	if err != nil {
		if os.IsNotExist(err) {
			resolvedPath = path
		} else {
			return fmt.Errorf("failed to resolve path: %w", err)
		}
	}

	mode := os.FileMode(0644)
	if info, err := os.Stat(resolvedPath); err == nil {
		mode = info.Mode()
		if mode&0200 == 0 {
			return fmt.Errorf("file is read-only: %s", resolvedPath)
		}
	}

	dir := filepath.Dir(resolvedPath)
	tmpFile, err := os.CreateTemp(dir, ".splice-*.tmp")
	if err != nil {
		return fmt.Errorf("failed to create temp file: %w", err)
	}
	tmpPath := tmpFile.Name()

	success := false
	defer func() {
		if !success {
			os.Remove(tmpPath)
		}
	}()

	if _, err := tmpFile.Write(data); err != nil {
		tmpFile.Close()
		return err
	}

	if err := tmpFile.Sync(); err != nil {
		tmpFile.Close()
		return fmt.Errorf("failed to sync file: %w", err)
	}

	if err := tmpFile.Close(); err != nil {
		return fmt.Errorf("failed to close temp file: %w", err)
	}

	if err := os.Chmod(tmpPath, mode); err != nil {
		return fmt.Errorf("failed to set permissions: %w", err)
	}

	if err := os.Rename(tmpPath, resolvedPath); err != nil {
		return fmt.Errorf("failed to rename temp file: %w", err)
	}

	success = true
	return nil
}

func countLines(s string) int {
	if s == "" {
		return 0
	}
	n := strings.Count(s, "\n")
	if s[len(s)-1] != '\n' {
		n++
	}
	return n
}

func buildSummary(config Config, lineCount int) string {
	lineWord := "line"
	if lineCount != 1 {
		lineWord = "lines"
	}

	switch config.Mode {
	case "append":
		return fmt.Sprintf("Appended %d %s from %s to %s", lineCount, lineWord, config.Source, config.Target)
	case "prepend":
		return fmt.Sprintf("Prepended %d %s from %s to %s", lineCount, lineWord, config.Source, config.Target)
	case "replace":
		return fmt.Sprintf("Replaced %s with %d %s from %s", config.Target, lineCount, lineWord, config.Source)
	case "insert":
		return fmt.Sprintf("Inserted %d %s from %s into %s at line %d", lineCount, lineWord, config.Source, config.Target, config.Line)
	}
	return ""
}