package agent

import (
	"encoding/json"
	"iter"
	"strings"

	"nub/internal/llm"
	"nub/internal/tool"
)

type blockAcc struct {
	kind llm.BlockKind
	id   string
	name string
	text strings.Builder
	args strings.Builder
}

// consume liest den Event-Stream eines Turns, emittiert Deltas als UIEvents
// und akkumuliert parallel die vollständigen Blöcke der Assistant-Message
// (Ablauf pro Turn, Schritt 2).
func (l *Loop) consume(events iter.Seq2[llm.Event, error], out chan<- tool.UIEvent) (llm.Message, llm.StopReason, llm.Usage, error) {
	blocks := map[int]*blockAcc{}
	var order []int
	stop := llm.StopEnd
	var usage llm.Usage
	var streamErr error

	events(func(ev llm.Event, err error) bool {
		if err != nil {
			streamErr = err
			return false
		}
		switch e := ev.(type) {
		case llm.BlockStart:
			b := &blockAcc{kind: e.Block.Kind, id: e.Block.ID, name: e.Block.Name}
			blocks[e.Index] = b
			order = append(order, e.Index)
			if e.Block.Kind == llm.KindToolUse {
				out <- tool.ToolCallStart{ID: e.Block.ID, Name: e.Block.Name}
			}
		case llm.BlockDelta:
			b := blocks[e.Index]
			if b == nil {
				return true
			}
			switch b.kind {
			case llm.KindText:
				b.text.WriteString(e.Text)
				out <- tool.TextDelta{Text: e.Text}
			case llm.KindThinking:
				b.text.WriteString(e.Text)
				out <- tool.ThinkingDelta{Text: e.Text}
			case llm.KindToolUse:
				b.args.WriteString(e.PartialJSON)
			}
		case llm.Done:
			stop = e.Stop
			usage = e.Usage
		}
		return true
	})

	if streamErr != nil {
		return llm.Message{}, "", llm.Usage{}, streamErr
	}

	msg := llm.Message{Role: llm.RoleAssistant}
	for _, idx := range order {
		b := blocks[idx]
		switch b.kind {
		case llm.KindText, llm.KindThinking:
			if b.text.Len() > 0 {
				msg.Content = append(msg.Content, llm.Block{Kind: b.kind, Text: b.text.String()})
			}
		case llm.KindToolUse:
			input := json.RawMessage(b.args.String())
			if len(input) == 0 {
				input = json.RawMessage("{}")
			}
			msg.Content = append(msg.Content, llm.Block{Kind: llm.KindToolUse, ID: b.id, Name: b.name, Input: input})
		}
	}
	return msg, stop, usage, nil
}