// Package mcp implements a small MCP-protocol server over HTTP. The wire
// format is JSON-RPC 2.0 with the MCP method set (initialize, tools/list,
// tools/call). Designed to be mounted under /mcp/rpc on the projax web
// binary; a stdio bridge (see ~/.claude/mcp/projax.sh) lets standard MCP
// clients talk to it transparently.
package mcp

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"
)

// ProtocolVersion is the MCP wire version this server speaks. Clients that
// initialize with a different version are still answered; they're expected
// to negotiate down.
const ProtocolVersion = "2024-11-05"

// JSON-RPC 2.0 error codes (subset).
const (
	codeParseError     = -32700
	codeInvalidRequest = -32600
	codeMethodNotFound = -32601
	codeInvalidParams  = -32602
	codeInternalError  = -32603
)

// Tool describes one callable tool exposed through tools/list and tools/call.
type Tool struct {
	Name        string          // e.g. "list_items"
	Description string          // one-line description for the client
	InputSchema json.RawMessage // JSON-schema object describing the params
	Handler     ToolHandler
}

// ToolHandler runs the actual work for a tool. params is the raw JSON object
// the client supplied as the tool arguments. A non-nil result is wrapped as
// a structured text content block; a non-nil *ToolError surfaces as the
// JSON-RPC error envelope (.error.code / .error.message / .error.data) so
// typed validation payloads round-trip cleanly to MCP clients.
type ToolHandler func(ctx context.Context, params json.RawMessage) (any, *ToolError)

// Server holds the registered tools + the auth token. Mount via Routes() on
// any *http.ServeMux.
type Server struct {
	Name    string
	Version string
	Token   string // Bearer token; empty means "no auth" (tests only)
	Logger  *slog.Logger
	tools   map[string]Tool
}

// New builds an MCP server with no tools registered.
func New(name, version, token string, logger *slog.Logger) *Server {
	if logger == nil {
		logger = slog.Default()
	}
	return &Server{
		Name:    name,
		Version: version,
		Token:   token,
		Logger:  logger,
		tools:   map[string]Tool{},
	}
}

// Register adds a tool. Duplicate names overwrite.
func (s *Server) Register(t Tool) { s.tools[t.Name] = t }

// Routes registers /rpc on the given mux prefix-relative path. Caller mounts
// at /mcp so the resulting URL is /mcp/rpc.
func (s *Server) Routes(mux *http.ServeMux) {
	mux.HandleFunc("POST /rpc", s.handleRPC)
	// Friendly GET for ops smoke-testing — never returns secrets.
	mux.HandleFunc("GET /rpc", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{
			"server":           s.Name,
			"version":          s.Version,
			"protocolVersion":  ProtocolVersion,
			"tools":            s.toolNames(),
			"authRequired":     s.Token != "",
		})
	})
}

func (s *Server) toolNames() []string {
	out := make([]string, 0, len(s.tools))
	for n := range s.tools {
		out = append(out, n)
	}
	return out
}

// jsonRPCReq mirrors the wire shape. Method "notifications/initialized" has
// no id (notification) — we tolerate the missing id field.
type jsonRPCReq struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

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

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

func (s *Server) handleRPC(w http.ResponseWriter, r *http.Request) {
	if s.Token != "" {
		if !s.checkAuth(r) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
	}
	body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 1<<20))
	if err != nil {
		http.Error(w, "request too large", http.StatusRequestEntityTooLarge)
		return
	}
	var req jsonRPCReq
	if err := json.Unmarshal(body, &req); err != nil {
		s.writeErr(w, nil, codeParseError, "invalid JSON: "+err.Error())
		return
	}
	if req.JSONRPC != "2.0" && req.JSONRPC != "" {
		s.writeErr(w, req.ID, codeInvalidRequest, "jsonrpc must be \"2.0\"")
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
	defer cancel()
	switch req.Method {
	case "initialize":
		s.writeOK(w, req.ID, s.initializeResult())
	case "notifications/initialized":
		// Notifications get no response on JSON-RPC, but if the client sent an
		// id we humour it with an empty result.
		if len(req.ID) > 0 {
			s.writeOK(w, req.ID, map[string]any{})
		} else {
			w.WriteHeader(http.StatusNoContent)
		}
	case "tools/list":
		s.writeOK(w, req.ID, s.toolsListResult())
	case "tools/call":
		s.handleToolsCall(ctx, w, req)
	case "ping":
		s.writeOK(w, req.ID, map[string]any{})
	default:
		s.writeErr(w, req.ID, codeMethodNotFound, "unknown method: "+req.Method)
	}
}

