agentu/pkg/config/config.go

package config

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"gopkg.in/yaml.v3"
)

const DefaultPath = "~/.agentu/config.yaml"

type Config struct {
	Providers map[string]ProviderConfig `yaml:"providers"`
	Agent     AgentConfig               `yaml:"agent"`
}

type ProviderConfig struct {
	Name               string
	BaseURL            string
	Model              string
	Models             []string
	ModelConfigs       map[string]ModelConfig
	ContextTokens      int
	APIKey             string
	Thinking           string
	ThinkingParam      string
	ThinkingConfigured bool
}

type ModelConfig struct {
	ID            string `yaml:"id"`
	Thinking      string `yaml:"thinking"`
	ContextTokens int    `yaml:"context"`
}

type AgentConfig struct {
	SystemPrompt     string        `yaml:"system_prompt"`
	WorkingDir       string        `yaml:"working_dir"`
	ToolTimeout      time.Duration `yaml:"tool_timeout"`
	MaxContextTokens int           `yaml:"max_context_tokens"`
}

type rawConfig struct {
	Providers map[string]rawProviderConfig `yaml:"providers"`
	Agent     struct {
		SystemPrompt     string `yaml:"system_prompt"`
		WorkingDir       string `yaml:"working_dir"`
		ToolTimeout      string `yaml:"tool_timeout"`
		MaxContextTokens int    `yaml:"max_context_tokens"`
	} `yaml:"agent"`
}

type rawProviderConfig struct {
	BaseURL string           `yaml:"base_url"`
	Models  []rawModelConfig `yaml:"models"`
	APIKey  string           `yaml:"api_key"`
}

type rawModelConfig struct {
	ID            string `yaml:"id"`
	Thinking      string `yaml:"thinking"`
	ContextTokens int    `yaml:"context"`
}

func Load(path string) (*Config, error) {
	path, err := ResolvePath(path)
	if err != nil {
		return nil, err
	}

	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read config: %w", err)
	}

	expanded := os.ExpandEnv(string(data))
	var raw rawConfig
	decoder := yaml.NewDecoder(strings.NewReader(expanded))
	decoder.KnownFields(true)
	if err := decoder.Decode(&raw); err != nil {
		return nil, fmt.Errorf("parse config: %w", err)
	}

	cfg := &Config{
		Providers: normalizeProviders(raw),
		Agent: AgentConfig{
			SystemPrompt:     raw.Agent.SystemPrompt,
			WorkingDir:       raw.Agent.WorkingDir,
			MaxContextTokens: raw.Agent.MaxContextTokens,
		},
	}

	if cfg.Agent.WorkingDir == "" {
		cfg.Agent.WorkingDir = "."
	}
	absWorkingDir, err := filepath.Abs(cfg.Agent.WorkingDir)
	if err != nil {
		return nil, fmt.Errorf("resolve working_dir: %w", err)
	}
	cfg.Agent.WorkingDir = absWorkingDir

	if raw.Agent.ToolTimeout == "" {
		cfg.Agent.ToolTimeout = 30 * time.Second
	} else {
		timeout, err := time.ParseDuration(raw.Agent.ToolTimeout)
		if err != nil {
			return nil, fmt.Errorf("parse agent.tool_timeout: %w", err)
		}
		cfg.Agent.ToolTimeout = timeout
	}

	if cfg.Agent.SystemPrompt == "" {
		cfg.Agent.SystemPrompt = "You are agentu, a concise local agent. Answer simple conversation directly. Use tools only when the user's request needs local project context or local command execution."
	}

	if err := cfg.Validate(); err != nil {
		return nil, err
	}

	return cfg, nil
}

func ResolvePath(path string) (string, error) {
	if strings.TrimSpace(path) == "" {
		path = DefaultPath
	}
	if path == "~" {
		home, err := os.UserHomeDir()
		if err != nil {
			return "", fmt.Errorf("resolve home dir: %w", err)
		}
		return home, nil
	}
	if strings.HasPrefix(path, "~/") {
		home, err := os.UserHomeDir()
		if err != nil {
			return "", fmt.Errorf("resolve home dir: %w", err)
		}
		return filepath.Join(home, strings.TrimPrefix(path, "~/")), nil
	}
	return path, nil
}

