CHORUS/pkg/ai/config_test.go

package ai

import (
	"io/ioutil"
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestNewConfigLoader(t *testing.T) {
	loader := NewConfigLoader("test.yaml", "development")

	assert.Equal(t, "test.yaml", loader.configPath)
	assert.Equal(t, "development", loader.environment)
}

func TestConfigLoaderExpandEnvVars(t *testing.T) {
	loader := NewConfigLoader("", "")

	// Set test environment variables
	os.Setenv("TEST_VAR", "test_value")
	os.Setenv("ANOTHER_VAR", "another_value")
	defer func() {
		os.Unsetenv("TEST_VAR")
		os.Unsetenv("ANOTHER_VAR")
	}()

	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{
			name:     "single variable",
			input:    "endpoint: ${TEST_VAR}",
			expected: "endpoint: test_value",
		},
		{
			name:     "multiple variables",
			input:    "endpoint: ${TEST_VAR}/api\nkey: ${ANOTHER_VAR}",
			expected: "endpoint: test_value/api\nkey: another_value",
		},
		{
			name:     "no variables",
			input:    "endpoint: http://localhost",
			expected: "endpoint: http://localhost",
		},
		{
			name:     "undefined variable",
			input:    "endpoint: ${UNDEFINED_VAR}",
			expected: "endpoint: ",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := loader.expandEnvVars(tt.input)
			assert.Equal(t, tt.expected, result)
		})
	}
}

func TestConfigLoaderApplyEnvironmentOverrides(t *testing.T) {
	loader := NewConfigLoader("", "production")

	config := &ModelConfig{
		DefaultProvider:  "ollama",
		FallbackProvider: "resetdata",
		Environments: map[string]EnvConfig{
			"production": {
				DefaultProvider:  "openai",
				FallbackProvider: "ollama",
			},
			"development": {
				DefaultProvider:  "ollama",
				FallbackProvider: "mock",
			},
		},
	}

	loader.applyEnvironmentOverrides(config)

	assert.Equal(t, "openai", config.DefaultProvider)
	assert.Equal(t, "ollama", config.FallbackProvider)
}

func TestConfigLoaderApplyEnvironmentOverridesNoMatch(t *testing.T) {
	loader := NewConfigLoader("", "testing")

	config := &ModelConfig{
		DefaultProvider:  "ollama",
		FallbackProvider: "resetdata",
		Environments: map[string]EnvConfig{
			"production": {
				DefaultProvider: "openai",
			},
		},
	}

	original := *config
	loader.applyEnvironmentOverrides(config)

	// Should remain unchanged
	assert.Equal(t, original.DefaultProvider, config.DefaultProvider)
	assert.Equal(t, original.FallbackProvider, config.FallbackProvider)
}

