package config

import (
	"io/ioutil"

	"gopkg.in/yaml.v3"
)

// Config 定义压测配置结构
// 可根据 production.md 扩展字段

type APIConfig struct {
	Endpoint   string `yaml:"endpoint"`
	APIKey     string `yaml:"api_key"`
	Model      string `yaml:"model"`
	Streaming  bool   `yaml:"streaming"`
	ClientType string `yaml:"client"`
}

type PromptTemplate struct {
	TargetTokens int      `yaml:"target_tokens"`
	Templates    []string `yaml:"templates"`
}

type RequestMix struct {
	Type   string  `yaml:"type"`
	Weight float64 `yaml:"weight"`
}

type ConcurrencyConfig struct {
	Steps           []int `yaml:"steps"`
	DurationPerStep int   `yaml:"duration_per_step"`
	MaxGoroutines   int   `yaml:"max_goroutines"`
}

type Config struct {
	API             APIConfig                 `yaml:"api"`
	PromptTemplates map[string]PromptTemplate `yaml:"prompt_templates"`
	Requests        []RequestMix              `yaml:"requests"`
	Concurrency     ConcurrencyConfig         `yaml:"concurrency"`
	Timeout         int                       `yaml:"timeout"`
	PoissonLambda   float64                   `yaml:"poisson_lambda"`
	TokenizerModel  string                    `yaml:"tokenizer_model"`
}

func LoadConfig(path string) (*Config, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var cfg Config
	if err := yaml.Unmarshal(data, &cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}