package config

import (
	"fmt"
	"os"
	"strings"

	"github.com/knadh/koanf/parsers/yaml"
	env "github.com/knadh/koanf/providers/env/v2"
	"github.com/knadh/koanf/providers/file"
	"github.com/knadh/koanf/v2"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent"
)

// LoadAgent reads an agent configuration, overlaying it on the defaults and
// then applying LLMGW_-prefixed environment overrides.
//
// An empty path loads the defaults, which with --server and --node-id on the
// command line is enough for a smoke test on a node with no IB.
func LoadAgent(path string) (agent.Config, error) {
	cfg := agent.DefaultConfig()
	k := koanf.New(".")

	if path != "" {
		if _, err := os.Stat(path); err != nil {
			return cfg, fmt.Errorf("config file: %w", err)
		}
		if err := k.Load(file.Provider(path), yaml.Parser()); err != nil {
			return cfg, fmt.Errorf("parsing %s: %w", path, err)
		}
	}

	err := k.Load(env.Provider(".", env.Opt{
		Prefix: "LLMGW_",
		TransformFunc: func(key, value string) (string, any) {
			key = strings.ToLower(strings.TrimPrefix(key, "LLMGW_"))
			return strings.ReplaceAll(key, "__", "."), value
		},
	}), nil)
	if err != nil {
		return cfg, fmt.Errorf("environment overrides: %w", err)
	}

	if err := k.UnmarshalWithConf("", &cfg, koanf.UnmarshalConf{Tag: "koanf"}); err != nil {
		return cfg, fmt.Errorf("decoding configuration: %w", err)
	}
	return cfg, validateAgent(cfg)
}

// validateAgent rejects configurations that would fail confusingly at runtime.
func validateAgent(c agent.Config) error {
	seen := map[string]struct{}{}
	for i, f := range c.Fabrics {
		if f.Fabric == "" {
			return fmt.Errorf("fabrics[%d]: name is required", i)
		}
		if _, dup := seen[f.Fabric]; dup {
			return fmt.Errorf("fabric %q is declared twice", f.Fabric)
		}
		seen[f.Fabric] = struct{}{}
		if len(f.HCAs) == 0 {
			// A fabric with no adapters is worse than no fabric at all: the
			// scheduler would place multi-node replicas here and every shard
			// would come up with NCCL unable to find a path (plan §5.3).
			return fmt.Errorf("fabric %q lists no HCAs", f.Fabric)
		}
	}

	// The same adapter on two fabrics means one of the labels is wrong, and the
	// consequence is a replica that starts and then hangs on its first
	// all-reduce — the failure mode this labelling exists to prevent.
	byHCA := map[string]string{}
	for _, f := range c.Fabrics {
		for _, h := range f.HCAs {
			if prev, dup := byHCA[h]; dup {
				return fmt.Errorf("HCA %q is claimed by both %q and %q", h, prev, f.Fabric)
			}
			byHCA[h] = f.Fabric
		}
	}
	return nil
}