func TestConfigLoaderValidateConfig(t *testing.T) {
	loader := NewConfigLoader("", "")

	tests := []struct {
		name      string
		config    *ModelConfig
		expectErr bool
		errMsg    string
	}{
		{
			name: "valid config",
			config: &ModelConfig{
				DefaultProvider:  "test",
				FallbackProvider: "backup",
				Providers: map[string]ProviderConfig{
					"test": {
						Type:         "ollama",
						Endpoint:     "http://localhost:11434",
						DefaultModel: "llama2",
					},
					"backup": {
						Type:         "resetdata",
						Endpoint:     "https://api.resetdata.ai",
						APIKey:       "key",
						DefaultModel: "llama2",
					},
				},
				Roles: map[string]RoleConfig{
					"developer": {
						Provider: "test",
					},
				},
			},
			expectErr: false,
		},
		{
			name: "default provider not found",
			config: &ModelConfig{
				DefaultProvider: "nonexistent",
				Providers: map[string]ProviderConfig{
					"test": {
						Type:         "ollama",
						Endpoint:     "http://localhost:11434",
						DefaultModel: "llama2",
					},
				},
			},
			expectErr: true,
			errMsg:    "default_provider 'nonexistent' not found",
		},
		{
			name: "fallback provider not found",
			config: &ModelConfig{
				FallbackProvider: "nonexistent",
				Providers: map[string]ProviderConfig{
					"test": {
						Type:         "ollama",
						Endpoint:     "http://localhost:11434",
						DefaultModel: "llama2",
					},
				},
			},
			expectErr: true,
			errMsg:    "fallback_provider 'nonexistent' not found",
		},
		{
			name: "invalid provider config",
			config: &ModelConfig{
				Providers: map[string]ProviderConfig{
					"invalid": {
						Type: "invalid_type",
					},
				},
			},
			expectErr: true,
			errMsg:    "invalid provider config 'invalid'",
		},
		{
			name: "invalid role config",
			config: &ModelConfig{
				Providers: map[string]ProviderConfig{
					"test": {
						Type:         "ollama",
						Endpoint:     "http://localhost:11434",
						DefaultModel: "llama2",
					},
				},
				Roles: map[string]RoleConfig{
					"developer": {
						Provider: "nonexistent",
					},
				},
			},
			expectErr: true,
			errMsg:    "invalid role config 'developer'",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := loader.validateConfig(tt.config)

			if tt.expectErr {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tt.errMsg)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func TestConfigLoaderValidateProviderConfig(t *testing.T) {
	loader := NewConfigLoader("", "")

	tests := []struct {
		name      string
		config    ProviderConfig
		expectErr bool
		errMsg    string
	}{
		{
			name: "valid ollama config",
			config: ProviderConfig{
				Type:         "ollama",
				Endpoint:     "http://localhost:11434",
				DefaultModel: "llama2",
				Temperature:  0.7,
				MaxTokens:    4096,
			},
			expectErr: false,
		},
		{
			name: "valid openai config",
			config: ProviderConfig{
				Type:         "openai",
				Endpoint:     "https://api.openai.com/v1",
				APIKey:       "test-key",
				DefaultModel: "gpt-4",
			},
			expectErr: false,
		},
		{
			name: "missing type",
			config: ProviderConfig{
				Endpoint: "http://localhost",
			},
			expectErr: true,
			errMsg:    "type is required",
		},
		{
			name: "invalid type",
			config: ProviderConfig{
				Type:     "invalid",
				Endpoint: "http://localhost",
			},
			expectErr: true,
			errMsg:    "invalid provider type 'invalid'",
		},
		{
			name: "missing endpoint",
			config: ProviderConfig{
				Type: "ollama",
			},
			expectErr: true,
			errMsg:    "endpoint is required",
		},
		{
			name: "openai missing api key",
			config: ProviderConfig{
				Type:         "openai",
				Endpoint:     "https://api.openai.com/v1",
				DefaultModel: "gpt-4",
			},
			expectErr: true,
			errMsg:    "api_key is required for openai provider",
		},
		{
			name: "missing default model",
			config: ProviderConfig{
				Type:     "ollama",
				Endpoint: "http://localhost:11434",
			},
			expectErr: true,
			errMsg:    "default_model is required",
		},
		{
			name: "invalid temperature",
			config: ProviderConfig{
				Type:         "ollama",
				Endpoint:     "http://localhost:11434",
				DefaultModel: "llama2",
				Temperature:  3.0, // Too high
			},
			expectErr: true,
			errMsg:    "temperature must be between 0 and 2.0",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := loader.validateProviderConfig("test", tt.config)

			if tt.expectErr {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tt.errMsg)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func TestConfigLoaderValidateRoleConfig(t *testing.T) {
	loader := NewConfigLoader("", "")

	providers := map[string]ProviderConfig{
		"test": {
			Type: "ollama",
		},
		"backup": {
			Type: "resetdata",
		},
	}

	tests := []struct {
		name      string
		config    RoleConfig
		expectErr bool
		errMsg    string
	}{
		{
			name: "valid role config",
			config: RoleConfig{
				Provider:    "test",
				Model:      "llama2",
				Temperature: 0.7,
				MaxTokens:   4096,
			},
			expectErr: false,
		},
		{
			name: "provider not found",
			config: RoleConfig{
				Provider: "nonexistent",
			},
			expectErr: true,
			errMsg:    "provider 'nonexistent' not found",
		},
		{
			name: "fallback provider not found",
			config: RoleConfig{
				Provider:         "test",
				FallbackProvider: "nonexistent",
			},
			expectErr: true,
			errMsg:    "fallback_provider 'nonexistent' not found",
		},
		{
			name: "invalid temperature",
			config: RoleConfig{
				Provider:    "test",
				Temperature: -1.0,
			},
			expectErr: true,
			errMsg:    "temperature must be between 0 and 2.0",
		},
		{
			name: "invalid max tokens",
			config: RoleConfig{
				Provider:  "test",
				MaxTokens: -100,
			},
			expectErr: true,
			errMsg:    "max_tokens cannot be negative",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := loader.validateRoleConfig("test-role", tt.config, providers)

			if tt.expectErr {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tt.errMsg)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func TestConfigLoaderLoadConfig(t *testing.T) {
	// Create a temporary config file
	configContent := `
providers:
  test:
    type: ollama
    endpoint: http://localhost:11434
    default_model: llama2
    temperature: 0.7

default_provider: test
fallback_provider: test

roles:
  developer:
    provider: test
    model: codellama
`

	tmpFile, err := ioutil.TempFile("", "test-config-*.yaml")
	require.NoError(t, err)
	defer os.Remove(tmpFile.Name())

	_, err = tmpFile.WriteString(configContent)
	require.NoError(t, err)
	tmpFile.Close()

	loader := NewConfigLoader(tmpFile.Name(), "")
	config, err := loader.LoadConfig()

	require.NoError(t, err)
	assert.Equal(t, "test", config.DefaultProvider)
	assert.Equal(t, "test", config.FallbackProvider)
	assert.Len(t, config.Providers, 1)
	assert.Contains(t, config.Providers, "test")
	assert.Equal(t, "ollama", config.Providers["test"].Type)
	assert.Len(t, config.Roles, 1)
	assert.Contains(t, config.Roles, "developer")
	assert.Equal(t, "codellama", config.Roles["developer"].Model)
}

func TestConfigLoaderLoadConfigWithEnvVars(t *testing.T) {
	// Set environment variables
	os.Setenv("TEST_ENDPOINT", "http://test.example.com")
	os.Setenv("TEST_MODEL", "test-model")
	defer func() {
		os.Unsetenv("TEST_ENDPOINT")
		os.Unsetenv("TEST_MODEL")
	}()

	configContent := `
providers:
  test:
    type: ollama
    endpoint: ${TEST_ENDPOINT}
    default_model: ${TEST_MODEL}

default_provider: test
`

	tmpFile, err := ioutil.TempFile("", "test-config-*.yaml")
	require.NoError(t, err)
	defer os.Remove(tmpFile.Name())

	_, err = tmpFile.WriteString(configContent)
	require.NoError(t, err)
	tmpFile.Close()

	loader := NewConfigLoader(tmpFile.Name(), "")
	config, err := loader.LoadConfig()

	require.NoError(t, err)
	assert.Equal(t, "http://test.example.com", config.Providers["test"].Endpoint)
	assert.Equal(t, "test-model", config.Providers["test"].DefaultModel)
}

func TestConfigLoaderLoadConfigFileNotFound(t *testing.T) {
	loader := NewConfigLoader("nonexistent.yaml", "")
	_, err := loader.LoadConfig()

	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to read config file")
}

func TestConfigLoaderLoadConfigInvalidYAML(t *testing.T) {
	// Create a file with invalid YAML
	tmpFile, err := ioutil.TempFile("", "invalid-config-*.yaml")
	require.NoError(t, err)
	defer os.Remove(tmpFile.Name())

	_, err = tmpFile.WriteString("invalid: yaml: content: [")
	require.NoError(t, err)
	tmpFile.Close()

	loader := NewConfigLoader(tmpFile.Name(), "")
	_, err = loader.LoadConfig()

	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to parse config file")
}

func TestDefaultConfigPath(t *testing.T) {
	// Test with environment variable
	os.Setenv("CHORUS_MODEL_CONFIG", "/custom/path/models.yaml")
	defer os.Unsetenv("CHORUS_MODEL_CONFIG")

	path := DefaultConfigPath()
	assert.Equal(t, "/custom/path/models.yaml", path)

	// Test without environment variable
	os.Unsetenv("CHORUS_MODEL_CONFIG")
	path = DefaultConfigPath()
	assert.Equal(t, "configs/models.yaml", path)
}

func TestGetEnvironment(t *testing.T) {
	// Test with CHORUS_ENVIRONMENT
	os.Setenv("CHORUS_ENVIRONMENT", "production")
	defer os.Unsetenv("CHORUS_ENVIRONMENT")

	env := GetEnvironment()
	assert.Equal(t, "production", env)

	// Test with NODE_ENV fallback
	os.Unsetenv("CHORUS_ENVIRONMENT")
	os.Setenv("NODE_ENV", "staging")
	defer os.Unsetenv("NODE_ENV")

	env = GetEnvironment()
	assert.Equal(t, "staging", env)

	// Test default
	os.Unsetenv("CHORUS_ENVIRONMENT")
	os.Unsetenv("NODE_ENV")

	env = GetEnvironment()
	assert.Equal(t, "development", env)
}

func TestModelConfig(t *testing.T) {
	config := ModelConfig{
		Providers: map[string]ProviderConfig{
			"test": {
				Type:         "ollama",
				Endpoint:     "http://localhost:11434",
				DefaultModel: "llama2",
			},
		},
		DefaultProvider:  "test",
		FallbackProvider: "test",
		Roles: map[string]RoleConfig{
			"developer": {
				Provider: "test",
				Model:   "codellama",
			},
		},
		Environments: map[string]EnvConfig{
			"production": {
				DefaultProvider: "openai",
			},
		},
		ModelPreferences: map[string]TaskPreference{
			"code_generation": {
				PreferredModels:  []string{"codellama", "gpt-4"},
				MinContextTokens: 8192,
			},
		},
	}

	assert.Len(t, config.Providers, 1)
	assert.Len(t, config.Roles, 1)
	assert.Len(t, config.Environments, 1)
	assert.Len(t, config.ModelPreferences, 1)
}

func TestEnvConfig(t *testing.T) {
	envConfig := EnvConfig{
		DefaultProvider:  "openai",
		FallbackProvider: "ollama",
	}

	assert.Equal(t, "openai", envConfig.DefaultProvider)
	assert.Equal(t, "ollama", envConfig.FallbackProvider)
}

func TestTaskPreference(t *testing.T) {
	pref := TaskPreference{
		PreferredModels:  []string{"gpt-4", "codellama:13b"},
		MinContextTokens: 8192,
	}

	assert.Len(t, pref.PreferredModels, 2)
	assert.Equal(t, 8192, pref.MinContextTokens)
	assert.Contains(t, pref.PreferredModels, "gpt-4")
}