package observation

import (
	"sync"

	"go.maximumdirect.net/promptkit/internal/domain"
)

// Group incrementally collects terminal facts for explicitly observed runs.
// Its zero value is ready; it must not be copied after first use.
type Group struct {
	mu                         sync.Mutex
	members                    map[string]bool
	terminals                  int
	attempts, retries, repairs uint64
	usage                      domain.ObservedTokenUsage
}

// Observe enrolls a run and incorporates its first terminal event only.
func (g *Group) Observe(event domain.ExecutionEvent) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.members == nil {
		g.members = make(map[string]bool)
	}
	terminal := g.members[event.RunID]
	if _, exists := g.members[event.RunID]; !exists {
		g.members[event.RunID] = false
	}
	if event.Kind != domain.ExecutionFinished || terminal {
		return
	}
	g.members[event.RunID] = true
	g.attempts += event.Attempts
	g.retries += event.TransientRetries
	g.repairs += event.RepairCalls
	if g.terminals == 0 {
		complete := domain.ObservedTokenCount{Complete: true}
		g.usage = domain.ObservedTokenUsage{PromptTokens: complete, CompletionTokens: complete, TotalTokens: complete, CachedTokens: complete, CacheWriteTokens: complete}
	}
	mergeGroupUsage(&g.usage, event.Usage)
	g.terminals++
}

// Snapshot returns a detached consistent view of currently observed members.
func (g *Group) Snapshot() domain.ExecutionGroupSummary {
	g.mu.Lock()
	defer g.mu.Unlock()
	result := domain.ExecutionGroupSummary{Members: make([]domain.ExecutionGroupMember, 0, len(g.members)), TerminalCoverageComplete: len(g.members) == g.terminals, Attempts: g.attempts, TransientRetries: g.retries, RepairCalls: g.repairs, Usage: g.usage}
	for id, terminal := range g.members {
		result.Members = append(result.Members, domain.ExecutionGroupMember{RunID: id, TerminalObserved: terminal})
	}
	if len(g.members) == 0 {
		zero := domain.ObservedTokenCount{Known: true, Complete: true}
		result.Usage = domain.ObservedTokenUsage{PromptTokens: zero, CompletionTokens: zero, TotalTokens: zero, CachedTokens: zero, CacheWriteTokens: zero}
	}
	return result
}

func mergeGroupUsage(total *domain.ObservedTokenUsage, next domain.ObservedTokenUsage) {
	merge := func(a *domain.ObservedTokenCount, b domain.ObservedTokenCount) {
		if b.Known {
			a.Value += b.Value
		}
		a.Known = a.Known || b.Known
		a.Complete = a.Complete && b.Complete
	}
	merge(&total.PromptTokens, next.PromptTokens)
	merge(&total.CompletionTokens, next.CompletionTokens)
	merge(&total.TotalTokens, next.TotalTokens)
	merge(&total.CachedTokens, next.CachedTokens)
	merge(&total.CacheWriteTokens, next.CacheWriteTokens)
}