package gateway

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"net/http"
)

var errRequestBodyTooLarge = errors.New("request body is too large")

func prepareTokenBudget(request *http.Request, maxBody int64) (int64, error) {
	if request.Body == nil || request.Method == http.MethodGet {
		return 0, nil
	}
	limited := io.LimitReader(request.Body, maxBody+1)
	body, err := io.ReadAll(limited)
	if err != nil {
		return 0, err
	}
	_ = request.Body.Close()
	if int64(len(body)) > maxBody {
		return 0, errRequestBodyTooLarge
	}
	restoreRequestBody(request, body)

	inputEstimate := max(int64((len(body)+3)/4), 1)
	if request.URL.Path == "/v1/embeddings" {
		return inputEstimate, nil
	}
	outputBudget := int64(1024)
	decoder := json.NewDecoder(bytes.NewReader(body))
	decoder.UseNumber()
	var payload map[string]any
	if decoder.Decode(&payload) == nil {
		for _, field := range []string{"max_output_tokens", "max_completion_tokens", "max_tokens"} {
			if value := jsonInt64(payload[field]); value > 0 {
				outputBudget = value
				break
			}
		}
	}
	// A per-request ceiling prevents malicious payloads from reserving unbounded Redis counters.
	return min(inputEstimate+outputBudget, int64(100_000_000)), nil
}