package pkulaw

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[EndpointID]*mcp.Client

// Executor resolves one installation's prepared topology and credential,
// selects one pre-registered endpoint/tool operation, and never performs
// invocation-time discovery or failover.
type Executor struct {
	connections runner.ConnectionSource
	clients     mcpClients
	validators  catalogValidators
}

func newExecutor(
	connections runner.ConnectionSource,
	clients mcpClients,
) (*Executor, error) {
	if connections == nil {
		return nil, fmt.Errorf("PKULaw connection source is required")
	}
	if err := validateClientSet(clients); err != nil {
		return nil, err
	}
	validators, err := compileCatalogValidators(Definition())
	if err != nil {
		return nil, err
	}
	return &Executor{
		connections: connections,
		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(aggregateRoutes))
	for name := range aggregateRoutes {
		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 || executor.connections == nil {
		return registry.BuiltinResult{}, fmt.Errorf(
			"PKULaw MCP executor is not configured",
		)
	}
	input, err := executor.validators.validateInput(
		invocation.Action,
		invocation.Input,
	)
	if err != nil {
		return registry.BuiltinResult{}, runner.InvocationError(err)
	}

	resolved, err := executor.connections.ResolveProviderConnection(
		ctx,
		invocation.WorkspaceID,
		invocation.ServiceInstallationID,
		ConnectionConfigContract(),
	)
	if err != nil {
		mapped := runner.InvocationError(
			runner.CredentialsUnavailable(err),
		)
		if mapped == nil {
			return registry.BuiltinResult{}, err
		}
		return registry.BuiltinResult{}, mapped
	}

	if err := ConnectionConfigContract().Validate(resolved.Config); err != nil {
		return registry.BuiltinResult{}, runner.InvocationError(
			invalidConnectionTopology(),
		)
	}
	topology := resolved.Config["topology"]
	var route ActionRoute
	var routeExists bool
	switch topology {
	case TopologyAggregate:
		route, routeExists = aggregateRoutes[invocation.Action]
	case TopologyFallback:
		route, routeExists = fallbackRoutes[invocation.Action]
	default:
		return registry.BuiltinResult{}, runner.InvocationError(
			invalidConnectionTopology(),
		)
	}
	if !routeExists {
		return registry.BuiltinResult{}, runner.InvocationError(
			actionNotFound(),
		)
	}
	client := executor.clients[route.Endpoint]
	if client == nil {
		return registry.BuiltinResult{}, runner.InvocationError(
			runner.Errorf(
				runner.ErrorConnectionUnavailable,
				"PKULaw MCP endpoint is unavailable",
				nil,
			),
		)
	}
	arguments, err := mapToolArguments(
		invocation.Action,
		topology,
		input,
	)
	if err != nil {
		return registry.BuiltinResult{}, runner.InvocationError(err)
	}
	result, err := client.CallTool(
		ctx,
		resolved.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 newOperationClients(
	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(
				"PKULaw: build %s MCP operation client: %w",
				policy.ID,
				err,
			)
		}
		clients[policy.ID] = client
	}
	return clients, nil
}

func validateClientSet(clients mcpClients) error {
	if len(clients) != len(endpointPolicies) {
		return fmt.Errorf("PKULaw MCP clients are incomplete")
	}
	for _, policy := range endpointPolicies {
		if clients[policy.ID] == nil {
			return fmt.Errorf(
				"PKULaw MCP client for %s is unavailable",
				policy.ID,
			)
		}
	}
	return nil
}

func invalidConnectionTopology() error {
	return runner.Errorf(
		runner.ErrorConnectionUnavailable,
		"PKULaw connection topology is unavailable",
		nil,
	)
}