package gateway

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"reflect"
	"slices"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"

	"github.com/txn2/mcp-data-platform/pkg/authevents"
	"github.com/txn2/mcp-data-platform/pkg/connoauth"
	"github.com/txn2/mcp-data-platform/pkg/query"
	"github.com/txn2/mcp-data-platform/pkg/semantic"
	"github.com/txn2/mcp-data-platform/pkg/toolkit"
	"github.com/txn2/mcp-data-platform/pkg/toolkits/gateway/enrichment"
)

// ErrConnectionExists is returned when AddConnection is called with a name
// already live in the toolkit.
var ErrConnectionExists = errors.New("gateway: connection already exists")

// ErrConnectionNotFound is returned when RemoveConnection is called for a
// name that is not present.
var ErrConnectionNotFound = errors.New("gateway: connection not found")

// Log key constants keep structured-slog field names consistent
// across the package — and across packages that compose with this
// one. LogKeyTokenURLHost is exported so the admin handler that
// performs the OAuth code-exchange (which lives outside this package)
// can use the same field name as the token-source's refresh /
// acquire logs, allowing operators to grep one connection's full
// auth lifecycle by `token_url_host=<host>`.
const (
	logKeyConnection = "connection"
	logKeyEndpoint   = "endpoint"
	logKeyError      = "error"

	// Config-map keys used by configToMap and AddConnection's parser.
	cfgKeyEndpoint       = "endpoint"
	cfgKeyAuthMode       = "auth_mode"
	cfgKeyCredential     = "credential"
	cfgKeyConnectionName = "connection_name"
	cfgKeyConnectTimeout = "connect_timeout"
	cfgKeyCallTimeout    = "call_timeout"
	cfgKeyTrustLevel     = "trust_level"

	// LogKeyTokenURLHost is the structured-log field name used when
	// emitting an IdP host. Exported so external packages don't
	// duplicate the literal and risk drift.
	LogKeyTokenURLHost = "token_url_host" // #nosec G101 -- structured-log key name, not a credential

	// LogKeyGrantType is the structured-log field name used when
	// emitting an OAuth grant_type. Exported so the admin handler
	// (which performs the authorization_code exchange) and the
	// gateway token source (which performs refresh / acquire) emit
	// the same field name across the full OAuth lifecycle.
	LogKeyGrantType = "grant_type"
)

// ToolListChangedNotifier publishes a tools/list_changed signal to
// downstream MCP clients. The gateway calls this whenever its
// aggregate tool inventory changes (a connection comes up after
// re-auth, a connection is removed, etc.). Decoupled as an interface
// so the gateway package doesn't depend on pkg/session — the platform
// wires an adapter at startup.
//
// Implementations must be safe for concurrent calls and must NOT
// block on slow downstream subscribers; the toolkit's lock may be held
// when this fires (debounced in practice via notifyToolListChanged).
type ToolListChangedNotifier interface {
	NotifyToolsListChanged(ctx context.Context)
}

// toolListChangeTimer is a tiny *time.Timer holder that lets callers
// reset the debounce window cheaply. Defined as a named type rather
// than a *time.Timer field so future evolution (e.g. swapping in a
// CountingTimer for tests) doesn't cascade across call sites.
type toolListChangeTimer struct {
	t *time.Timer
}

// notifyDebounceWindow batches a flurry of AddTool calls (e.g. when a
// single upstream registers all 12 of its tools at once) into a
// single tools/list_changed notification. 50ms is large enough to
// absorb a typical multi-tool registration burst on a busy node and
// small enough that downstream agents see fresh inventory within
// human-perceivable latency. The MCP SDK's own internal debounce is
// 10ms (mcp/server.go:notificationDelay) — we deliberately run a
// longer window because our path goes through an extra postgres
// LISTEN/NOTIFY round-trip, so collapsing more aggressively reduces
// fan-out cost across the cluster.
const notifyDebounceWindow = 50 * time.Millisecond

// notifyDispatchTimeout bounds the AfterFunc-spawned dispatch goroutine
// so a partitioned downstream (e.g. postgres LISTEN/NOTIFY connection
// hung) cannot leak a goroutine per inventory change. tools/list_changed
// is best-effort — a missed delivery just means downstream agents see
// a stale tool list until the next gateway state change retries the
// notify; an indefinite hang is strictly worse.
const notifyDispatchTimeout = 5 * time.Second

// Toolkit is a gateway that proxies tools from one or more upstream MCP
// servers through the platform's registry, auth, persona, and audit pipeline.
//
// A single Toolkit manages multiple named upstream connections. Each upstream
// tool is re-exposed under a namespaced local name:
// "<connection_name>__<remote_tool_name>". Connections can be added or
// removed at runtime; the MCP server is notified of tool-list changes
// through AddTool / RemoveTools.
//
// Startup is failure-isolated: an unreachable upstream is logged and
// skipped — it does not block platform startup. Other connections remain
// functional.
type Toolkit struct {
	name        string
	defaultName string

	mu               sync.RWMutex
	server           *mcp.Server
	connections      map[string]*upstream
	enrichmentEngine *enrichment.Engine
	connOAuthStore   connoauth.Store
	authEvents       *authevents.Writer

	// listChangedNotifier (when non-nil) is fired (debounced) every
	// time the toolkit's aggregate tool inventory changes — i.e. after
	// every AddTool or RemoveTools call. The platform plugs this into
	// the session broadcaster so SSE long-poll subscribers receive a
	// live notifications/tools/list_changed signal even when the SDK
	// runs in stateless streamable-HTTP mode (where the SDK's own
	// notification path is disabled).
	//
	// Stored as atomic.Pointer so notifyToolListChanged can read it
	// without acquiring t.mu — many of the call sites that need to
	// notify already hold t.mu, and grabbing it twice would deadlock.
	listChangedNotifier  atomic.Pointer[ToolListChangedNotifier]
	pendingNotifyTimer   *toolListChangeTimer
	pendingNotifyTimerMu sync.Mutex

	semanticProvider semantic.Provider
	queryProvider    query.Provider
}

// SetAuthEvents wires the audit-event writer into the toolkit so every
// outbound refresh / revocation / rotation event from the tool-call hot
// path lands in the connection_auth_events history. Connections that
// were dialed before this call pick up the writer on their next
// outbound request (the round-tripper builds a fresh connoauth.Source
// per call from the toolkit's current store + events pair).
func (t *Toolkit) SetAuthEvents(w *authevents.Writer) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.authEvents = w
}

// ConnOAuthStore returns the unified OAuth token store wired into
// this toolkit, or nil when none has been wired yet. Mirrors
// apigateway.Toolkit.ConnOAuthStore for cross-toolkit parity.
func (t *Toolkit) ConnOAuthStore() connoauth.Store {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.connOAuthStore
}

