package ai

import (
	"os"
	"sort"
	"strconv"

	"github.com/shopspring/decimal"
)

// BenchSuppressionThreshold is the small cell suppression floor k for aggregate structural answers about
// a tenant bench. Cells contributed to by fewer than k distinct individuals are suppressed. Set by the
// 2026-06-30 de identification research (NIST SP 800-188, FCSM WP22, l diversity, NYC LL144): k = 10 is
// the hard floor; the load bearing protection is the architecture, not k. l diversity is intentionally
// NOT computed here: the graph carries no protected attribute (the EEO module is unreachable), so there
// is no sensitive value for the homogeneity attack to infer. The heavyweight differencing defenses (a per
// querier DP epsilon budget, tracker and overlap refusal, calibrated noise) live on the cross tenant
// market surface (flag off today), not on this owner queries own data path.
const BenchSuppressionThreshold = 10

// minSuppressionThreshold is the floor below which k may never be configured (the research dissent's hard
// line). Any MM_BENCH_SUPPRESSION_K under this is ignored.
const minSuppressionThreshold = 5

// SuppressionSentinel is the single value free token rendered for every suppressed cell. It is identical
// for a true zero and any cell of 1 to k-1, so the token itself carries no count signal: a real 0 and a
// small nonzero are indistinguishable.
const SuppressionSentinel = "below reporting threshold"

// Cell is one aggregate cell in a bench answer: a generalized quasi identifier bucket and the count of
// DISTINCT individuals contributing to it. Value is an optional decimal aggregate (for example a median
// readiness); Suppress zeroes it on a suppressed cell so it carries no residual small cell signal, and
// RenderValue renders the sentinel for it. Cells of a single intent form a PARTITION of the matched set
// (each matched individual lands in exactly one bucket), so the counts sum to the matched total.
type Cell struct {
	// Count is json excluded: it is the raw headcount, retained on a suppressed cell only for the internal
	// partition sum invariant, and must reach a response ONLY through RenderCount (which returns the sentinel
	// when suppressed). Tagging it json:"-" means a suppressed sub k count can never leak even if a Cell or a
	// row that embeds it is marshaled directly by a future caller.
	Key        []string        `json:"key"`
	Count      int             `json:"-"`
	Value      decimal.Decimal `json:"value"`
	HasValue   bool            `json:"has_value"`
	Suppressed bool            `json:"suppressed"`
}

// ConfiguredSuppressionThreshold reads MM_BENCH_SUPPRESSION_K, clamps it to never drop below
// minSuppressionThreshold, and defaults to BenchSuppressionThreshold.
func ConfiguredSuppressionThreshold() int {
	k := BenchSuppressionThreshold
	if v := os.Getenv("MM_BENCH_SUPPRESSION_K"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n >= minSuppressionThreshold {
			k = n
		}
	}
	return k
}

// Suppress applies deterministic small cell suppression to one intent's partition of cells. The input
// order is preserved; only the Suppressed flag is set. The rules:
//
//  1. Refuse the whole answer when the matched total is below k: even the existence of the set is sub
//     threshold, so every cell is suppressed.
//  2. Primary: suppress any cell with Count < k (Count of 0 to k-1).
//  3. Complementary: if exactly one cell ends up suppressed by primary suppression, its value is
//     recoverable as the total minus the published cells, so suppress the smallest published cell as well.
//     This guarantees at least two hidden cells whenever primary suppression hides one among MULTIPLE
//     buckets, the FCSM Working Paper 22 secondary suppression rule. The rule 1 refuse branch may leave a
//     single hidden bucket when the whole answer is one sub k bucket; that is safe because there is no
//     sibling cell and Suppress emits no grand total to subtract against.
//
// Suppress is pure and deterministic: the complementary choice breaks ties by smallest count then by Key
// order, so the same input always yields the same suppression pattern. It is the sole authority on the
// Suppressed flag and clears any stale flags at entry, so a reused slice is handled safely.
func Suppress(cells []Cell, k int) []Cell {
	if k < minSuppressionThreshold {
		k = minSuppressionThreshold
	}
	// Clear any stale flags so the count is the only truth and a reused slice cannot keep an old decision.
	for i := range cells {
		cells[i].Suppressed = false
	}
	total := 0
	for _, c := range cells {
		total += c.Count
	}
	switch {
	case total < k:
		for i := range cells {
			cells[i].Suppressed = true
		}
	default:
		suppressedCount := 0
		for i := range cells {
			if cells[i].Count < k {
				cells[i].Suppressed = true
				suppressedCount++
			}
		}
		if suppressedCount == 1 {
			// Choose the smallest published cell to also hide (least utility lost), tie broken by Key order
			// so the pattern is deterministic. Index order of published cells.
			published := make([]int, 0, len(cells))
			for i := range cells {
				if !cells[i].Suppressed {
					published = append(published, i)
				}
			}
			sort.SliceStable(published, func(a, b int) bool {
				ia, ib := published[a], published[b]
				if cells[ia].Count != cells[ib].Count {
					return cells[ia].Count < cells[ib].Count
				}
				return keyLess(cells[ia].Key, cells[ib].Key)
			})
			if len(published) > 0 {
				cells[published[0]].Suppressed = true
			}
		}
	}
	// A suppressed cell must carry NO residual signal: clear its decimal aggregate so a caller that renders
	// Value cannot leak a small cell average or median. RenderValue also returns the sentinel; this is
	// defense in depth so the raw struct never holds a hidden value.
	for i := range cells {
		if cells[i].Suppressed {
			cells[i].Value = decimal.Zero
			cells[i].HasValue = false
		}
	}
	return cells
}

// keyLess is a total order over generalized bucket keys for deterministic tie breaking.
func keyLess(a, b []string) bool {
	for i := 0; i < len(a) && i < len(b); i++ {
		if a[i] != b[i] {
			return a[i] < b[i]
		}
	}
	return len(a) < len(b)
}

// RenderCount returns the displayable count for a cell: the sentinel when suppressed, else the number.
// The audit and the faithfulness narration use this so a suppressed count never leaves the process.
func RenderCount(c Cell) string {
	if c.Suppressed {
		return SuppressionSentinel
	}
	return strconv.Itoa(c.Count)
}

// RenderValue returns the displayable decimal aggregate for a cell: the sentinel when suppressed, an empty
// string when the cell carries no value, else the value. Like RenderCount, this keeps a suppressed
// aggregate from leaving the process; Suppress also zeroes Value on a suppressed cell as defense in depth.
func RenderValue(c Cell) string {
	if c.Suppressed {
		return SuppressionSentinel
	}
	if !c.HasValue {
		return ""
	}
	return c.Value.String()
}