// Package mcp is the client side of the Model Context Protocol: it starts the
// servers the config names, asks each what tools it has, and offers those
// tools to the agent as if they were Codey's own.
//
// Only the stdio transport is spoken. A server is a child process with
// JSON-RPC 2.0 on its stdin and stdout, one message per line — the same
// arrangement as the language servers in package diag, save that the LSP
// frames its messages with a Content-Length header and MCP with a newline.
// This file is the sibling of diag/lsp.go for that reason, and where the two
// look alike it is because the same thing goes wrong in both.
package mcp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sync"
	"time"
)

// protocolVersion is what the client asks for. supportedVersions are the
// answers it will take: every one of them has the same initialize, tools/list
// and tools/call shapes over stdio, and those are the only methods this client
// uses, so a server answering with an older one is a server it can still
// talk to. A version outside the list is one whose messages may mean
// something else, and the spec's instruction is to disconnect rather than
// guess.
const protocolVersion = "2025-06-18"

var supportedVersions = []string{"2025-06-18", "2025-03-26", "2024-11-05"}

// initTimeout bounds the whole handshake: initialize, initialized and the
// first tools/list. A server started with npx may spend most of it
// downloading itself on first run, which is why it is longer than a call. A
// server that has not answered by then is not going to, and every one of its
// tools would otherwise be advertised to the model and then fail.
const initTimeout = 30 * time.Second

// callTimeout is how long one tools/call may take. It matches what the bash
// tool gives a command by default: an MCP tool is a command the model runs
// with different plumbing, and the model's patience is the same either way.
const callTimeout = 60 * time.Second

// maxMessageBytes is the largest line the client will read. Sixteen megabytes
// is far past any tool result the model could use and short of the point
// where a server that streams a file at us shows up in the harness's memory.
// bufio.Scanner enforces it by failing the read, which ends the session; a
// stream that has lost sync cannot be recovered anyway.
const maxMessageBytes = 16 << 20

// errServerDead is what a call gets from a server that has exited or been
// closed. The server is not restarted: a tool that fails this way once fails
// this way for the rest of the session, and the message says so.
var errServerDead = errors.New("MCP server is not running")

// client is one running server and the connection to it.
type client struct {
	name string

	// stdin is written under writeMu so two tool calls from a parallel batch
	// cannot interleave their lines.
	writeMu sync.Mutex
	stdin   io.WriteCloser
	// kill ends the process and wait reaps it; both are no-ops for the
	// in-process fake the tests use.
	kill func()
	wait func()

	// dead is closed when the connection is gone for any reason; deadErr says
	// why, for the error a tool call then gets.
	dead    chan struct{}
	dieOnce sync.Once
	deadErr error

	mu      sync.Mutex
	nextID  int64
	pending map[int64]chan rpcResponse

	// serverVersion is what the server answered initialize with, kept for
	// error messages: "the server speaks 2024-11-05" is a better complaint
	// than "protocol error".
	serverVersion string
	// hasTools is whether the server declared the tools capability. Without
	// it tools/list is a method the server may not have, and the message it
	// sends back for that is less clear than saying so ourselves.
	hasTools bool
}

// rpcResponse is the part of a response the client uses.
type rpcResponse struct {
	Result json.RawMessage
	Error  *rpcError
}

// rpcError is a JSON-RPC error object.
type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

func (e *rpcError) Error() string { return fmt.Sprintf("%s (%d)", e.Message, e.Code) }

// newClient wires a client over an already-open connection and starts the
// reader. It does not run the handshake: that is the caller's, under the
// caller's deadline. It is the seam the tests use to speak to an in-process
// fake server through pipes.
func newClient(name string, stdin io.WriteCloser, stdout io.Reader, kill, wait func()) *client {
	if kill == nil {
		kill = func() {}
	}
	if wait == nil {
		wait = func() {}
	}
	c := &client{
		name:    name,
		stdin:   stdin,
		kill:    kill,
		wait:    wait,
		dead:    make(chan struct{}),
		pending: map[int64]chan rpcResponse{},
	}
	go c.readLoop(stdout)
	return c
}

// Dead reports whether the connection is gone.
func (c *client) Dead() bool {
	select {
	case <-c.dead:
		return true
	default:
		return false
	}
}

