package middleware

import (
	"context"
	"maps"

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

// Upstream tool names referenced across multiple middleware files. These are
// fully-qualified MCP tool names exposed by the Trino and DataHub toolkits.
const (
	toolNameTrinoQuery       = "trino_query"
	toolNameTrinoExecute     = "trino_execute"
	toolNameDatahubSearch    = "datahub_search"
	toolNameDatahubGetEntity = "datahub_get_entity"
)

// defaultDescriptionOverrides contains built-in description overrides that
// guide agents toward DataHub discovery before running queries.
var defaultDescriptionOverrides = map[string]string{
	toolNameTrinoQuery: "Execute a read-only SQL query against Trino and return results. " +
		"IMPORTANT: Before writing SQL, call datahub_search to discover the table and " +
		"understand its business context (descriptions, owners, tags, glossary terms). " +
		"The search results include schema previews with column names and types. " +
		"Only SELECT, SHOW, DESCRIBE, EXPLAIN, and WITH statements are allowed.",
	toolNameTrinoExecute: "Execute a SQL statement against Trino, including write operations. " +
		"IMPORTANT: Before writing SQL, call datahub_search to discover the table and " +
		"understand its business context (descriptions, owners, tags, glossary terms). " +
		"Use trino_query for read-only SELECT queries. This tool should be used when " +
		"you need to modify data or schema.",
}

// MergedDescriptionOverrides merges the built-in default overrides with
// user-provided config overrides. Config overrides take precedence.
func MergedDescriptionOverrides(configOverrides map[string]string) map[string]string {
	merged := make(map[string]string, len(defaultDescriptionOverrides)+len(configOverrides))
	maps.Copy(merged, defaultDescriptionOverrides)
	maps.Copy(merged, configOverrides)
	return merged
}

// MCPDescriptionOverrideMiddleware creates MCP protocol-level middleware that
// replaces tool descriptions in tools/list responses. This is used to inject
// workflow guidance (e.g., "call datahub_search first") into tool descriptions
// that agents see when discovering available tools.
//
// Deprecated: prefer MCPDescriptionOverrideMiddlewareDynamic so admin-API
// edits to tool.<name>.description take effect on the next tools/list call
// instead of requiring a platform restart.
func MCPDescriptionOverrideMiddleware(overrides map[string]string) mcp.Middleware {
	return MCPDescriptionOverrideMiddlewareDynamic(func() map[string]string {
		return overrides
	})
}

// MCPDescriptionOverrideMiddlewareDynamic is a hot-reloading variant that
// re-resolves the override map on every tools/list call. The getter is
// expected to merge built-in defaults with the live config overrides.
// Used by Platform so per-tool description overrides authored from the
// admin portal take effect immediately.
func MCPDescriptionOverrideMiddlewareDynamic(getOverrides func() map[string]string) mcp.Middleware {
	return func(next mcp.MethodHandler) mcp.MethodHandler {
		return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
			result, err := next(ctx, method, req)
			if err != nil {
				return result, err
			}
			if method != methodToolsList {
				return result, nil
			}
			return applyDescriptionOverrides(getOverrides(), result), nil
		}
	}
}

// applyDescriptionOverrides replaces tool descriptions for matching names.
func applyDescriptionOverrides(overrides map[string]string, result mcp.Result) mcp.Result {
	listResult, ok := result.(*mcp.ListToolsResult)
	if !ok || listResult == nil || len(overrides) == 0 {
		return result
	}

	for _, tool := range listResult.Tools {
		if desc, found := overrides[tool.Name]; found {
			tool.Description = desc
		}
	}

	return listResult
}