// SetConnOAuthStore wires the unified OAuth token store into the
// gateway. Required for authorization_code grants to survive process
// restarts and to refresh tokens silently between operator interactions.
//
// When called after AddConnection, any authorization_code placeholder
// connections (those that were registered as "awaiting reauth" because
// the store wasn't yet wired) are automatically retried so persisted
// tokens are picked up without requiring a manual refresh. This makes
// startup wiring order-independent.
//
// Concurrency: the toolkit lock is HELD ONLY for the snapshot of
// placeholders to retry and again for the final pointer-swap on each
// success. The discover() network I/O happens with the lock RELEASED
// so concurrent Status / ListConnections / Tools / AddConnection
// callers don't block on potentially-slow upstream dials. The
// placeholder remains in the connections map throughout, so the
// admin UI keeps showing "Connect" while a retry is in flight.
func (t *Toolkit) SetConnOAuthStore(s connoauth.Store) {
	type pending struct {
		name string
		cfg  Config
	}
	t.mu.Lock()
	t.connOAuthStore = s
	var retry []pending
	for name, u := range t.connections {
		if u.client != nil ||
			u.config.AuthMode != AuthModeOAuth ||
			u.config.OAuth.Grant != OAuthGrantAuthorizationCode {
			continue
		}
		retry = append(retry, pending{name: name, cfg: u.config})
	}
	t.mu.Unlock()

	// Retries run concurrently — each placeholder's discover() is bounded
	// by cfg.ConnectTimeout, so a single sick upstream no longer blocks
	// the others from coming up. installLiveConnection is concurrency-
	// safe (it re-acquires the toolkit lock for its atomic pointer swap).
	var wg sync.WaitGroup
	for _, p := range retry {
		wg.Add(1)
		go func(p pending) {
			defer wg.Done()
			ctx, cancel := dialContext(p.cfg)
			defer cancel()
			res, err := t.discoverFor(ctx, p.name, p.cfg)
			if err != nil {
				slog.Warn("gateway: retry placeholder after token store wired",
					logKeyConnection, p.cfg.ConnectionName,
					logKeyError, err)
				// Update the placeholder's lastError so Status() reflects
				// the most recent rejection reason. Without this, the UI
				// would still show the original placeholder error from
				// AddConnection time, even after a retry surfaced a new
				// upstream rejection.
				t.recordPlaceholderError(p.name, err.Error())
				return
			}
			t.installLiveConnection(p.name, p.cfg, res)
		}(p)
	}
	wg.Wait()
}

// discoverResult bundles the products of a discover run that the
// installer needs to thread into the connection's live state — the
// upstream client, its tool catalog, and (for cc grants) the
// in-memory token provider that Status/Reacquire later read.
type discoverResult struct {
	client     *upstreamClient
	tools      []*mcp.Tool
	ccProvider *clientCredentialsTokenProvider
}

// discoverFor builds the tokenProvider for the connection and runs
// the dial + listTools pair. Centralized so every call site that
// materializes an upstream client (initial registration, post-Connect
// rebuild, SetConnOAuthStore retry) picks the same provider strategy.
//
// authorization_code → connoauthTokenProvider, which re-reads the
// toolkit's fields and the persisted token row on every outbound
// request. client_credentials → clientCredentialsTokenProvider, a
// mutex-guarded in-memory cache (cc has no rotation contract, so
// caching is correct and the provider is the source of truth for
// Status / Reacquire).
func (t *Toolkit) discoverFor(ctx context.Context, name string, cfg Config) (*discoverResult, error) {
	provider, ccProvider := t.tokenProviderFor(name, cfg.OAuth)
	dialCtx, cancel := context.WithTimeout(ctx, cfg.ConnectTimeout)
	defer cancel()
	client, err := dial(dialCtx, cfg, dialDeps{TokenProvider: provider})
	if err != nil {
		return nil, err
	}
	remoteTools, err := client.listTools(dialCtx)
	if err != nil {
		_ = client.close()
		return nil, fmt.Errorf("list tools: %w", err)
	}
	return &discoverResult{client: client, tools: remoteTools, ccProvider: ccProvider}, nil
}

// tokenProviderFor returns the correct tokenProvider for the
// connection's OAuth grant, or nil for non-OAuth connections.
// Returns nil for authorization_code when no connoauth.Store has been
// wired yet so dial() can surface a clear "OAuth store not wired"
// error to the operator rather than constructing a placeholder
// provider whose first Token call would fail vaguely.
//
// The second return is the cc-specific provider when grant is
// client_credentials, so the toolkit can stash it on the upstream
// and route Status / Reacquire through it. Returns nil otherwise.
func (t *Toolkit) tokenProviderFor(name string, cfg OAuthConfig) (tokenProvider, *clientCredentialsTokenProvider) {
	switch cfg.Grant {
	case OAuthGrantClientCredentials:
		cc := newClientCredentialsTokenProvider(cfg)
		return cc, cc
	case OAuthGrantAuthorizationCode:
		t.mu.RLock()
		hasStore := t.connOAuthStore != nil
		t.mu.RUnlock()
		if !hasStore {
			return nil, nil
		}
		return connoauthTokenProvider{tk: t, connection: name, cfg: cfg}, nil
	}
	return nil, nil
}

// dialContext returns a context bounded by cfg.ConnectTimeout (or the
// package default when cfg.ConnectTimeout is zero/negative). Used for
// every discover() call so a hung MCP-protocol handshake against an
// unhealthy upstream cannot block the caller (admin OAuth callback,
// SetTokenStore retry, AddConnection) for longer than the operator-
// configured timeout.
func dialContext(cfg Config) (context.Context, context.CancelFunc) {
	timeout := cfg.ConnectTimeout
	if timeout <= 0 {
		timeout = DefaultConnectTimeout
	}
	return context.WithTimeout(context.Background(), timeout)
}

// recordPlaceholderError updates a placeholder's lastError so Status()
// reflects the most recent dial / discover failure. No-op when the
// connection no longer exists (concurrent removal) or has already been
// promoted to a live client.
//
// Takes the message as a string rather than an error so the call site
// owns formatting decisions (e.g. wrapping, context prefixing); the
// helper itself stores raw display text without further interpretation.
func (t *Toolkit) recordPlaceholderError(name, msg string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	u, ok := t.connections[name]
	if !ok || u.client != nil {
		return
	}
	u.lastError = msg
}

// installLiveConnection promotes a placeholder to a live upstream by
// swapping the connection-map pointer atomically under the lock. If
// the connection was removed concurrently (or already replaced by
// another retry), the freshly-dialed client is closed and the result
// discarded — last-writer-wins for the rare race.
func (t *Toolkit) installLiveConnection(name string, cfg Config, res *discoverResult) {
	t.mu.Lock()
	existing, ok := t.connections[name]
	if !ok || existing.client != nil {
		t.mu.Unlock()
		// Connection was removed or already promoted; drop our work.
		if cerr := res.client.close(); cerr != nil {
			slog.Warn("gateway: discarded retry client close error",
				logKeyConnection, cfg.ConnectionName,
				logKeyError, cerr)
		}
		return
	}
	u := &upstream{
		name:       name,
		config:     cfg,
		client:     res.client,
		tools:      res.tools,
		toolNames:  makeLocalNames(cfg.ConnectionName, res.tools),
		desc:       "Gateway to " + cfg.Endpoint,
		ccProvider: res.ccProvider,
	}
	t.connections[name] = u
	if t.server != nil {
		t.addToolsToServerLocked(u)
	}
	t.mu.Unlock()
	slog.Info("gateway: upstream connected",
		logKeyConnection, cfg.ConnectionName,
		logKeyEndpoint, cfg.Endpoint,
		"tools", len(u.toolNames))
}