// initialize runs the handshake. The server's protocol version is checked
// against supportedVersions and its capabilities read for tools; the
// initialized notification follows, after which the server may be asked for
// its tools.
func (c *client) initialize(ctx context.Context) error {
	params := map[string]any{
		"protocolVersion": protocolVersion,
		// No roots, no sampling, no elicitation: the client offers the server
		// nothing it can ask of Codey. A server that asks anyway is told the
		// method is not found, which the spec allows it to expect.
		"capabilities": map[string]any{},
		"clientInfo": map[string]any{
			"name":    "codey",
			"version": clientVersion,
		},
	}
	raw, err := c.call(ctx, "initialize", params)
	if err != nil {
		return fmt.Errorf("initialize: %w", err)
	}
	var res struct {
		ProtocolVersion string `json:"protocolVersion"`
		Capabilities    struct {
			Tools json.RawMessage `json:"tools"`
		} `json:"capabilities"`
	}
	if err := json.Unmarshal(raw, &res); err != nil {
		return fmt.Errorf("initialize: server sent a malformed result: %w", err)
	}
	if res.ProtocolVersion == "" {
		return fmt.Errorf("initialize: server did not say which protocol version it speaks")
	}
	if !supported(res.ProtocolVersion) {
		return fmt.Errorf("initialize: server speaks MCP %s, which this client does not (it speaks %s)",
			res.ProtocolVersion, protocolVersion)
	}
	c.mu.Lock()
	c.serverVersion = res.ProtocolVersion
	c.hasTools = len(res.Capabilities.Tools) > 0 && string(res.Capabilities.Tools) != "null"
	c.mu.Unlock()
	if err := c.notify("notifications/initialized", nil); err != nil {
		return fmt.Errorf("initialized: %w", err)
	}
	return nil
}

func supported(v string) bool {
	for _, s := range supportedVersions {
		if s == v {
			return true
		}
	}
	return false
}

