package runtime

import (
	"context"
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
)

type RuntimeDefaults struct {
	Model           string
	ReasoningEffort string
}

// LoadMachineRuntimeCapabilities loads supported runtime names and defaults for machines.
func LoadMachineRuntimeCapabilities(ctx context.Context, queryer coredb.Queryer, machineIDs []string) (map[string][]string, map[string]map[string]RuntimeDefaults, error) {
	normalizedMachineIDs := uniqueTrimmedStrings(machineIDs)
	if len(normalizedMachineIDs) == 0 {
		return map[string][]string{}, map[string]map[string]RuntimeDefaults{}, nil
	}

	requestedMachineIDsCTE, requestedMachineIDArgs, _, err := coredb.BuildStringValuesCTE("requested_machine_ids", "machine_id", normalizedMachineIDs, 1)
	if err != nil {
		return nil, nil, fmt.Errorf("build machine runtime capability query: %w", err)
	}

	rows, err := queryer.Query(ctx, fmt.Sprintf(`
WITH %s
SELECT requested.machine_id, capabilities.runtime
     , COALESCE(capabilities.default_model, '')
     , COALESCE(capabilities.default_reasoning_effort, '')
FROM requested_machine_ids AS requested
LEFT JOIN machine_runtime_capabilities AS capabilities
  ON capabilities.machine_id = requested.machine_id
ORDER BY requested.machine_id ASC, capabilities.runtime ASC
`, requestedMachineIDsCTE), requestedMachineIDArgs...)
	if err != nil {
		return nil, nil, fmt.Errorf("load machine runtime capabilities: %w", err)
	}
	defer rows.Close()

	capabilityMap := make(map[string][]string, len(normalizedMachineIDs))
	defaultsMap := make(map[string]map[string]RuntimeDefaults, len(normalizedMachineIDs))
	for _, machineID := range normalizedMachineIDs {
		capabilityMap[machineID] = []string{}
		defaultsMap[machineID] = map[string]RuntimeDefaults{}
	}

	for rows.Next() {
		var machineID string
		var runtime *string
		var defaultModel string
		var defaultReasoningEffort string
		if err := rows.Scan(&machineID, &runtime, &defaultModel, &defaultReasoningEffort); err != nil {
			return nil, nil, fmt.Errorf("scan machine runtime capability: %w", err)
		}
		if runtime == nil {
			continue
		}
		trimmedRuntime := strings.TrimSpace(*runtime)
		if trimmedRuntime == "" {
			continue
		}
		capabilityMap[machineID] = append(capabilityMap[machineID], trimmedRuntime)
		defaultsMap[machineID][trimmedRuntime] = RuntimeDefaults{
			Model:           strings.TrimSpace(defaultModel),
			ReasoningEffort: strings.TrimSpace(defaultReasoningEffort),
		}
	}
	if err := rows.Err(); err != nil {
		return nil, nil, fmt.Errorf("iterate machine runtime capabilities: %w", err)
	}
	return capabilityMap, defaultsMap, nil
}

func uniqueTrimmedStrings(values []string) []string {
	result := make([]string, 0, len(values))
	seen := make(map[string]struct{}, len(values))
	for _, value := range values {
		trimmed := strings.TrimSpace(value)
		if trimmed == "" {
			continue
		}
		if _, exists := seen[trimmed]; exists {
			continue
		}
		seen[trimmed] = struct{}{}
		result = append(result, trimmed)
	}
	return result
}