// SetEnrichmentEngine wires a cross-enrichment engine into this gateway.
// When set, every successful forwarded tool result is run through the
// engine before being returned to the client. Safe to call before or
// after RegisterTools — handlers fetch the current engine on each call.
func (t *Toolkit) SetEnrichmentEngine(e *enrichment.Engine) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.enrichmentEngine = e
}

// SetToolListChangedNotifier wires the tools/list_changed publisher.
// Pass nil to disable (default). Both bare-nil and typed-nil interface
// values are detected via isNilNotifier and treated as "disable"; a
// `var p *somePtrAdapter; SetToolListChangedNotifier(p)` is safe and
// behaves identically to the bare-nil call.
// The notifier is fired (debounced) on every aggregate tool-inventory
// change that registers or removes at least one tool — RegisterTools,
// addParsed successes, SetTokenStore retry promotions, and
// RemoveConnection. Connections that resolve to zero tools (placeholder
// upstreams that never finished discovery, RemoveConnection on a
// connection with no live tools) are silently skipped to keep the
// downstream notification budget proportional to operator-visible
// inventory state.
//
// Safe to call before or after the toolkit has connections. The
// production wiring path: main.go calls WireGatewayBroadcaster AFTER
// p.Start() runs RegisterTools, so the notifier is set after the
// connections in the initial config are already loaded — those
// initial RegisterTools calls fire BEFORE the notifier is wired, so
// they do not produce a notification. A caller that wires the
// notifier earlier (e.g., a custom integration test, or a hot-reload
// path that reinstalls connections) WILL see notifications for
// every connection coming up.
func (t *Toolkit) SetToolListChangedNotifier(n ToolListChangedNotifier) {
	if isNilNotifier(n) {
		t.listChangedNotifier.Store(nil)
		return
	}
	t.listChangedNotifier.Store(&n)
}

// isNilNotifier reports whether n is nil OR a typed-nil interface
// value (e.g., a nil pointer to a struct that implements
// ToolListChangedNotifier). The bare `n == nil` check misses
// typed-nil — a `var p *somePtrAdapter; SetToolListChangedNotifier(p)`
// would otherwise install a non-nil interface wrapping a nil pointer
// and panic with a nil-receiver dereference at fire time.
//
// The Kind list covers the typed-nil cases where dispatching a method
// would actually panic on the nil receiver: Ptr (most common — e.g.,
// `var p *somePtrAdapter`), Interface (interface holding another
// interface), and Func (named func types like the test's
// `notifierFunc`, where the receiver IS the function value and a nil
// func panics on call). Map/Chan/Slice are deliberately omitted —
// although named map/chan/slice types CAN define methods and satisfy
// ToolListChangedNotifier, calling a method on a nil map/chan/slice
// receiver does not panic on the receiver itself (the receiver is
// the value, not a dereference target). Adding them would be a no-op
// guard.
func isNilNotifier(n ToolListChangedNotifier) bool {
	if n == nil {
		return true
	}
	v := reflect.ValueOf(n)
	switch v.Kind() {
	case reflect.Pointer, reflect.Interface, reflect.Func:
		return v.IsNil()
	default:
		return false
	}
}

// notifyToolListChanged schedules a debounced notification through the
// configured notifier. Multiple calls within notifyDebounceWindow
// collapse into a single fire — matches the SDK's internal batching
// behavior and avoids fan-out floods when an upstream registers many
// tools in rapid succession.
//
// Safe for concurrent calls and safe to call while t.mu is already
// held by the caller — the notifier slot is read atomically (no
// t.mu) and the timer firing happens on a separate goroutine via
// time.AfterFunc.
func (t *Toolkit) notifyToolListChanged() {
	if t.listChangedNotifier.Load() == nil {
		return
	}
	t.pendingNotifyTimerMu.Lock()
	defer t.pendingNotifyTimerMu.Unlock()
	if t.pendingNotifyTimer == nil {
		t.pendingNotifyTimer = &toolListChangeTimer{
			// Read the notifier from the atomic slot at FIRE time, not
			// at scheduling time. SetToolListChangedNotifier may swap
			// the notifier between scheduling and firing — we must
			// dispatch to the currently-installed notifier or the
			// callback could call into a stale reference.
			t: time.AfterFunc(notifyDebounceWindow, func() {
				// Clear the timer slot before invoking the notifier so
				// concurrent callers arriving DURING dispatch construct
				// a fresh AfterFunc. Per Go's Timer.Reset documentation:
				// Reset on an AfterFunc-created timer reschedules the
				// callback. Without this clear, a concurrent
				// notifyToolListChanged would call Reset on the
				// already-fired timer and queue a SECOND callback while
				// this one is still running — at most one extra
				// dispatch (bounded), but it weakens the debounce
				// contract. The narrow race that remains: a caller
				// arriving between timer expiry and the AfterFunc body
				// taking the mutex can still Reset; that path produces
				// at most ≤2 dispatches per burst, which the burst test
				// asserts as the loose bound.
				t.pendingNotifyTimerMu.Lock()
				t.pendingNotifyTimer = nil
				t.pendingNotifyTimerMu.Unlock()
				if p := t.listChangedNotifier.Load(); p != nil {
					// Bound the dispatch context so a partitioned
					// downstream (postgres LISTEN connection hung,
					// remote receiver blocked) cannot leak a goroutine
					// per inventory change.
					ctx, cancel := context.WithTimeout(context.Background(), notifyDispatchTimeout)
					(*p).NotifyToolsListChanged(ctx)
					cancel()
				}
			}),
		}
		return
	}
	t.pendingNotifyTimer.t.Reset(notifyDebounceWindow)
}

