// Package mcp implements a minimal Model Context Protocol server over stdio.
//
// MCP is a JSON-RPC 2.0 protocol for AI agents to discover and call tools.
// The stdio transport is newline-delimited JSON: one JSON object per line
// on stdin, one JSON response per line on stdout. All diagnostics go to
// stderr — stdout is reserved for the protocol channel.
//
// Spec: https://modelcontextprotocol.io
package mcp

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sync"
)

// ProtocolVersion — the MCP version we implement. Clients negotiate via
// the initialize handshake; if they ask for a different version we accept
// theirs and hope for the best (graceful degradation).
const ProtocolVersion = "2024-11-05"

// ServerInfo identifies the server to the client.
type ServerInfo struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

// ToolHandler is the signature for a registered tool's Go implementation.
// args is the decoded arguments map (types per the tool's input schema);
// the returned value is JSON-encoded as the tool's result.
type ToolHandler func(ctx context.Context, args map[string]any) (any, error)

// Tool describes a single callable function — name, description, JSON
// schema for inputs, and the Go handler.
type Tool struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema map[string]any  `json:"inputSchema"`
	Handler     ToolHandler     `json:"-"`
}

// Server is the stdio MCP loop.
type Server struct {
	info  ServerInfo
	tools []Tool
	mu    sync.Mutex // guards writes to stdout
	w     *bufio.Writer
}

func New(info ServerInfo) *Server {
	return &Server{
		info: info,
		w:    bufio.NewWriter(os.Stdout),
	}
}

// RegisterTool adds a tool to the server's tool list.
func (s *Server) RegisterTool(t Tool) {
	s.tools = append(s.tools, t)
}

// Run reads line-delimited JSON-RPC 2.0 requests from stdin and writes
// responses to stdout until EOF or the context is canceled.
func (s *Server) Run(ctx context.Context) error {
	scanner := bufio.NewScanner(os.Stdin)
	// 10 MiB line buffer — MCP payloads can get chunky (big tool lists,
	// long tool results); default 64 KiB is too tight.
	buf := make([]byte, 0, 1<<16)
	scanner.Buffer(buf, 10*1024*1024)

	logf := func(format string, a ...any) {
		fmt.Fprintf(os.Stderr, "[mcp] "+format+"\n", a...)
	}

	for scanner.Scan() {
		if err := ctx.Err(); err != nil {
			return err
		}
		line := scanner.Bytes()
		if len(line) == 0 {
			continue
		}
		s.dispatch(ctx, line, logf)
	}
	return scanner.Err()
}

// jsonrpcRequest is the subset of JSON-RPC 2.0 we care about.
type jsonrpcRequest struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"` // may be absent (notification) or null/number/string
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params"`
}

type jsonrpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Data    any    `json:"data,omitempty"`
}

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

const (
	errParse        = -32700
	errInvalidReq   = -32600
	errMethodNotFnd = -32601
	errInvalidParms = -32602
	errInternal     = -32603
)

func (s *Server) writeResponse(resp jsonrpcResponse) {
	resp.JSONRPC = "2.0"
	b, err := json.Marshal(resp)
	if err != nil {
		fmt.Fprintf(os.Stderr, "[mcp] marshal response: %v\n", err)
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.w.Write(b)
	s.w.WriteByte('\n')
	s.w.Flush()
}

func (s *Server) writeError(id json.RawMessage, code int, msg string) {
	s.writeResponse(jsonrpcResponse{ID: id, Error: &jsonrpcError{Code: code, Message: msg}})
}

func (s *Server) dispatch(ctx context.Context, line []byte, logf func(string, ...any)) {
	var req jsonrpcRequest
	if err := json.Unmarshal(line, &req); err != nil {
		s.writeError(nil, errParse, "parse error: "+err.Error())
		return
	}
	if req.JSONRPC != "2.0" {
		s.writeError(req.ID, errInvalidReq, "jsonrpc must be 2.0")
		return
	}

	// Notifications (no id) don't get responses.
	isNotification := len(req.ID) == 0

	switch req.Method {
	case "initialize":
		s.handleInitialize(req)
	case "initialized", "notifications/initialized":
		// Acknowledge; no response for notifications.
		if !isNotification {
			s.writeResponse(jsonrpcResponse{ID: req.ID, Result: struct{}{}})
		}
	case "tools/list":
		s.handleToolsList(req)
	case "tools/call":
		s.handleToolsCall(ctx, req, logf)
	case "ping":
		s.writeResponse(jsonrpcResponse{ID: req.ID, Result: struct{}{}})
	default:
		if !isNotification {
			s.writeError(req.ID, errMethodNotFnd, "method not found: "+req.Method)
		}
	}
}

func (s *Server) handleInitialize(req jsonrpcRequest) {
	result := map[string]any{
		"protocolVersion": ProtocolVersion,
		"serverInfo":      s.info,
		"capabilities": map[string]any{
			"tools": map[string]any{
				// Announce we ship a tool list but don't dynamically change it.
				"listChanged": false,
			},
		},
	}
	s.writeResponse(jsonrpcResponse{ID: req.ID, Result: result})
}

func (s *Server) handleToolsList(req jsonrpcRequest) {
	// MCP expects `{"tools": [...]}`. Strip the Go-only Handler field via the JSON tag `-`.
	s.writeResponse(jsonrpcResponse{ID: req.ID, Result: map[string]any{"tools": s.tools}})
}

func (s *Server) handleToolsCall(ctx context.Context, req jsonrpcRequest, logf func(string, ...any)) {
	var p struct {
		Name      string         `json:"name"`
		Arguments map[string]any `json:"arguments"`
	}
	if err := json.Unmarshal(req.Params, &p); err != nil {
		s.writeError(req.ID, errInvalidParms, "parse params: "+err.Error())
		return
	}
	var tool *Tool
	for i := range s.tools {
		if s.tools[i].Name == p.Name {
			tool = &s.tools[i]
			break
		}
	}
	if tool == nil {
		s.writeError(req.ID, errMethodNotFnd, "unknown tool: "+p.Name)
		return
	}
	result, err := tool.Handler(ctx, p.Arguments)
	if err != nil {
		// MCP convention: tool errors are result payloads with isError=true,
		// not JSON-RPC errors (those are for protocol failures).
		s.writeResponse(jsonrpcResponse{ID: req.ID, Result: map[string]any{
			"isError": true,
			"content": []map[string]any{
				{"type": "text", "text": err.Error()},
			},
		}})
		return
	}
	// Serialize the result to JSON and wrap as a text content block.
	payload, err := json.MarshalIndent(result, "", "  ")
	if err != nil {
		s.writeError(req.ID, errInternal, "marshal result: "+err.Error())
		return
	}
	s.writeResponse(jsonrpcResponse{ID: req.ID, Result: map[string]any{
		"content": []map[string]any{
			{"type": "text", "text": string(payload)},
		},
	}})
}

// asString is a small helper for tool handlers pulling string args.
func asString(args map[string]any, key string) string {
	if v, ok := args[key]; ok {
		if s, ok := v.(string); ok {
			return s
		}
	}
	return ""
}

// ArgString is the public form; exposed for tool handlers defined outside
// this package.
var ArgString = asString

// Discard is a handy helper to silence unused-import complaints.
var _ io.Writer = io.Discard