package runtime

import (
	agentllm "agent-native/runtime/llm"
	core "agent-system-core"
	"context"
	"fmt"
)

// Stop before a further request when reported total usage exhausts the run budget.
// A single in-flight request can exceed the remaining input allowance; output is capped.
func budgetedStream(stream agentllm.StreamFunc, budget core.RunBudget, outputCap int) agentllm.StreamFunc {
	if budget.MaxTokens == nil || *budget.MaxTokens <= 0 {
		return stream
	}
	limit := *budget.MaxTokens
	used := 0
	unknown := false
	return func(ctx context.Context, model string, messages []agentllm.Message, tools []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		left := limit - used
		if left <= 0 || unknown {
			out := make(chan agentllm.Event, 1)
			reason := "run token budget exhausted"
			if unknown {
				reason = "provider usage unavailable; further requests stopped"
			}
			out <- agentllm.Event{Kind: agentllm.KindError, Err: fmt.Errorf("budget_exceeded: %s", reason), ErrorMessage: reason}
			close(out)
			return out
		}
		if opts.MaxTokens <= 0 || opts.MaxTokens > left {
			opts.MaxTokens = left
		}
		if outputCap > 0 && opts.MaxTokens > outputCap {
			opts.MaxTokens = outputCap
		}
		in := stream(ctx, model, messages, tools, opts)
		out := make(chan agentllm.Event)
		go func() {
			defer close(out)
			reported := false
			for e := range in {
				if e.Kind == agentllm.KindMessageEnd {
					n := e.Usage.TotalTokens
					if n == 0 {
						n = e.Usage.PromptTokens + e.Usage.CompletionTokens
					}
					if n > 0 {
						used += n
						reported = true
					}
				}
				select {
				case out <- e:
				case <-ctx.Done():
					return
				}
			}
			if !reported {
				unknown = true
			}
		}()
		return out
	}
}