package utils

import (
	"os"

	toml "github.com/BurntSushi/toml"
	"github.com/charmbracelet/log"
)

type WorkerConfig struct {
	NumberOfWorkers int    `toml:"number_of_workers"`
	Pulling         string `toml:"pulling_time"`
}

type Config struct {
	Hostname        string
	Port            int
	NumberOfWorkers int `toml:"number_of_workers"`
	SupressCuda     int `toml:"supress_cuda"`

	GpuWorker WorkerConfig `toml:"Worker"`
}

func LoadConfig() Config {

	log.Info("Loading the config file")

	dat, err := os.ReadFile("./config.toml")
	if err != nil {
		log.Error("Failed to load config file", "err", err)
		// Use default values
		return Config{
			Hostname:        "localhost",
			Port:            8000,
			NumberOfWorkers: 10,
			GpuWorker: WorkerConfig{
				NumberOfWorkers: 1,
				Pulling:         "500ms",
			},
		}
	}

	var conf Config
	_, err = toml.Decode(string(dat), &conf)

    if conf.SupressCuda == 1 {
        log.Warn("Supressing Cuda Messages!")
        os.Setenv("TF_CPP_MIN_VLOG_LEVEL", "3") 
        os.Setenv("TF_CPP_MIN_LOG_LEVEL", "3")
    }

	return conf
}