// upstream tracks a single live connection to a remote MCP server.
type upstream struct {
	name      string
	config    Config
	client    *upstreamClient
	tools     []*mcp.Tool // cached definitions from discovery
	toolNames []string
	desc      string
	// ccProvider is the live in-memory client_credentials token
	// provider for this connection. Non-nil ONLY for live oauth
	// client_credentials upstreams; nil for authorization_code (which
	// reads from the unified connoauth.Store on every call), bearer,
	// api_key, and none modes. Held here so Status() can synthesize a
	// truthful token-acquired snapshot and ReacquireOAuthToken can
	// force a fresh fetch — neither operation can go through the
	// connoauth.Store path for cc, since cc tokens are intentionally
	// in-memory only.
	ccProvider *clientCredentialsTokenProvider
	// lastError captures the most recent dial / discover error when the
	// connection is in placeholder state (client == nil). Surfaced via
	// Status().OAuth.LastError so the admin UI can show the operator the
	// actual upstream rejection reason — instead of leaving them to guess
	// at a silent "awaiting reauth" warning.
	//
	// Lifecycle: set by addParsedConnection at placeholder creation time
	// and refreshed by recordPlaceholderError when SetTokenStore's retry
	// path produces a new failure. On successful promotion to live,
	// installLiveConnection REPLACES the entire upstream struct with a
	// fresh value (so lastError starts at the zero string ""); there is
	// no in-place clear. Live connections never set or read this field.
	lastError string

	// claiming is set briefly while addParsedConnection is performing
	// network I/O (discover) WITHOUT holding the toolkit's mutex. The
	// entry exists in t.connections so concurrent AddConnection calls
	// for the same name see ErrConnectionExists, but other readers
	// (Status, ListConnections, Tools) treat this as a placeholder.
	// Cleared when addParsedConnection completes (entry replaced with
	// either a live connection or a real placeholder).
	claiming bool
}

// New builds a Toolkit with the given default connection name and no
// initial connections.
func New(defaultName string) *Toolkit {
	return &Toolkit{
		name:        defaultName,
		defaultName: defaultName,
		connections: make(map[string]*upstream),
	}
}

// NewMulti builds a Toolkit and pre-loads the given parsed connection
// configs. Unreachable upstreams are logged and skipped so platform startup
// is never blocked.
func NewMulti(cfg MultiConfig) *Toolkit {
	name := cfg.DefaultName
	if name == "" {
		name = Kind
	}
	t := New(name)
	t.defaultName = cfg.DefaultName // may be empty; only "name" always set
	for instanceName, c := range cfg.Instances {
		if c.ConnectionName == "" {
			c.ConnectionName = instanceName
		}
		if err := t.addParsedConnection(instanceName, c); err != nil {
			slog.Warn("gateway: initial connection failed",
				logKeyConnection, instanceName, logKeyError, err)
		}
	}
	return t
}

// Kind returns the toolkit kind.
func (*Toolkit) Kind() string {
	return Kind
}

// Name returns the toolkit instance name.
func (t *Toolkit) Name() string {
	return t.name
}

// Connection returns the default connection name used in audit logs when a
// request does not carry one. Empty if no default is configured.
func (t *Toolkit) Connection() string {
	return t.defaultName
}

// Tools returns the aggregate, sorted set of namespaced local tool names
// across all live connections.
func (t *Toolkit) Tools() []string {
	t.mu.RLock()
	defer t.mu.RUnlock()
	var out []string
	for _, u := range t.connections {
		out = append(out, u.toolNames...)
	}
	sort.Strings(out)
	return out
}

// ConnectionForTool maps a namespaced local tool name (e.g.
// "vendor__list_contacts") back to its source connection name. Used by
// the platform's audit middleware to attribute proxied tool calls to
// the specific upstream connection they were routed to.
func (t *Toolkit) ConnectionForTool(toolName string) string {
	t.mu.RLock()
	defer t.mu.RUnlock()
	for _, u := range t.connections {
		if slices.Contains(u.toolNames, toolName) {
			return u.config.ConnectionName
		}
	}
	return ""
}

// RegisterTools captures the server reference and registers every tool from
// every already-loaded connection. Must be called exactly once, after the
// toolkit is registered in the platform registry.
func (t *Toolkit) RegisterTools(s *mcp.Server) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.server = s
	for _, u := range t.connections {
		if u.client == nil || len(u.tools) == 0 {
			continue
		}
		t.addToolsToServerLocked(u)
	}
}

// SetSemanticProvider stores the semantic provider (not consumed directly in v1).
func (t *Toolkit) SetSemanticProvider(provider semantic.Provider) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.semanticProvider = provider
}

// SetQueryProvider stores the query provider (not consumed directly in v1).
func (t *Toolkit) SetQueryProvider(provider query.Provider) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.queryProvider = provider
}

// AddConnection parses the raw config, dials the upstream, discovers its
// tools, and registers them on the server (if one is already set). The
// admin layer's hotAddConnection helper logs any returned error as a
// structured warning.
func (t *Toolkit) AddConnection(name string, config map[string]any) error {
	cfg, err := ParseConfig(config)
	if err != nil {
		return err
	}
	if cfg.ConnectionName == "" {
		cfg.ConnectionName = name
	}
	return t.addParsedConnection(name, cfg)
}

// addParsedConnection runs the dial + registration. The toolkit's mutex
// is held ONLY for the brief slot-claim and slot-install steps; the
// network I/O (dial + ListTools) runs WITHOUT the lock so other
// operations (Status, ListConnections, Tools, RemoveConnection of a
// different name) proceed in parallel.
//
// A "claiming" sentinel entry is inserted in the connections map before
// the dial so concurrent AddConnection calls for the same name see
// ErrConnectionExists. The claim sentinel is identified by POINTER
// IDENTITY (not just the claiming bool) so that a Remove + Add cycle
// during a slow dial cannot install our result into a later caller's
// slot — see installDialResult.
func (t *Toolkit) addParsedConnection(name string, cfg Config) error {
	claim, err := t.claimConnectionSlot(name, cfg)
	if err != nil {
		return err
	}
	// Network I/O happens here, OUTSIDE the lock. A slow or hung
	// upstream blocks only this call — other toolkit operations
	// (Status polls, listing connections, calling tools on other
	// connections) proceed without contention.
	ctx, cancel := dialContext(cfg)
	defer cancel()
	res, dialErr := t.discoverFor(ctx, name, cfg)
	dr := dialResult{
		claim:   claim,
		name:    name,
		cfg:     cfg,
		dialErr: dialErr,
	}
	if res != nil {
		dr.client = res.client
		dr.tools = res.tools
		dr.ccProvider = res.ccProvider
	}
	return t.installDialResult(dr)
}

// claimConnectionSlot inserts a sentinel "claiming" entry under the
// lock so concurrent AddConnection calls for the same name see
// ErrConnectionExists. Returns ErrConnectionExists if a real or
// claiming entry already occupies the slot. The returned *upstream is
// the inserted sentinel — callers MUST pass it to installDialResult
// so the install step can verify the slot still holds OUR claim
// (and not a fresh claim from a concurrent Remove + Add cycle).
func (t *Toolkit) claimConnectionSlot(name string, cfg Config) (*upstream, error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if _, exists := t.connections[name]; exists {
		return nil, fmt.Errorf("gateway: %s: %w", name, ErrConnectionExists)
	}
	claim := &upstream{
		name:     name,
		config:   cfg,
		desc:     "Connecting to " + cfg.Endpoint,
		claiming: true,
	}
	t.connections[name] = claim
	return claim, nil
}

