package validators

import (
	"encoding/json"
	"fmt"

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

type LLMValidationBatch struct {
	BatchIndex int                 `json:"batch_index"`
	Items      []LLMValidationItem `json:"items"`
	TokenCount int                 `json:"token_count"`
}

func ChunkLLMValidationItems(items []LLMValidationItem, maxPromptTokens int, estimator chunking.TokenEstimator) ([]LLMValidationBatch, error) {
	if maxPromptTokens <= 0 {
		return nil, fmt.Errorf("validation max prompt tokens must be greater than zero")
	}
	if len(items) == 0 {
		return nil, fmt.Errorf("validation input must contain at least one proposal")
	}
	if estimator == nil {
		estimator = chunking.NewSimpleTokenEstimator()
	}

	batches := make([]LLMValidationBatch, 0)
	current := make([]LLMValidationItem, 0)
	currentTokens := 0

	for _, item := range items {
		singleTokens, err := estimateBatchTokens(estimator, []LLMValidationItem{item})
		if err != nil {
			return nil, err
		}
		if singleTokens > maxPromptTokens {
			return nil, fmt.Errorf("single validation proposal exceeds max prompt tokens")
		}

		candidate := append(append([]LLMValidationItem(nil), current...), item)
		candidateTokens, err := estimateBatchTokens(estimator, candidate)
		if err != nil {
			return nil, err
		}

		if len(current) > 0 && candidateTokens > maxPromptTokens {
			batches = append(batches, LLMValidationBatch{
				BatchIndex: len(batches),
				Items:      append([]LLMValidationItem(nil), current...),
				TokenCount: currentTokens,
			})
			current = []LLMValidationItem{item}
			currentTokens = singleTokens
			continue
		}

		current = candidate
		currentTokens = candidateTokens
	}

	if len(current) > 0 {
		batches = append(batches, LLMValidationBatch{
			BatchIndex: len(batches),
			Items:      append([]LLMValidationItem(nil), current...),
			TokenCount: currentTokens,
		})
	}

	return batches, nil
}

func estimateBatchTokens(estimator chunking.TokenEstimator, batch []LLMValidationItem) (int, error) {
	payload, err := json.Marshal(batch)
	if err != nil {
		return 0, fmt.Errorf("marshal batch payload: %w", err)
	}
	return estimator.EstimateTokens(string(payload)), nil
}