package semantic

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

// 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
	}

	// Group by split family.
	families := make(map[string][]RouteExample)
	for _, e := range exs {
		key := e.SplitGroup
		if key == "" {
			key = e.SourceID
		}
		families[key] = append(families[key], e)
	}

	// Sort families for determinism.
	famKeys := make([]string, 0, len(families))
	for k := range families {
		famKeys = append(famKeys, k)
	}
	sort.Strings(famKeys)

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

// 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])
}