// installDialResult finishes a claim by replacing the "claiming"
// sentinel with the dial's result. Holds the lock only briefly to
// mutate the map and register tools.
//
// Identity check: the install only proceeds if the slot still holds
// the SAME *upstream pointer that claimConnectionSlot inserted. A
// Remove + Add cycle during our dial replaces the entry with a fresh
// claim sentinel that has the same `claiming` bool but a different
// pointer; without this check, our (potentially stale) dial result
// would silently overwrite the second caller's claim.
// dialResult bundles the inputs to installDialResult so the function
// signature stays under revive's argument-limit ceiling without losing
// any of the data the install step needs.
type dialResult struct {
	claim      *upstream
	name       string
	cfg        Config
	client     *upstreamClient
	tools      []*mcp.Tool
	ccProvider *clientCredentialsTokenProvider
	dialErr    error
}

func (t *Toolkit) installDialResult(r dialResult) error {
	claim, name, cfg, client, tools, dialErr := r.claim, r.name, r.cfg, r.client, r.tools, r.dialErr
	t.mu.Lock()
	if t.connections[name] != claim {
		// Slot was removed (and possibly re-claimed) during our dial —
		// discard the result. The other caller's dial owns the slot.
		t.mu.Unlock()
		if client != nil {
			_ = client.close()
		}
		return nil
	}

	if dialErr != nil {
		// authorization_code connections without a stored token are
		// expected to fail discovery — keep a placeholder so the admin
		// status surface renders a Connect button. Other failure modes
		// drop the slot and bubble the error so retries can succeed.
		if cfg.AuthMode == AuthModeOAuth && cfg.OAuth.Grant == OAuthGrantAuthorizationCode {
			t.connections[name] = &upstream{
				name:      name,
				config:    cfg,
				desc:      "Awaiting OAuth authorization",
				lastError: dialErr.Error(),
			}
			t.mu.Unlock()
			slog.Warn("gateway: oauth authorization_code connection awaiting reauth",
				logKeyConnection, cfg.ConnectionName,
				logKeyEndpoint, cfg.Endpoint,
				logKeyError, dialErr)
			return nil
		}
		delete(t.connections, name)
		t.mu.Unlock()
		slog.Warn("gateway: upstream unavailable",
			logKeyConnection, cfg.ConnectionName,
			logKeyEndpoint, cfg.Endpoint,
			logKeyError, dialErr)
		return dialErr
	}

	u := &upstream{
		name:       name,
		config:     cfg,
		client:     client,
		tools:      tools,
		toolNames:  makeLocalNames(cfg.ConnectionName, tools),
		desc:       "Gateway to " + cfg.Endpoint,
		ccProvider: r.ccProvider,
	}
	t.connections[name] = u
	if t.server != nil {
		t.addToolsToServerLocked(u)
	}
	t.mu.Unlock()
	slog.Info("gateway: upstream connected",
		logKeyConnection, cfg.ConnectionName,
		logKeyEndpoint, cfg.Endpoint,
		"tools", len(u.toolNames))
	return nil
}

// RemoveConnection unregisters a connection's tools from the MCP server,
// closes its upstream session, and removes it from the toolkit.
//
// The toolkit's mutex is held only for the brief map mutation +
// server.RemoveTools call. The actual client.close() — which performs
// network I/O (DELETE the MCP session against the upstream) — runs
// WITHOUT the lock so a slow upstream cannot block other toolkit
// operations.
func (t *Toolkit) RemoveConnection(name string) error {
	t.mu.Lock()
	u, ok := t.connections[name]
	if !ok {
		t.mu.Unlock()
		return fmt.Errorf(errFmtConnection, name, ErrConnectionNotFound)
	}
	notify := false
	if t.server != nil && len(u.toolNames) > 0 {
		t.server.RemoveTools(u.toolNames...)
		notify = true
	}
	delete(t.connections, name)
	client := u.client
	connectionName := u.config.ConnectionName
	t.mu.Unlock()
	if notify {
		t.notifyToolListChanged()
	}

	if client != nil {
		if err := client.close(); err != nil {
			slog.Warn("gateway: error closing upstream session",
				logKeyConnection, connectionName,
				logKeyError, err)
		}
	}
	return nil
}

// HasConnection reports whether a connection with the given name is live.
func (t *Toolkit) HasConnection(name string) bool {
	t.mu.RLock()
	defer t.mu.RUnlock()
	_, ok := t.connections[name]
	return ok
}

// ConnectionStatus is a per-connection health snapshot exposed by the
// admin status endpoint.
type ConnectionStatus struct {
	Name     string                 `json:"name"`
	Healthy  bool                   `json:"healthy"`
	AuthMode string                 `json:"auth_mode"`
	Tools    []string               `json:"tools,omitempty"`
	OAuth    *connoauth.OAuthStatus `json:"oauth,omitempty"`
}

// Status returns a status snapshot for the named connection. Returns
// nil when the connection is not registered.
//
// For OAuth connections the OAuth field is populated by reading the
// persisted token row via connoauth.Source.Status — identical to
// what the general
// /api/v1/admin/connections/{kind}/{name}/oauth-status endpoint
// returns, so the gateway-specific status panel and the generic
// OAuthStatusCard never disagree about token state.
//
// ctx is propagated to the persistence read so a client disconnect
// during admin polling cancels the DB query. The toolkit's read lock
// is released BEFORE the DB roundtrip so the storage call cannot
// stall concurrent writers (AddConnection, RemoveConnection, etc.).
func (t *Toolkit) Status(ctx context.Context, name string) *ConnectionStatus {
	// Snapshot under the lock — fields used below must be a coherent
	// view, but the slow part (the DB call) runs after the unlock.
	t.mu.RLock()
	u, ok := t.connections[name]
	if !ok {
		t.mu.RUnlock()
		return nil
	}
	healthy := u.client != nil
	cfg := u.config
	tools := append([]string(nil), u.toolNames...)
	lastError := u.lastError
	ccProvider := u.ccProvider
	store := t.connOAuthStore
	events := t.authEvents
	t.mu.RUnlock()

	cs := &ConnectionStatus{
		Name:     name,
		Healthy:  healthy,
		AuthMode: cfg.AuthMode,
		Tools:    tools,
	}
	if cfg.AuthMode == AuthModeOAuth {
		var s connoauth.OAuthStatus
		switch {
		case cfg.OAuth.Grant == OAuthGrantClientCredentials && ccProvider != nil:
			// cc tokens live only in memory by design — the
			// connoauth.Store has no row to read. Synthesize the
			// status from the live provider's cache.
			s = ccProvider.Status()
		case cfg.OAuth.Grant == OAuthGrantClientCredentials:
			// cc connection with no live provider (dial failed
			// earlier). Report Configured but no token yet.
			s = connoauth.OAuthStatus{
				Configured: true,
				Grant:      OAuthGrantClientCredentials,
				TokenURL:   cfg.OAuth.TokenURL,
			}
		default:
			s = authzCodeStatusSnapshot(ctx, store, events, name, cfg.OAuth)
		}
		// Surface the connection's last dial error so operators see
		// WHY the upstream rejected — distinct from a pure
		// "needs reauth" caused by an absent token row.
		if lastError != "" {
			s.LastError = lastError
		}
		cs.OAuth = &s
	}
	return cs
}