// toolInfo is one entry of tools/list as the server describes it.
type toolInfo struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"inputSchema"`
}

// listTools asks for every tool the server has, following the cursor until
// the server stops sending one. The spec allows a server to page; none of the
// common ones do, but one that did would otherwise silently show up with
// only its first page.
func (c *client) listTools(ctx context.Context) ([]toolInfo, error) {
	c.mu.Lock()
	hasTools := c.hasTools
	c.mu.Unlock()
	if !hasTools {
		return nil, fmt.Errorf("server offers no tools (it did not declare the tools capability)")
	}
	var out []toolInfo
	cursor := ""
	for {
		params := map[string]any{}
		if cursor != "" {
			params["cursor"] = cursor
		}
		raw, err := c.call(ctx, "tools/list", params)
		if err != nil {
			return nil, fmt.Errorf("tools/list: %w", err)
		}
		var res struct {
			Tools      []toolInfo `json:"tools"`
			NextCursor string     `json:"nextCursor"`
		}
		if err := json.Unmarshal(raw, &res); err != nil {
			return nil, fmt.Errorf("tools/list: server sent a malformed result: %w", err)
		}
		out = append(out, res.Tools...)
		if res.NextCursor == "" || res.NextCursor == cursor {
			return out, nil
		}
		cursor = res.NextCursor
	}
}

// callTool runs one tool and returns the raw result for the adapter to
// render. The deadline is the caller's or callTimeout, whichever comes first.
func (c *client) callTool(ctx context.Context, name string, args json.RawMessage) (json.RawMessage, error) {
	ctx, cancel := context.WithTimeout(ctx, callTimeout)
	defer cancel()
	if len(args) == 0 {
		args = json.RawMessage("{}")
	}
	return c.call(ctx, "tools/call", map[string]any{"name": name, "arguments": args})
}

// Close ends the server. Stdio MCP has no shutdown request: the client
// closes stdin, which is the server's signal to exit, and gives it a moment
// before the process group is killed. The group rather than the leader,
// because npx and uvx wrappers start the real server as a child and a killed
// wrapper leaves that child holding the pipes.
func (c *client) Close() error {
	if !c.Dead() {
		_ = c.stdin.Close()
		select {
		case <-c.dead:
		case <-time.After(shutdownGrace):
		}
	}
	c.die(errServerDead)
	c.wait()
	return nil
}

// shutdownGrace is how long a server has to exit on its own after stdin is
// closed. Half a second is what the language servers get, for the same
// reason: one that has not managed it by then is not going to.
const shutdownGrace = 500 * time.Millisecond

// die marks the connection gone, fails every request still waiting on it,
// and ends the process. Everything that can end a client funnels through
// here so it happens once.
func (c *client) die(err error) {
	c.dieOnce.Do(func() {
		c.mu.Lock()
		c.deadErr = err
		pending := c.pending
		c.pending = map[int64]chan rpcResponse{}
		c.mu.Unlock()
		for _, ch := range pending {
			close(ch)
		}
		close(c.dead)
		c.kill()
		_ = c.stdin.Close()
	})
}

// deadError is the error a call gets after death: what killed the server,
// wrapped so callers can still test for errServerDead.
func (c *client) deadError() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.deadErr == nil || c.deadErr == errServerDead {
		return fmt.Errorf("%w: %s", errServerDead, c.name)
	}
	return fmt.Errorf("%w: %s: %v", errServerDead, c.name, c.deadErr)
}

// call sends a request and waits for its response, for ctx, or for the
// server to die.
func (c *client) call(ctx context.Context, method string, params any) (json.RawMessage, error) {
	c.mu.Lock()
	if c.Dead() {
		c.mu.Unlock()
		return nil, c.deadError()
	}
	c.nextID++
	id := c.nextID
	ch := make(chan rpcResponse, 1)
	c.pending[id] = ch
	c.mu.Unlock()

	msg := map[string]any{"jsonrpc": "2.0", "id": id, "method": method}
	if params != nil {
		msg["params"] = params
	}
	if err := c.write(msg); err != nil {
		c.mu.Lock()
		delete(c.pending, id)
		c.mu.Unlock()
		return nil, err
	}
	select {
	case resp, ok := <-ch:
		if !ok {
			return nil, c.deadError()
		}
		if resp.Error != nil {
			return nil, resp.Error
		}
		return resp.Result, nil
	case <-ctx.Done():
		c.mu.Lock()
		delete(c.pending, id)
		c.mu.Unlock()
		return nil, ctx.Err()
	}
}

// notify sends a notification: a request with no id and no reply.
func (c *client) notify(method string, params any) error {
	msg := map[string]any{"jsonrpc": "2.0", "method": method}
	if params != nil {
		msg["params"] = params
	}
	return c.write(msg)
}

// reply answers a request the server made of us. It writes from its own
// goroutine rather than from the read loop that called it, for the reason
// diag/lsp.go gives: a server blocked writing its next request while we are
// blocked writing this answer is a deadlock the pipe buffer hides until a
// message outgrows it.
func (c *client) reply(id json.RawMessage, result any, err *rpcError) {
	msg := map[string]any{"jsonrpc": "2.0", "id": id}
	if err != nil {
		msg["error"] = err
	} else {
		msg["result"] = result
	}
	go func() { _ = c.write(msg) }()
}

// write frames msg as one line and writes it in one call. json.Marshal never
// emits a literal newline — it escapes them inside strings — so the
// trailing one is the only one on the line, which is what makes the framing
// hold.
func (c *client) write(msg any) error {
	if c.Dead() {
		return c.deadError()
	}
	body, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	body = append(body, '\n')
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if _, err := c.stdin.Write(body); err != nil {
		// A write that fails is a pipe with nobody on the other end.
		c.die(fmt.Errorf("write: %w", err))
		return c.deadError()
	}
	return nil
}

// incoming is the superset of the three JSON-RPC message shapes; which one
// it is depends on which fields are set.
type incoming struct {
	ID     json.RawMessage `json:"id"`
	Method string          `json:"method"`
	Params json.RawMessage `json:"params"`
	Result json.RawMessage `json:"result"`
	Error  *rpcError       `json:"error"`
}

// readLoop reads lines until the connection fails, dispatching responses to
// their callers. A line that is not JSON ends the session: a server that
// prints its logging to stdout instead of stderr has broken the framing for
// good, and the model is better told its tools are gone than left waiting on
// a reply that will never be recognised.
func (c *client) readLoop(r io.Reader) {
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 64<<10), maxMessageBytes)
	for sc.Scan() {
		line := sc.Bytes()
		if len(bytes.TrimSpace(line)) == 0 {
			continue
		}
		var msg incoming
		if err := json.Unmarshal(line, &msg); err != nil {
			c.die(fmt.Errorf("server sent invalid JSON: %w", err))
			return
		}
		switch {
		case msg.Method == "" && len(msg.ID) > 0:
			c.dispatchResponse(msg)
		case msg.Method != "" && len(msg.ID) > 0 && string(msg.ID) != "null":
			c.handleRequest(msg)
		default:
			// Notifications — list_changed, logging, progress — carry nothing
			// this client acts on. Tools are listed once at startup.
		}
	}
	if err := sc.Err(); err != nil {
		c.die(fmt.Errorf("read: %w", err))
		return
	}
	c.die(errors.New("server exited"))
}

func (c *client) dispatchResponse(msg incoming) {
	var id int64
	if err := json.Unmarshal(msg.ID, &id); err != nil {
		return
	}
	c.mu.Lock()
	ch, ok := c.pending[id]
	delete(c.pending, id)
	c.mu.Unlock()
	if ok {
		ch <- rpcResponse{Result: msg.Result, Error: msg.Error}
	}
}

// handleRequest answers what a server asks of its client. ping gets its
// empty answer because a server that pings and hears nothing may take the
// silence for a dead client and exit. Everything else — roots, sampling,
// elicitation — is a capability the client did not offer in initialize, and
// the spec's answer for that is method-not-found.
func (c *client) handleRequest(msg incoming) {
	switch msg.Method {
	case "ping":
		c.reply(msg.ID, map[string]any{}, nil)
	default:
		c.reply(msg.ID, nil, &rpcError{Code: -32601, Message: "method not found"})
	}
}