package aifin

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sort"

	"github.com/tidbcloud/AgentTeam/internal/extension/registry"
	"github.com/tidbcloud/AgentTeam/internal/extension/runner"
	"github.com/tidbcloud/AgentTeam/internal/extension/runner/mcp"
)

type mcpClients map[Family]*mcp.Client

// Executor resolves one installation's credential and dispatches one static
// semantic Action to its code-owned endpoint/tool route.
type Executor struct {
	credentials runner.CredentialSource
	clients     mcpClients
	validators  catalogValidators
}

func newExecutor(
	credentials runner.CredentialSource,
	clients mcpClients,
) (*Executor, error) {
	if credentials == nil {
		return nil, fmt.Errorf(
			"wind AIFin Market credential source is required",
		)
	}
	if len(clients) != len(endpointPolicies) {
		return nil, fmt.Errorf(
			"wind AIFin Market MCP clients are incomplete",
		)
	}
	for _, policy := range endpointPolicies {
		if clients[policy.Family] == nil {
			return nil, fmt.Errorf(
				"wind AIFin Market MCP client for %s is unavailable",
				policy.Family,
			)
		}
	}
	validators, err := compileCatalogValidators(Definition())
	if err != nil {
		return nil, err
	}
	return &Executor{
		credentials: credentials,
		clients:     clients,
		validators:  validators,
	}, nil
}

func (*Executor) Strategy() string {
	return "mcp"
}

func (executor *Executor) ActionNames() []string {
	if executor == nil {
		return nil
	}
	names := make([]string, 0, len(actionRoutes))
	for name := range actionRoutes {
		names = append(names, name)
	}
	sort.Strings(names)
	return names
}

func (executor *Executor) Invoke(
	ctx context.Context,
	invocation registry.BuiltinInvocation,
) (registry.BuiltinResult, error) {
	if executor == nil {
		return registry.BuiltinResult{}, fmt.Errorf(
			"wind AIFin Market MCP executor is not configured",
		)
	}
	input, err := executor.validators.validateInput(
		invocation.Action,
		invocation.Input,
	)
	if err != nil {
		return registry.BuiltinResult{}, runner.InvocationError(err)
	}
	route, ok := actionRoutes[invocation.Action]
	if !ok {
		return registry.BuiltinResult{}, runner.InvocationError(
			actionNotFound(),
		)
	}
	client := executor.clients[route.Family]
	if client == nil {
		return registry.BuiltinResult{}, runner.InvocationError(
			runner.Errorf(
				runner.ErrorConnectionUnavailable,
				"Wind AIFin Market MCP endpoint is unavailable",
				nil,
			),
		)
	}
	arguments, err := mapToolArguments(invocation.Action, input)
	if err != nil {
		return registry.BuiltinResult{}, runner.InvocationError(err)
	}
	credentials, err := executor.credentials.ResolveProviderCredentials(
		ctx,
		invocation.WorkspaceID,
		invocation.ServiceInstallationID,
	)
	if err != nil {
		mapped := runner.InvocationError(
			runner.CredentialsUnavailable(err),
		)
		if mapped == nil {
			return registry.BuiltinResult{}, err
		}
		return registry.BuiltinResult{}, mapped
	}
	result, err := client.CallTool(
		ctx,
		credentials,
		route.Tool,
		arguments,
	)
	if err != nil {
		if mapped := runner.InvocationError(err); mapped != nil {
			return registry.BuiltinResult{}, mapped
		}
		return registry.BuiltinResult{}, err
	}
	output, err := normalizeToolResult(result)
	if err != nil {
		return registry.BuiltinResult{}, runner.InvocationError(err)
	}
	normalized, err := executor.validators.validateOutput(
		invocation.Action,
		output,
	)
	if err != nil {
		return registry.BuiltinResult{}, runner.InvocationError(err)
	}
	return registry.BuiltinResult{
		Data: append(json.RawMessage(nil), normalized...),
	}, nil
}

func newMCPClients(transport http.RoundTripper) (mcpClients, error) {
	if err := validateRouting(); err != nil {
		return nil, err
	}
	clients := make(mcpClients, len(endpointPolicies))
	for _, policy := range endpointPolicies {
		client, err := mcp.NewClient(mcp.Options{
			Endpoint:     policy.Endpoint,
			Auth:         AuthSpec(),
			AllowedTools: append([]string(nil), policy.AllowedTools...),
			Transport:    transport,
		})
		if err != nil {
			return nil, fmt.Errorf(
				"wind AIFin Market: build %s MCP client: %w",
				policy.Family,
				err,
			)
		}
		clients[policy.Family] = client
	}
	return clients, nil
}

func validateRouting() error {
	toolsByFamily := make(map[Family]map[string]struct{}, len(endpointPolicies))
	for _, policy := range endpointPolicies {
		if policy.Family == "" ||
			policy.Endpoint == "" ||
			len(policy.AllowedTools) == 0 {
			return fmt.Errorf(
				"wind AIFin Market endpoint policy is invalid",
			)
		}
		if _, duplicate := toolsByFamily[policy.Family]; duplicate {
			return fmt.Errorf(
				"wind AIFin Market endpoint family is duplicated",
			)
		}
		tools := make(map[string]struct{}, len(policy.AllowedTools))
		for _, tool := range policy.AllowedTools {
			if tool == "" {
				return fmt.Errorf(
					"wind AIFin Market tool policy is invalid",
				)
			}
			if _, duplicate := tools[tool]; duplicate {
				return fmt.Errorf(
					"wind AIFin Market tool policy is duplicated",
				)
			}
			tools[tool] = struct{}{}
		}
		toolsByFamily[policy.Family] = tools
	}
	if len(actionRoutes) != 13 {
		return fmt.Errorf("wind AIFin Market action routing is incomplete")
	}
	routedTools := make(map[Family]map[string]struct{}, len(toolsByFamily))
	for action, route := range actionRoutes {
		if action == "" {
			return fmt.Errorf("wind AIFin Market action routing is invalid")
		}
		tools, ok := toolsByFamily[route.Family]
		if !ok {
			return fmt.Errorf(
				"wind AIFin Market action family is not allowed",
			)
		}
		if _, ok := tools[route.Tool]; !ok {
			return fmt.Errorf(
				"wind AIFin Market action tool is not allowed",
			)
		}
		if routedTools[route.Family] == nil {
			routedTools[route.Family] = make(map[string]struct{})
		}
		if _, duplicate := routedTools[route.Family][route.Tool]; duplicate {
			return fmt.Errorf(
				"wind AIFin Market action tool is duplicated",
			)
		}
		routedTools[route.Family][route.Tool] = struct{}{}
	}
	for family, tools := range toolsByFamily {
		if len(routedTools[family]) != len(tools) {
			return fmt.Errorf(
				"wind AIFin Market tool allowlist has no static Action",
			)
		}
	}
	return nil
}