// authzCodeStatusSnapshot returns the persisted authorization_code
// token snapshot for a connection. When no OAuth store is wired the
// snapshot reports the configured-but-unwired state; otherwise it
// reads through the unified connoauth.Source.Status path (same shape
// as the generic admin oauth-status endpoint). Only valid for
// authorization_code connections — cc never persists a row and so
// must take the in-memory provider's Status instead.
func authzCodeStatusSnapshot(ctx context.Context, store connoauth.Store, events *authevents.Writer,
	name string, cfg OAuthConfig,
) connoauth.OAuthStatus {
	if store == nil {
		return connoauth.OAuthStatus{
			Configured:  true,
			NeedsReauth: true,
			TokenURL:    cfg.TokenURL,
			Grant:       cfg.Grant,
		}
	}
	return connoauthSourceFor(store, events, name, cfg).Status(ctx)
}

// ReacquireOAuthToken forces a fresh token mint for the named
// connection. For authorization_code, that's a refresh-token exchange
// against the IdP via connoauth.Source.Reacquire (the cached row in
// connection_oauth_tokens is rolled forward). For client_credentials,
// the in-memory token is dropped and the grant is re-run against the
// IdP. Returns an error if the connection is missing or not configured
// for OAuth.
func (t *Toolkit) ReacquireOAuthToken(ctx context.Context, name string) error {
	t.mu.RLock()
	u, ok := t.connections[name]
	store := t.connOAuthStore
	events := t.authEvents
	t.mu.RUnlock()
	if !ok {
		return fmt.Errorf(errFmtConnection, name, ErrConnectionNotFound)
	}
	if u.config.AuthMode != AuthModeOAuth {
		return fmt.Errorf("gateway: %s: not configured for OAuth", name)
	}
	switch u.config.OAuth.Grant {
	case OAuthGrantClientCredentials:
		if u.ccProvider == nil {
			return fmt.Errorf("gateway: %s: client_credentials provider not initialized (connection not yet dialed)", name)
		}
		if err := u.ccProvider.Reacquire(ctx); err != nil {
			return fmt.Errorf("gateway: %s: reacquire client_credentials: %w", name, err)
		}
		return nil
	case OAuthGrantAuthorizationCode:
		if store == nil {
			return fmt.Errorf("gateway: %s: oauth token store not wired", name)
		}
		src := connoauthSourceFor(store, events, name, u.config.OAuth)
		if err := src.Reacquire(ctx); err != nil {
			return fmt.Errorf("gateway: %s: reacquire oauth token: %w", name, err)
		}
		return nil
	default:
		return fmt.Errorf("gateway: %s: unknown OAuth grant %q", name, u.config.OAuth.Grant)
	}
}

// IngestOAuthTokenInput is the parameter set for IngestOAuthToken.
// Defined as a struct to keep the public method's argument list under
// the project's revive limit.
//
// RefreshExpiresIn is optional: only some IdPs (Keycloak) emit it on
// their token-endpoint response. Zero leaves the refresh deadline
// unknown — the admin UI renders an em dash.
type IngestOAuthTokenInput struct {
	Name             string
	AccessToken      string
	RefreshToken     string
	ExpiresIn        int
	RefreshExpiresIn int
	Scope            string
	AuthenticatedBy  string
}

// IngestOAuthToken stores tokens obtained from an authorization_code
// callback into the unified connection_oauth_tokens row AND rebuilds
// the upstream so the previously "needs reauth" connection becomes
// live with its discovered tools registered on the MCP server.
//
// In production this is reached via the legacy admin
// gateway_oauth_handler — the unified connection OAuth handler in
// pkg/admin writes the same row directly through connoauth.Store.Set
// and then calls OAuthKindHandler.AfterConnect for the rebuild.
func (t *Toolkit) IngestOAuthToken(ctx context.Context, in IngestOAuthTokenInput) error {
	slog.Info("gateway: IngestOAuthToken — start",
		logKeyConnection, in.Name,
		"authenticated_by", in.AuthenticatedBy,
		"access_token_len", len(in.AccessToken),
		"refresh_token_present", in.RefreshToken != "",
		"expires_in", in.ExpiresIn,
		"scope", in.Scope)
	if in.AccessToken == "" {
		return fmt.Errorf("gateway: %s: ingest token: access_token is required", in.Name)
	}
	t.mu.RLock()
	u, ok := t.connections[in.Name]
	store := t.connOAuthStore
	t.mu.RUnlock()
	if !ok {
		slog.Warn("gateway: IngestOAuthToken — connection not found",
			logKeyConnection, in.Name)
		return fmt.Errorf(errFmtConnection, in.Name, ErrConnectionNotFound)
	}
	cfg := u.config
	if store == nil {
		return fmt.Errorf("gateway: %s: ingest token: oauth token store not wired", in.Name)
	}
	now := time.Now()
	persisted := connoauth.PersistedToken{
		Key:              connoauth.Key{Kind: connoauth.KindMCP, Name: in.Name},
		AccessToken:      in.AccessToken,
		RefreshToken:     in.RefreshToken,
		ExpiresAt:        expiresAtFromSeconds(now, in.ExpiresIn),
		RefreshExpiresAt: expiresAtFromSeconds(now, in.RefreshExpiresIn),
		Scope:            in.Scope,
		AuthenticatedBy:  in.AuthenticatedBy,
		AuthenticatedAt:  now,
	}
	if err := store.Set(ctx, persisted); err != nil {
		slog.Error("gateway: IngestOAuthToken — persist failed",
			logKeyConnection, in.Name, logKeyError, err)
		return fmt.Errorf("gateway: %s: ingest token: %w", in.Name, err)
	}
	slog.Debug("gateway: IngestOAuthToken — tokens persisted, rebuilding connection",
		logKeyConnection, in.Name)

	// Replace the connection with a fresh one so it gets re-dialed
	// using the now-valid tokens and its tools registered on the live
	// server. RemoveConnection is a no-op when the connection is
	// already gone (a concurrent admin operation, etc.) — that's why
	// the not-found error is intentionally swallowed.
	if err := t.RemoveConnection(in.Name); err != nil && !errors.Is(err, ErrConnectionNotFound) {
		slog.Error("gateway: IngestOAuthToken — RemoveConnection failed",
			logKeyConnection, in.Name, logKeyError, err)
		return fmt.Errorf("gateway: %s: remove during reauth: %w", in.Name, err)
	}
	rawCfg := configToMap(cfg)
	if err := t.AddConnection(in.Name, rawCfg); err != nil {
		slog.Error("gateway: IngestOAuthToken — AddConnection failed",
			logKeyConnection, in.Name, logKeyError, err)
		return err
	}
	slog.Info("gateway: IngestOAuthToken — complete",
		logKeyConnection, in.Name,
		"authenticated_by", in.AuthenticatedBy)
	return nil
}