func (s *Server) checkAuth(r *http.Request) bool {
	h := r.Header.Get("Authorization")
	if !strings.HasPrefix(h, "Bearer ") {
		return false
	}
	return strings.TrimSpace(h[len("Bearer "):]) == s.Token
}

func (s *Server) initializeResult() map[string]any {
	return map[string]any{
		"protocolVersion": ProtocolVersion,
		"capabilities": map[string]any{
			"tools": map[string]any{},
		},
		"serverInfo": map[string]any{
			"name":    s.Name,
			"version": s.Version,
		},
	}
}

type toolDescriptor struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"inputSchema"`
}

func (s *Server) toolsListResult() map[string]any {
	out := make([]toolDescriptor, 0, len(s.tools))
	for _, t := range s.tools {
		schema := t.InputSchema
		if len(schema) == 0 {
			schema = json.RawMessage(`{"type":"object","properties":{}}`)
		}
		out = append(out, toolDescriptor{
			Name:        t.Name,
			Description: t.Description,
			InputSchema: schema,
		})
	}
	return map[string]any{"tools": out}
}

type toolsCallParams struct {
	Name      string          `json:"name"`
	Arguments json.RawMessage `json:"arguments"`
}

func (s *Server) handleToolsCall(ctx context.Context, w http.ResponseWriter, req jsonRPCReq) {
	var p toolsCallParams
	if err := json.Unmarshal(req.Params, &p); err != nil {
		s.writeErr(w, req.ID, codeInvalidParams, "tools/call params: "+err.Error())
		return
	}
	tool, ok := s.tools[p.Name]
	if !ok {
		s.writeErr(w, req.ID, codeMethodNotFound, "unknown tool: "+p.Name)
		return
	}
	result, te := tool.Handler(ctx, p.Arguments)
	if te != nil {
		// Tool errors surface through the JSON-RPC error envelope with
		// {code, message, data} so MCP clients receive typed payloads on
		// .error.data alongside the human-readable message. Transport-level
		// failures (auth, parse, unknown method) use the same envelope but
		// with the standard JSON-RPC codes only — those callers never set
		// Data.
		s.Logger.Warn("mcp tool error", "tool", p.Name, "code", te.Code, "msg", te.Msg)
		s.writeToolErr(w, req.ID, te)
		return
	}
	payload, err := json.Marshal(result)
	if err != nil {
		s.writeErr(w, req.ID, codeInternalError, "marshal result: "+err.Error())
		return
	}
	s.writeOK(w, req.ID, map[string]any{
		"content": []map[string]any{{
			"type": "text",
			"text": string(payload),
		}},
		"isError": false,
	})
}

func (s *Server) writeOK(w http.ResponseWriter, id json.RawMessage, result any) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(jsonRPCResp{JSONRPC: "2.0", ID: id, Result: result})
}

func (s *Server) writeErr(w http.ResponseWriter, id json.RawMessage, code int, message string) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(jsonRPCResp{JSONRPC: "2.0", ID: id, Error: &rpcError{Code: code, Message: message}})
}

func (s *Server) writeToolErr(w http.ResponseWriter, id json.RawMessage, te *ToolError) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(jsonRPCResp{
		JSONRPC: "2.0", ID: id,
		Error: &rpcError{Code: te.Code, Message: te.Msg, Data: te.Data},
	})
}

// ToolError is the only failure shape a ToolHandler may return. Code is one
// of the JSON-RPC 2.0 codes (codeInternalError / codeInvalidParams / …),
// Msg is the human-readable message, Data is an optional typed payload
// that lands in .error.data verbatim (omitted when nil).
type ToolError struct {
	Code int
	Msg  string
	Data any
}

func (e *ToolError) Error() string { return e.Msg }

// InternalError wraps a plain Go error as a -32603 ToolError. Returns nil
// for a nil input so callers can `return nil, InternalError(err)` without a
// preceding nil-check.
func InternalError(err error) *ToolError {
	if err == nil {
		return nil
	}
	return &ToolError{Code: codeInternalError, Msg: err.Error()}
}

// InvalidParamsError builds a -32602 ToolError for tool-level argument
// failures (missing required field, wrong shape). Slice B's
// ValidationToolError is the typed-payload counterpart for itemwrite
// rejections.
func InvalidParamsError(msg string) *ToolError {
	return &ToolError{Code: codeInvalidParams, Msg: msg}
}