package validators

import (
	"encoding/json"
	"fmt"

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/chunking"
)

// Compact JSON fragments in this budget begin and end with non-whitespace.
// When fragments join, their touching word runs merge. Counting only extra
// runs inside each fragment makes costs additive, with one run for the whole
// payload and punctuation rounded only after assembling the array.
type validationCost struct {
	extraRuns   int
	punctuation int
}

func (c validationCost) add(other validationCost) validationCost {
	return validationCost{c.extraRuns + other.extraRuns, c.punctuation + other.punctuation}
}

func (c validationCost) arrayTokens() int {
	return 1 + c.extraRuns + (c.punctuation+2)/2 // opening and closing brackets
}

func measureJSON(estimator *chunking.SimpleTokenEstimator, fragment []byte) validationCost {
	words, punctuation := estimator.CountComponents(string(fragment))
	return validationCost{words - 1, punctuation}
}

func measureValidationItem(estimator *chunking.SimpleTokenEstimator, item LLMValidationItem) (validationCost, error) {
	// Every positive local reference has the same cost, regardless of digit count.
	item.CorrectionIndex = 1
	payload, err := json.Marshal(item)
	if err != nil {
		return validationCost{}, fmt.Errorf("marshal batch payload: %w", err)
	}
	return measureJSON(estimator, payload), nil
}

func measureContextPrefixes[T any](estimator *chunking.SimpleTokenEstimator, field string, entries []T) ([]validationCost, error) {
	prefixes := make([]validationCost, len(entries)+1)
	cost := measureJSON(estimator, []byte(`,"`+field+`":[]`))
	for i, entry := range entries {
		payload, err := json.Marshal(entry)
		if err != nil {
			return nil, fmt.Errorf("marshal validation %s: %w", field, err)
		}
		if i > 0 {
			cost.punctuation++ // separator between entries
		}
		cost = cost.add(measureJSON(estimator, payload))
		prefixes[i+1] = cost
	}
	return prefixes, nil
}

func fitSimpleValidationContext(item LLMValidationItem, budget int, estimator *chunking.SimpleTokenEstimator) (LLMValidationItem, validationCost, error) {
	required := item
	required.NeighboringSegments, required.GlossaryEntries = nil, nil
	base, err := measureValidationItem(estimator, required)
	if err != nil {
		return item, validationCost{}, err
	}
	if len(item.NeighboringSegments)+len(item.GlossaryEntries) == 0 {
		return item, base, nil
	}
	neighbors, err := measureContextPrefixes(estimator, "neighboring_segments", item.NeighboringSegments)
	if err != nil {
		return item, validationCost{}, err
	}
	glossary, err := measureContextPrefixes(estimator, "glossary_entries", item.GlossaryEntries)
	if err != nil {
		return item, validationCost{}, err
	}
	n, g := len(item.NeighboringSegments), len(item.GlossaryEntries)
	cost := base.add(neighbors[n]).add(glossary[g])
	for n > 0 && cost.arrayTokens() > budget {
		n--
		cost = base.add(neighbors[n]).add(glossary[g])
	}
	for g > 0 && cost.arrayTokens() > budget {
		g--
		cost = base.add(glossary[g])
	}
	item.NeighboringSegments = item.NeighboringSegments[:n]
	item.GlossaryEntries = item.GlossaryEntries[:g]
	return item, cost, nil
}