// expiresAtFromSeconds converts an IdP-disclosed expires_in (seconds
// from "now") to an absolute deadline. Returns the zero time when the
// IdP omitted the field — admin UI renders zero as "unknown".
func expiresAtFromSeconds(now time.Time, seconds int) time.Time {
	if seconds <= 0 {
		return time.Time{}
	}
	return now.Add(time.Duration(seconds) * time.Second)
}

// errFmtConnection is the standard "gateway: <name>: <wrapped err>"
// format. Centralized to satisfy revive's add-constant rule.
const errFmtConnection = "gateway: %s: %w"

// configToMap rebuilds the raw config map for a Config — used by the
// OAuth ingest path to round-trip a parsed Config back through
// AddConnection's parser. Only the fields that AddConnection reads need
// to round-trip; ConnectionName is set directly to preserve the original.
func configToMap(c Config) map[string]any {
	m := map[string]any{
		cfgKeyEndpoint:       c.Endpoint,
		cfgKeyAuthMode:       c.AuthMode,
		cfgKeyCredential:     c.Credential,
		cfgKeyConnectionName: c.ConnectionName,
		cfgKeyConnectTimeout: c.ConnectTimeout.String(),
		cfgKeyCallTimeout:    c.CallTimeout.String(),
		cfgKeyTrustLevel:     c.TrustLevel,
	}
	if c.OAuth.Grant != "" {
		m["oauth_grant"] = c.OAuth.Grant
		m["oauth_token_url"] = c.OAuth.TokenURL
		m["oauth_authorization_url"] = c.OAuth.AuthorizationURL
		m["oauth_client_id"] = c.OAuth.ClientID
		m["oauth_client_secret"] = c.OAuth.ClientSecret
		m["oauth_scope"] = c.OAuth.Scope
	}
	return m
}

