package semantic

import (
	"crypto/sha256"
	"encoding/hex"
	"sort"
	"strings"
)

// SplitRatio controls the frozen holdout fraction of the total corpus.
// 0.15 means 15% of families go to the frozen holdout, 85% to the
// development pool.
const DefaultFrozenRatio = 0.15

// SplitByFamily divides examples into train/eval splits such that all
// paraphrases or contrastive variants of one seed stay in the same split.
//
// The algorithm is deterministic: each SplitGroup gets a hash, and the hash
// bucket determines the split. This prevents the failure mode where
// "выключи свет" lands in train and "пожалуйста выключи свет" lands in eval.
//
// splitRatio controls the train fraction (0.0–1.0). 0.8 means 80% train.
func SplitByFamily(exs []RouteExample, splitRatio float64) (train, eval []RouteExample) {
	if splitRatio <= 0 || splitRatio >= 1 {
		splitRatio = 0.8
	}

	families := groupByFamily(exs)

	famKeys := sortedFamilyKeys(families)

	for _, k := range famKeys {
		members := families[k]
		h := sha256.Sum256([]byte(k))
		bucket := float64(h[0]) / 256.0
		if bucket < splitRatio {
			train = append(train, members...)
		} else {
			eval = append(eval, members...)
		}
	}
	return
}

// FrozenHoldoutSplit splits the corpus into a frozen holdout and a
// development pool. The frozen holdout must never be used for weight fitting,
// hyperparameter selection, threshold tuning, or generating near-duplicate
// training examples.
//
// The split is deterministic and identified by FrozenHoldoutHash.
func FrozenHoldoutSplit(exs []RouteExample) (frozen, development []RouteExample, holdoutHash string) {
	families := groupByFamily(exs)
	famKeys := sortedFamilyKeys(families)

	// Use a dedicated hash seed for the frozen split so changing the
	// train/eval ratio does not move the holdout.
	const splitSeed = "semantic-router-frozen-v1"
	for _, k := range famKeys {
		members := families[k]
		h := sha256.Sum256([]byte(splitSeed + ":" + k))
		bucket := float64(h[0]) / 256.0
		if bucket < DefaultFrozenRatio {
			frozen = append(frozen, members...)
		} else {
			development = append(development, members...)
		}
	}

	// Compute holdout hash from the group IDs in the frozen set.
	groupIDs := make(map[string]bool)
	for _, e := range frozen {
		groupIDs[e.SplitGroup] = true
	}
	ids := make([]string, 0, len(groupIDs))
	for id := range groupIDs {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	h := sha256.Sum256([]byte(hex.EncodeToString([]byte(strings.Join(ids, "|")))))
	holdoutHash = hex.EncodeToString(h[:16])

	return
}

// GroupedCVFolds splits a development pool into k folds, keeping all
// examples with the same SplitGroup in one fold. Returns k folds; each fold
// is the eval set, the rest are the training set.
func GroupedCVFolds(exs []RouteExample, k int) []CVFold {
	if k <= 1 {
		k = 5
	}
	families := groupByFamily(exs)
	famKeys := sortedFamilyKeys(families)

	// Assign families to folds round-robin for balance.
	foldFamilies := make([][]string, k)
	for i, key := range famKeys {
		foldFamilies[i%k] = append(foldFamilies[i%k], key)
	}

	folds := make([]CVFold, k)
	for i := 0; i < k; i++ {
		evalGroups := make(map[string]bool)
		for _, g := range foldFamilies[i] {
			evalGroups[g] = true
		}
		var train, eval []RouteExample
		for _, e := range exs {
			if evalGroups[e.SplitGroup] {
				eval = append(eval, e)
			} else {
				train = append(train, e)
			}
		}
		folds[i] = CVFold{Fold: i, Train: train, Eval: eval}
	}
	return folds
}

// CVFold — one fold of a grouped cross-validation split.
type CVFold struct {
	Fold  int
	Train []RouteExample
	Eval  []RouteExample
}

// FamilyID derives a stable family ID from a base source ID and a transform
// name, ensuring contrastive variants share the base's family.
func FamilyID(baseSourceID, transform string) string {
	if transform == "" {
		return baseSourceID
	}
	h := sha256.Sum256([]byte(baseSourceID + ":" + transform))
	return "family:" + hex.EncodeToString(h[:8])
}

// groupByFamily groups examples by SplitGroup.
func groupByFamily(exs []RouteExample) map[string][]RouteExample {
	families := make(map[string][]RouteExample)
	for _, e := range exs {
		key := e.SplitGroup
		if key == "" {
			key = e.SourceID
		}
		families[key] = append(families[key], e)
	}
	return families
}

// sortedFamilyKeys returns the sorted keys of a families map.
func sortedFamilyKeys(families map[string][]RouteExample) []string {
	keys := make([]string, 0, len(families))
	for k := range families {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}