func (c *Config) Validate() error {
	if len(c.Providers) == 0 {
		return errors.New("missing required config: providers")
	}
	for name, provider := range c.Providers {
		var missing []string
		prefix := "providers." + name
		if strings.TrimSpace(provider.BaseURL) == "" {
			missing = append(missing, prefix+".base_url")
		}
		if len(provider.Models) == 0 {
			missing = append(missing, prefix+".models")
		}
		if strings.TrimSpace(provider.APIKey) == "" {
			missing = append(missing, prefix+".api_key")
		}
		if len(missing) > 0 {
			return fmt.Errorf("missing required config: %s", strings.Join(missing, ", "))
		}
		for _, model := range provider.Models {
			if strings.TrimSpace(model) == "" {
				return fmt.Errorf("%s.models.id is required", prefix)
			}
		}
		for _, model := range provider.ModelConfigs {
			if model.ContextTokens < 0 {
				return fmt.Errorf("%s.models.context must be greater than or equal to zero", prefix)
			}
			if model.Thinking != "" && !IsThinkingLevel(model.Thinking) {
				return fmt.Errorf("%s.models.thinking must be one of: %s", prefix, strings.Join(ThinkingLevels(), ", "))
			}
		}
		if provider.ThinkingConfigured && !IsThinkingLevel(provider.Thinking) {
			return fmt.Errorf("%s.thinking must be one of: %s", prefix, strings.Join(ThinkingLevels(), ", "))
		}
	}
	if c.Agent.ToolTimeout <= 0 {
		return errors.New("agent.tool_timeout must be greater than zero")
	}
	if c.Agent.MaxContextTokens < 0 {
		return errors.New("agent.max_context_tokens must be greater than or equal to zero")
	}
	return nil
}

func (c *Config) DefaultProviderName() string {
	names := c.ProviderNames()
	if len(names) == 0 {
		return ""
	}
	return names[0]
}

func (c *Config) ProviderNames() []string {
	names := make([]string, 0, len(c.Providers))
	for name := range c.Providers {
		names = append(names, name)
	}
	sort.Strings(names)
	return names
}

func normalizeProviders(raw rawConfig) map[string]ProviderConfig {
	providers := make(map[string]ProviderConfig)
	for name, provider := range raw.Providers {
		providers[name] = normalizeProvider(name, provider)
	}
	return providers
}

func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
	models, modelConfigs := normalizeModelConfigs(raw.Models)
	provider := ProviderConfig{
		Name:         name,
		BaseURL:      raw.BaseURL,
		Models:       models,
		ModelConfigs: modelConfigs,
		APIKey:       raw.APIKey,
	}
	if len(provider.Models) > 0 {
		provider.Model = provider.Models[0]
	}
	applyModelDefaults(&provider)
	if provider.ThinkingConfigured && provider.ThinkingParam == "" {
		provider.ThinkingParam = "thinking"
	}
	return provider
}

func normalizeModelConfigs(rawModels []rawModelConfig) ([]string, map[string]ModelConfig) {
	models := make([]string, 0, len(rawModels))
	configs := make(map[string]ModelConfig, len(rawModels))
	for _, raw := range rawModels {
		model := ModelConfig{
			ID:            strings.TrimSpace(raw.ID),
			Thinking:      strings.ToLower(strings.TrimSpace(raw.Thinking)),
			ContextTokens: raw.ContextTokens,
		}
		models = append(models, model.ID)
		if model.ID != "" {
			configs[model.ID] = model
		}
	}
	return models, configs
}

func applyModelDefaults(provider *ProviderConfig) {
	model, ok := provider.ModelConfigs[provider.Model]
	if !ok {
		return
	}
	provider.ContextTokens = model.ContextTokens
	if model.Thinking != "" {
		provider.Thinking = model.Thinking
		provider.ThinkingConfigured = true
	}
}

func (p ProviderConfig) WithModel(model string) ProviderConfig {
	p.Model = strings.TrimSpace(model)
	applyModelDefaults(&p)
	return p
}

func ThinkingLevels() []string {
	return []string{"none", "middle", "high", "xhigh", "max"}
}

func IsThinkingLevel(value string) bool {
	value = strings.ToLower(strings.TrimSpace(value))
	for _, level := range ThinkingLevels() {
		if value == level {
			return true
		}
	}
	return false
}