// ListConnections returns metadata for every live connection, sorted by name.
func (t *Toolkit) ListConnections() []toolkit.ConnectionDetail {
	t.mu.RLock()
	defer t.mu.RUnlock()
	out := make([]toolkit.ConnectionDetail, 0, len(t.connections))
	for name, u := range t.connections {
		out = append(out, toolkit.ConnectionDetail{
			Name:        name,
			Description: u.desc,
			IsDefault:   name == t.defaultName,
		})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
	return out
}

// Close closes every upstream session. Safe to call on a never-registered
// toolkit.
//
// Snapshots the live client pointers under the lock and closes them
// outside the lock so a slow upstream cannot block other toolkit
// operations during shutdown. A concurrent AddConnection that lands
// after the snapshot will leak its client at process exit — by design,
// Close is called once at process shutdown and external callers are
// expected to stop dispatching new AddConnection calls before invoking
// Close (the platform's lifecycle layer enforces this).
func (t *Toolkit) Close() error {
	// Best-effort: stop any debounce timer that hasn't fired yet.
	// time.Timer.Stop returns false if the AfterFunc body is already
	// running on its own goroutine — Stop does NOT wait for it. The
	// in-flight body may still call NotifyToolsListChanged on the
	// (possibly-already-closed) broadcaster after Close returns; that
	// is safe because the broadcaster's Closed state returns
	// ErrBroadcasterClosed and the gatewayListChangedNotifier adapter
	// just logs the failure. So this Stop saves a notification when
	// it can, but does not guarantee no-fire-after-Close.
	t.pendingNotifyTimerMu.Lock()
	if t.pendingNotifyTimer != nil {
		t.pendingNotifyTimer.t.Stop()
		t.pendingNotifyTimer = nil
	}
	t.pendingNotifyTimerMu.Unlock()

	t.mu.Lock()
	clients := make([]*upstreamClient, 0, len(t.connections))
	for _, u := range t.connections {
		if u.client != nil {
			clients = append(clients, u.client)
		}
	}
	t.mu.Unlock()

	var firstErr error
	for _, c := range clients {
		if err := c.close(); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

// addToolsToServerLocked registers each remote tool on the server under a
// namespaced local name. Caller must hold t.mu (write lock) and ensure
// t.server is non-nil. Schedules a debounced tools/list_changed
// notification when at least one tool was registered, so SSE long-poll
// subscribers learn the inventory grew without a manual reconnect.
func (t *Toolkit) addToolsToServerLocked(u *upstream) {
	registered := 0
	for _, rt := range u.tools {
		if rt == nil || rt.Name == "" {
			continue
		}
		local := &mcp.Tool{
			Name:        u.config.ConnectionName + NamespaceSeparator + rt.Name,
			Description: rt.Description,
			InputSchema: rt.InputSchema,
			Title:       rt.Title,
			Annotations: rt.Annotations,
		}
		t.server.AddTool(local, t.makeForwarder(u, rt.Name, local.Name))
		registered++
	}
	if registered > 0 {
		t.notifyToolListChanged()
	}
}

// makeForwarder returns a handler that forwards the call upstream and, if
// an enrichment engine is configured, applies enrichment rules to the
// upstream response.
func (t *Toolkit) makeForwarder(u *upstream, remoteName, localName string) mcp.ToolHandler {
	connection := u.config.ConnectionName
	callTimeout := u.config.CallTimeout
	client := u.client
	return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		if client == nil {
			return upstreamErr(connection, "upstream unavailable"), nil
		}
		callCtx, cancel := context.WithTimeout(ctx, callTimeout)
		defer cancel()

		args := argumentsFromRequest(req)
		res, err := client.callTool(callCtx, remoteName, args)
		if err != nil {
			return upstreamErr(connection, err.Error()), nil
		}
		if !res.IsError {
			t.applyEnrichment(ctx, connection, localName, req, res)
		}
		return res, nil
	}
}

// applyEnrichment runs the configured engine against the upstream response.
// It mutates res.StructuredContent in place when enrichment fires; warnings
// are surfaced as additional TextContent entries appended to res.Content.
// A nil engine, missing structured input, or any internal failure leaves
// the result unchanged so enrichment is never load-bearing for correctness.
func (t *Toolkit) applyEnrichment(ctx context.Context, connection, localName string,
	req *mcp.CallToolRequest, res *mcp.CallToolResult,
) {
	t.mu.RLock()
	engine := t.enrichmentEngine
	t.mu.RUnlock()
	if engine == nil {
		return
	}

	input, ok := structuredInput(res)
	if !ok {
		return
	}

	call := enrichment.CallContext{
		Connection: connection,
		ToolName:   localName,
		Args:       argsFromRequest(req),
	}
	out := engine.Apply(ctx, call, input)
	if out.Response != nil {
		res.StructuredContent = out.Response
	}
	for _, w := range out.Warnings {
		res.Content = append(res.Content, &mcp.TextContent{Text: "warning: " + w})
	}
}

// structuredInput pulls a structured object out of an upstream
// CallToolResult, preferring StructuredContent and falling back to JSON
// embedded in the first TextContent.
func structuredInput(res *mcp.CallToolResult) (any, bool) {
	if res.StructuredContent != nil {
		return res.StructuredContent, true
	}
	if len(res.Content) == 0 {
		return nil, false
	}
	tc, ok := res.Content[0].(*mcp.TextContent)
	if !ok {
		return nil, false
	}
	var parsed any
	if err := json.Unmarshal([]byte(tc.Text), &parsed); err != nil {
		return nil, false
	}
	return parsed, true
}

// argsFromRequest decodes the tool call's raw arguments into a generic
// any value for use as an evaluation-context input.
func argsFromRequest(req *mcp.CallToolRequest) any {
	if req == nil || req.Params == nil || len(req.Params.Arguments) == 0 {
		return nil
	}
	var parsed any
	if err := json.Unmarshal(req.Params.Arguments, &parsed); err != nil {
		return nil
	}
	return parsed
}

// makeLocalNames builds the namespaced tool-name slice for a set of
// discovered remote tools.
func makeLocalNames(connection string, remote []*mcp.Tool) []string {
	out := make([]string, 0, len(remote))
	for _, rt := range remote {
		if rt == nil || rt.Name == "" {
			continue
		}
		out = append(out, connection+NamespaceSeparator+rt.Name)
	}
	return out
}

// dialWithoutAuth dials a non-OAuth upstream (bearer, api_key, none)
// for one-shot Probe-style operations that don't need persistent
// token state. OAuth connections must go through Toolkit.discoverFor
// so the SourceFactory carries the toolkit's connoauth.Store.
func dialWithoutAuth(ctx context.Context, cfg Config) (*upstreamClient, []*mcp.Tool, error) {
	if cfg.AuthMode == AuthModeOAuth {
		return nil, nil, errors.New("gateway: dialWithoutAuth cannot serve oauth connections")
	}
	dialCtx, cancel := context.WithTimeout(ctx, cfg.ConnectTimeout)
	defer cancel()
	client, err := dial(dialCtx, cfg, dialDeps{})
	if err != nil {
		return nil, nil, err
	}
	remoteTools, err := client.listTools(dialCtx)
	if err != nil {
		_ = client.close()
		return nil, nil, fmt.Errorf("list tools: %w", err)
	}
	return client, remoteTools, nil
}

// argumentsFromRequest extracts raw arguments from a tool request. Empty
// arguments are forwarded as nil so the outbound request omits the field.
func argumentsFromRequest(req *mcp.CallToolRequest) any {
	if req == nil || req.Params == nil || len(req.Params.Arguments) == 0 {
		return nil
	}
	return req.Params.Arguments
}

// upstreamErr wraps a message in a tool-error CallToolResult attributed to
// the given connection.
func upstreamErr(connection, msg string) *mcp.CallToolResult {
	return &mcp.CallToolResult{
		IsError: true,
		Content: []mcp.Content{&mcp.TextContent{
			Text: fmt.Sprintf("upstream:%s: %s", connection, msg),
		}},
	}
}

// ProbeTool is a summary of a single discovered upstream tool, used by the
// admin test endpoint to preview what a connection would expose.
type ProbeTool struct {
	Name        string `json:"name"`
	LocalName   string `json:"local_name"`
	Description string `json:"description,omitempty"`
}

// ErrConnectionNotLive is returned by TestLiveConnection when the named
// connection is registered but does not have a live client (typically an
// authorization_code OAuth connection awaiting the operator's first
// Connect, or a connection whose initial dial failed and is sitting as a
// placeholder). Callers convert this to an actionable user-facing message
// rather than a generic upstream error.
var ErrConnectionNotLive = errors.New("gateway: connection has no live client")

// TestLiveConnection exercises the live upstreamClient for the named
// connection by issuing a `tools/list` request — the same call any tool
// invocation would trigger to satisfy auth and reach the upstream. It is
// the correct implementation of an admin "Test connection" button on a
// connection that is already wired into the toolkit: it uses the live
// auth round-tripper (loading and refreshing OAuth tokens via the
// toolkit's TokenStore), so it sees the SAME state that real tool calls
// see. Probe (the config-only one-shot dial) cannot do this because it
// builds an isolated client with no token-store access and would always
// fail for `authorization_code` connections that depend on a persisted
// refresh token.
//
// Returns ErrConnectionNotFound when no row exists for name, and
// ErrConnectionNotLive when the row is a placeholder awaiting reauth.
func (t *Toolkit) TestLiveConnection(ctx context.Context, name string) ([]ProbeTool, error) {
	t.mu.RLock()
	u, ok := t.connections[name]
	if !ok {
		t.mu.RUnlock()
		return nil, fmt.Errorf(errFmtConnection, name, ErrConnectionNotFound)
	}
	if u.client == nil {
		t.mu.RUnlock()
		return nil, fmt.Errorf(errFmtConnection, name, ErrConnectionNotLive)
	}
	client := u.client
	prefix := u.config.ConnectionName
	if prefix == "" {
		prefix = name
	}
	t.mu.RUnlock()

	tools, err := client.listTools(ctx)
	if err != nil {
		return nil, err
	}
	out := make([]ProbeTool, 0, len(tools))
	for _, rt := range tools {
		if rt == nil || rt.Name == "" {
			continue
		}
		out = append(out, ProbeTool{
			Name:        rt.Name,
			LocalName:   prefix + NamespaceSeparator + rt.Name,
			Description: rt.Description,
		})
	}
	return out, nil
}

// Probe dials the upstream described by cfg, lists its tools, and closes the
// session. It does NOT mutate any live toolkit state, so it's safe to call
// from the admin "test connection" endpoint before persisting a config.
func Probe(ctx context.Context, cfg Config) ([]ProbeTool, error) {
	if cfg.ConnectionName == "" {
		cfg.ConnectionName = "probe"
	}
	// Probe is a one-shot dial — never persists tokens. OAuth grants
	// can't be probed without a stored refresh token (which would
	// require Connect first), so the caller is expected to filter to
	// non-OAuth modes before reaching here; dialWithoutAuth refuses
	// AuthModeOAuth defensively.
	client, tools, err := dialWithoutAuth(ctx, cfg)
	if err != nil {
		return nil, err
	}
	defer func() { _ = client.close() }()

	out := make([]ProbeTool, 0, len(tools))
	for _, rt := range tools {
		if rt == nil || rt.Name == "" {
			continue
		}
		out = append(out, ProbeTool{
			Name:        rt.Name,
			LocalName:   cfg.ConnectionName + NamespaceSeparator + rt.Name,
			Description: rt.Description,
		})
	}
	return out, nil
}

// Verify interface compliance at compile time.
var (
	_ interface {
		Kind() string
		Name() string
		Connection() string
		RegisterTools(s *mcp.Server)
		Tools() []string
		SetSemanticProvider(provider semantic.Provider)
		SetQueryProvider(provider query.Provider)
		Close() error
	} = (*Toolkit)(nil)
	_ toolkit.ConnectionManager = (*Toolkit)(nil)
	_ toolkit.ConnectionLister  = (*Toolkit)(nil)
)
