package dagui

import (
	"context"
	"iter"
	"slices"
	"sort"
	"time"

	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	otellog "go.opentelemetry.io/otel/log"
	sdklog "go.opentelemetry.io/otel/sdk/log"
	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
	"go.opentelemetry.io/otel/sdk/metric/metricdata"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/trace"

	"github.com/dagger/dagger/dagql/call/callpbv1"
	"github.com/dagger/dagger/engine/slog"
	telemetry "github.com/dagger/otel-go"
)

type DB struct {
	PrimarySpan SpanID
	PrimaryLogs map[SpanID][]sdklog.Record

	Epoch, End time.Time

	Spans    *OrderedSet[SpanID, *Span]
	RootSpan *Span

	Resources map[attribute.Distinct]*resource.Resource

	CallPayloads map[string]string
	Calls        map[string]*callpbv1.Call

	Outputs   map[string]map[string]struct{}
	OutputOf  map[string]map[string]struct{}
	Intervals map[string]map[time.Time]*Span

	CreatorSpans map[string]SpanSet

	// Map of call digest -> metric name -> data points
	// NOTE: this is hard coded for Gauge int64 metricdata essentially right now,
	// needs generalization as more metric types get added
	MetricsByCall map[string]map[string][]metricdata.DataPoint[int64]
	MetricsBySpan map[SpanID]map[string][]metricdata.DataPoint[int64]

	// updatedSpans is a set of spans that have been updated since the last
	// sync, which includes any parent spans whose overall active time intervals
	// or status were modified via a child or linked span.
	updatedSpans SpanSet

	// seenSpans keeps track of which spans have been observed via
	// UpdatedSnapshots so that we can know whether we need to send them when we
	// finally see them
	seenSpans map[SpanID]struct{}

	pendingResumeOutputs map[resumeOutputKey]SpanSet
	pendingLogsByOutput  map[resumeOutputKey][]sdklog.Record
	resolvedLogsBySpan   map[SpanID][]sdklog.Record

	testIndex *TestIndex
}

type resumeOutputKey struct {
	TraceID TraceID
	Output  string
}

func NewDB() *DB {
	return &DB{
		PrimaryLogs: make(map[SpanID][]sdklog.Record),

		Spans:     NewSpanSet(),
		Resources: make(map[attribute.Distinct]*resource.Resource),

		CallPayloads: make(map[string]string),
		Calls:        make(map[string]*callpbv1.Call),

		OutputOf:  make(map[string]map[string]struct{}),
		Outputs:   make(map[string]map[string]struct{}),
		Intervals: make(map[string]map[time.Time]*Span),

		CreatorSpans: make(map[string]SpanSet),

		updatedSpans: NewSpanSet(),
		seenSpans:    make(map[SpanID]struct{}),

		pendingResumeOutputs: make(map[resumeOutputKey]SpanSet),
		pendingLogsByOutput:  make(map[resumeOutputKey][]sdklog.Record),
		resolvedLogsBySpan:   make(map[SpanID][]sdklog.Record),
	}
}

func (db *DB) seen(spanID SpanID) {
	db.seenSpans[spanID] = struct{}{}
}

func (db *DB) hasSeen(spanID SpanID) bool {
	_, seen := db.seenSpans[spanID]
	return seen
}

func (db *DB) UpdatedSnapshots(filter map[SpanID]bool) []SpanSnapshot {
	snapshots := snapshotSpans(db.updatedSpans.Order, func(span *Span) bool {
		if !span.Received {
			// don't send along any stubs; let the client-side create its own stubs
			return false
		}
		if filter == nil || filter[span.ParentID] {
			// include subscribed (or all) spans
			return true
		}
		if span.IsFailedOrCausedFailure() {
			// include failed spans so we can summarize them without having to
			// deep-dive.
			return true
		}
		if span.Reveal || len(span.RevealedSpans.Order) > 0 {
			// always include revealed spans and their parents
			return true
		}
		if span.Passthrough {
			// include any passthrough spans to ensure failures are collected.
			// the POST /query span for example never fails on its own.
			for _, child := range span.ChildSpans.Order {
				if child.IsFailedOrCausedFailure() {
					return true
				}
			}
		}
		return false
	})
	for spanID := range filter {
		span := db.Spans.Map[spanID]
		if span == nil {
			continue
		}
		if !db.hasSeen(spanID) {
			snapshots = append(snapshots, span.Snapshot())
		}
		for p := range span.Parents {
			if !db.hasSeen(p.ID) {
				snapshots = append(snapshots, p.Snapshot())
			}
		}
	}
	for _, snapshot := range snapshots {
		db.seen(snapshot.ID)
	}
	db.updatedSpans = NewSpanSet()
	return snapshots
}

func (db *DB) ImportSnapshots(snapshots []SpanSnapshot) {
	spans := make([]*Span, len(snapshots))
	for i, snapshot := range snapshots {
		span := db.findOrAllocSpan(snapshot.ID)
		span.Received = true
		snapshot.Version += span.Version // don't reset the version
		span.SpanSnapshot = snapshot
		db.integrateSpan(span)
		spans[i] = span
	}
	for _, span := range spans {
		span.PropagateStatusToParentsAndLinks()
	}
}

func (db *DB) update(span *Span) {
	db.noteTestSpanUpdated(span)
	if span.Final {
		// don't bump versions for final spans; leave the remote as the
		// source of truth, lest we stray forward and miss an actual version bump
		return
	}
	span.Version++
	db.updatedSpans.Add(span)
}

// Matches returns true if the span matches the filter, looking through
// Passthrough span parents until a match is found or a non-Passthrough span
// is reached.
func (span *Span) Matches(match func(*Span) bool) bool {
	if match(span) {
		return true
	}
	if span.ParentSpan != nil && span.ParentSpan.Passthrough {
		return span.ParentSpan.Matches(match)
	}
	return false
}

func snapshotSpans(spans []*Span, filter func(*Span) bool) []SpanSnapshot {
	var filtered []SpanSnapshot
	for _, span := range spans {
		if span.Matches(filter) {
			filtered = append(filtered, span.Snapshot())
		}
	}
	return filtered
}

func (db *DB) SpanSnapshots(id SpanID) []SpanSnapshot {
	snaps := snapshotSpans(db.Spans.Order, func(span *Span) bool {
		return span.ParentID == id || span.ID == id
	})
	if span := db.Spans.Map[id]; span != nil {
		for p := range span.Parents {
			if !db.hasSeen(p.ID) {
				snaps = append(snaps, p.Snapshot())
			}
		}
	}
	for _, snapshot := range snaps {
		db.seen(snapshot.ID)
	}
	return snaps
}

func (db *DB) RemainingSnapshots() []SpanSnapshot {
	return snapshotSpans(db.Spans.Order, func(span *Span) bool {
		return span.Received && !db.hasSeen(span.ID)
	})
}

var _ sdktrace.SpanExporter = (*DB)(nil)

func (db *DB) ExportSpans(ctx context.Context, otelSpans []sdktrace.ReadOnlySpan) error {
	spans := make([]*Span, len(otelSpans))
	for i, otelSpan := range otelSpans {
		spans[i] = db.recordOTelSpan(otelSpan)
	}
	for _, span := range spans {
		span.PropagateStatusToParentsAndLinks()
	}
	return nil
}

func (db *DB) LogExporter() sdklog.Exporter {
	return DBLogExporter{db}
}

type DBLogExporter struct {
	*DB
}

func (db DBLogExporter) Export(ctx context.Context, logs []sdklog.Record) error {
	for _, log := range logs {
		if log.Body().AsString() == "" {
			// eof; ignore
			continue
		}
		spanID, pendingKey := db.routeLog(log)
		if pendingKey != nil {
			db.pendingLogsByOutput[*pendingKey] = append(db.pendingLogsByOutput[*pendingKey], log)
			continue
		}
		if spanID == db.PrimarySpan {
			// buffer raw logs so we can replay them later
			db.PrimaryLogs[spanID] = append(db.PrimaryLogs[spanID], log)
		}
		// flag that the span has received logs
		db.initSpan(spanID).HasLogs = true
	}
	return nil
}

func (db *DB) Shutdown(ctx context.Context) error {
	return nil // noop
}

func (db *DB) ForceFlush(ctx context.Context) error {
	return nil // noop
}

func (db *DB) MetricExporter() sdkmetric.Exporter {
	return DBMetricExporter{db}
}

func (db *DB) Temporality(sdkmetric.InstrumentKind) metricdata.Temporality {
	return metricdata.DeltaTemporality
}

func (db *DB) Aggregation(sdkmetric.InstrumentKind) sdkmetric.Aggregation {
	return sdkmetric.AggregationDefault{}
}

type DBMetricExporter struct {
	*DB
}

func (db DBMetricExporter) Export(ctx context.Context, resourceMetrics *metricdata.ResourceMetrics) error {
	for _, scopeMetric := range resourceMetrics.ScopeMetrics {
		for _, metric := range scopeMetric.Metrics {
			// TODO: don't lose track of whether it's a Sum or a Gauge - it matters!
			if metricData, ok := metric.Data.(metricdata.Sum[int64]); ok {
				db.exportDataPoints(metric, metricData.DataPoints)
			}
			if metricData, ok := metric.Data.(metricdata.Gauge[int64]); ok {
				db.exportDataPoints(metric, metricData.DataPoints)
			}
		}
	}

	return nil
}

func (db DBMetricExporter) exportDataPoints(metric metricdata.Metrics, dataPoints []metricdata.DataPoint[int64]) {
	for _, point := range dataPoints {
		var metricsByName map[string][]metricdata.DataPoint[int64]
		if callDigest, ok := point.Attributes.Value(telemetry.DagDigestAttr); ok {
			if db.MetricsByCall == nil {
				db.MetricsByCall = make(map[string]map[string][]metricdata.DataPoint[int64])
			}
			var ok bool
			metricsByName, ok = db.MetricsByCall[callDigest.AsString()]
			if !ok {
				metricsByName = make(map[string][]metricdata.DataPoint[int64])
				db.MetricsByCall[callDigest.AsString()] = metricsByName
			}
		} else if spanIDHex, ok := point.Attributes.Value(telemetry.MetricsSpanIDAttr); ok {
			var spanID SpanID
			var err error
			if spanID.SpanID, err = trace.SpanIDFromHex(spanIDHex.AsString()); err != nil {
				continue
			}
			if db.MetricsBySpan == nil {
				db.MetricsBySpan = make(map[SpanID]map[string][]metricdata.DataPoint[int64])
			}
			var ok bool
			metricsByName, ok = db.MetricsBySpan[spanID]
			if !ok {
				metricsByName = make(map[string][]metricdata.DataPoint[int64])
				db.MetricsBySpan[spanID] = metricsByName
			}
		} else {
			continue
		}

		metricsByName[metric.Name] = append(metricsByName[metric.Name], point)
	}
}

// SetPrimarySpan allows the primary span to be explicitly set to a particular
// span. normally we assume the root span is the primary span, but in a nested
// scenario we never actually see the root span, so the CLI explicitly sets it
// to the span it created.
func (db *DB) SetPrimarySpan(span SpanID) {
	db.PrimarySpan = span
}

func (db *DB) initSpan(spanID SpanID) *Span {
	spanData, found := db.Spans.Map[spanID]
	if !found {
		spanData = db.newSpan(spanID)
		db.Spans.Add(spanData)
	}
	return spanData
}

func (db *DB) findOrAllocSpan(spanID SpanID) *Span {
	spanData, found := db.Spans.Map[spanID]
	if found {
		return spanData
	}
	return db.newSpan(spanID)
}

func (db *DB) newSpan(spanID SpanID) *Span {
	// TODO: this fools things into thinking they're a root span...?
	return &Span{
		SpanSnapshot: SpanSnapshot{
			ID: spanID,
		},
		ChildSpans:      NewSpanSet(),
		RunningSpans:    NewSpanSet(),
		RevealedSpans:   NewSpanSet(),
		FailedLinks:     NewSpanSet(),
		CanceledLinks:   NewSpanSet(),
		ErrorOrigins:    NewSpanSet(),
		causesViaLinks:  NewSpanSet(),
		effectsViaLinks: NewSpanSet(),
		db:              db,
	}
}

func (db *DB) recordOTelSpan(span sdktrace.ReadOnlySpan) *Span {
	spanID := SpanID{span.SpanContext().SpanID()}

	// mark the span as updated so we sync it to the frontend,
	// if this database is running on a backend
	//
	// NOTE: any updates to *other* spans should also be marked! for example, if
	// we update the status or active time of any linked/parent spans, those
	// should go in here too.

	// create or update the span itself
	spanData := db.findOrAllocSpan(spanID)
	spanData.Received = true
	spanData.TraceID = TraceID{span.SpanContext().TraceID()}
	spanData.ParentID.SpanID = span.Parent().SpanID()
	spanData.Name = span.Name()
	spanData.StartTime = span.StartTime()
	spanData.EndTime = span.EndTime()
	spanData.Status = span.Status()
	spanData.Links = make([]SpanLink, len(span.Links()))
	for i, link := range span.Links() {
		var purpose string
		for _, linkAttr := range link.Attributes {
			if linkAttr.Key == telemetry.LinkPurposeAttr {
				purpose = linkAttr.Value.AsString()
				break
			}
		}
		spanData.Links[i] = SpanLink{
			SpanContext: SpanContext{
				TraceID: TraceID{link.SpanContext.TraceID()},
				SpanID:  SpanID{link.SpanContext.SpanID()},
			},
			Purpose: purpose,
		}
	}

	if resource := span.Resource(); resource != nil {
		db.Resources[resource.Equivalent()] = resource
	}

	// populate snapshot from otel attributes
	for _, attr := range span.Attributes() {
		spanData.ProcessAttribute(
			string(attr.Key),
			attr.Value.AsInterface(),
		)
	}

	// integrate the span's data into the DB's live objects
	db.integrateSpan(spanData)
	return spanData
}

type Activity struct {
	CompletedIntervals []Interval
	EarliestRunning    time.Time

	// Keep track of the full set of running spans so we can update
	// EarliestRunning as they complete.
	//
	// This needs to be synced to the frontend so it doesn't lose track of the
	// running status in Activity.Add. We exclude from JSON marshalling since
	// the map key is incompatible. Syncing to the frontend uses encoding/gob,
	// which accepts the map key type.
	AllRunning map[SpanID]time.Time `json:"-"`
}

func (activity *Activity) Intervals(now time.Time) iter.Seq[Interval] {
	return func(yield func(Interval) bool) {
		var latestEnd time.Time
		yieldRunning := func() {
			runningStart := activity.EarliestRunning
			if latestEnd.After(runningStart) {
				runningStart = latestEnd
			}
			if runningStart.Before(now) {
				yield(Interval{Start: runningStart, End: now})
			}
		}
		for _, ival := range activity.CompletedIntervals {
			if !activity.EarliestRunning.IsZero() &&
				activity.EarliestRunning.Before(ival.Start) {
				yieldRunning()
				return
			}
			if !yield(ival) {
				return
			}
			if ival.End.After(latestEnd) {
				latestEnd = ival.End
			}
		}
		if !activity.EarliestRunning.IsZero() {
			yieldRunning()
		}
	}
}

func (activity *Activity) Duration(now time.Time) time.Duration {
	var dur time.Duration
	for ival := range activity.Intervals(now) {
		dur += ival.End.Sub(ival.Start)
	}
	return dur
}

type Interval struct {
	Start time.Time
	End   time.Time
}

func (activity *Activity) Add(span *Span) bool {
	var changed bool

	if span.IsRunning() {
		if activity.AllRunning == nil {
			activity.AllRunning = map[SpanID]time.Time{}
		}
		if _, found := activity.AllRunning[span.ID]; !found {
			activity.AllRunning[span.ID] = span.StartTime
			changed = true
		}
		// O(1): just check if this span is earlier than current earliest.
		if activity.EarliestRunning.IsZero() || span.StartTime.Before(activity.EarliestRunning) {
			activity.EarliestRunning = span.StartTime
			changed = true
		}
		return changed
	}

	wasEarliest := span.StartTime.Equal(activity.EarliestRunning)
	delete(activity.AllRunning, span.ID)
	if len(activity.AllRunning) == 0 {
		if !activity.EarliestRunning.IsZero() {
			activity.EarliestRunning = time.Time{}
			changed = true
		}
	} else if wasEarliest {
		// Only rescan if we removed the earliest — O(R) but rare.
		activity.EarliestRunning = time.Time{}
		for _, t := range activity.AllRunning {
			if activity.EarliestRunning.IsZero() || t.Before(activity.EarliestRunning) {
				activity.EarliestRunning = t
			}
		}
		changed = true
	}

	ival := Interval{
		Start: span.StartTime,
		End:   span.EndTime,
	}

	if len(activity.CompletedIntervals) == 0 {
		activity.CompletedIntervals = append(activity.CompletedIntervals, ival)
		changed = true
		return changed
	}

	idx, _ := slices.BinarySearchFunc(activity.CompletedIntervals, ival, func(a, b Interval) int {
		if a.Start.Before(b.Start) {
			return -1
		} else if a.Start.After(b.Start) {
			return 1
		} else {
			return 0
		}
	})

	// optimization: if the new interval is wholly subsumed by an existing
	// interval we can skip adding it. this is also handled by mergeIntervals,
	// but it's harder to return false after the fact.
	for _, existing := range activity.CompletedIntervals {
		if ival.Start.After(existing.Start) && ival.End.Before(existing.End) {
			return changed
		}
	}

	activity.CompletedIntervals = slices.Insert(activity.CompletedIntervals, idx, ival)
	activity.mergeIntervals()
	changed = true
	return changed
}

func (activity *Activity) IsRunning() bool {
	return !activity.EarliestRunning.IsZero()
}

func (activity *Activity) EndTimeOrFallback(now time.Time) time.Time {
	if !activity.EarliestRunning.IsZero() {
		return now
	}
	if len(activity.CompletedIntervals) == 0 {
		return time.Time{}
	}
	return activity.CompletedIntervals[len(activity.CompletedIntervals)-1].End
}

// mergeIntervals merges overlapping intervals in the activity.
func (activity *Activity) mergeIntervals() {
	// If there are no intervals, there's nothing to merge.
	if len(activity.CompletedIntervals) == 0 {
		return
	}

	// Keep track of the index of the last merged interval.
	lastIndex := 0
	for i := 1; i < len(activity.CompletedIntervals); i++ {
		ival := activity.CompletedIntervals[i]
		// If the current interval overlaps with the last one, merge them.
		if ival.Start.Before(activity.CompletedIntervals[lastIndex].End) {
			if ival.End.After(activity.CompletedIntervals[lastIndex].End) {
				// Extend the last interval.
				activity.CompletedIntervals[lastIndex].End = ival.End
			}
			// If ival is wholly subsumed, do nothing (continue).
		} else {
			// No overlap, move the lastIndex forward.
			lastIndex++
			activity.CompletedIntervals[lastIndex] = ival
		}
	}

	// Resize the slice to only include the merged intervals.
	activity.CompletedIntervals = activity.CompletedIntervals[:lastIndex+1]
}

// integrateSpan takes a possibly newly created span and updates
// database relationships and state
func (db *DB) integrateSpan(span *Span) { //nolint: gocyclo
	// track the span's own interval
	span.Activity.Add(span)
	db.update(span)

	// keep track of the time boundary
	if db.Epoch.IsZero() ||
		(!span.StartTime.IsZero() &&
			span.StartTime.Before(db.Epoch)) {
		db.Epoch = span.StartTime
	}
	if span.EndTime.After(db.End) {
		db.End = span.EndTime
	}

	// associate the span to its parent
	if span.ParentID.IsValid() {
		span.ParentSpan = db.initSpan(span.ParentID)
		if span.ParentSpan.ChildSpans.Add(span) {
			// if we're a new child, take a new snapshot for ChildCount
			db.update(span.ParentSpan)
		}
	}

	// associate the span to its links
	for _, linked := range span.Links {
		linkedCtx := linked.SpanContext
		switch linked.Purpose {
		case telemetry.LinkPurposeCause,
			// By default, links imply a causal relationship.
			//
			// Two reasons:
			// 1. Backward compatibility - this is how links are already interpreted.
			// 2. Span links are almost always representing a causal relationship
			// where the target of the link completed before the linking span.
			// (Otherwise the linking span could just be a child span.)
			"":
			linked := db.initSpan(linkedCtx.SpanID)
			linked.ChildSpans.Add(span)
			linked.effectsViaLinks.Add(span)
			span.causesViaLinks.Add(linked)
		case telemetry.LinkPurposeErrorOrigin:
			if linkedCtx.SpanID == span.ID {
				// defense in depth; it's technically possible to link to yourself, and
				// we don't want to double-init the span since that'll leave a zombie
				// span
				continue
			}
			linked := db.initSpan(linkedCtx.SpanID)
			span.ErrorOrigins.Add(linked)
		}
	}

	// Extract error origins from span error descriptions
	if span.Status.Code == codes.Error {
		for _, origin := range telemetry.ParseErrorOrigins(span.Status.Description) {
			originID := SpanID{SpanID: origin.SpanID()}
			if originID == span.ID {
				// renderStepError early-returns on any non-empty ErrorOrigins,
				// so a self-origin would suppress the leaf error message.
				continue
			}
			linked := db.initSpan(originID)
			span.ErrorOrigins.Add(linked)
		}
	}

	// keep track of intervals seen for a digest
	if span.CallDigest != "" {
		if db.Intervals[span.CallDigest] == nil {
			db.Intervals[span.CallDigest] = make(map[time.Time]*Span)
		}
		db.Intervals[span.CallDigest][span.StartTime] = span
	}

	if span.CallDigest != "" && span.CallPayload != "" {
		db.CallPayloads[span.CallDigest] = span.CallPayload
	}

	if !span.ParentID.IsValid() && span.Received {
		// keep track of the trace's root span
		if db.RootSpan == nil {
			db.RootSpan = span
		}

		if !db.PrimarySpan.IsValid() {
			// default primary to root span, though we might never see a "root
			// span" in a nested scenario.
			db.PrimarySpan = span.ID
		}

		if span == db.RootSpan && !span.IsRunning() {
			// If the root span is completed, we should mark any still-running spans
			// as canceled.
			for _, span := range db.Spans.Order {
				if span.IsRunning() {
					span.Canceled = true
					span.LeftRunning = true
					span.EndTime = db.RootSpan.EndTime
					span.PropagateStatusToParentsAndLinks()
					db.update(span)
				}
			}
		}
	} else if db.RootSpan != nil && !db.RootSpan.IsRunning() && span.IsRunning() {
		// Same as above (cancel running spans when root ends), but handled for
		// incoming span updates too.
		span.Canceled = true
		span.LeftRunning = true
		span.EndTime = db.RootSpan.EndTime
		span.PropagateStatusToParentsAndLinks()
	}

	if span.CallDigest != "" && span.Output != "" {
		// parent -> child
		if db.Outputs[span.CallDigest] == nil {
			db.Outputs[span.CallDigest] = make(map[string]struct{})
		}
		db.Outputs[span.CallDigest][span.Output] = struct{}{}

		// child -> parent
		if db.OutputOf[span.Output] == nil {
			db.OutputOf[span.Output] = make(map[string]struct{})
		}
		db.OutputOf[span.Output][span.CallDigest] = struct{}{}

		// output -> creator
		if db.CreatorSpans[span.Output] == nil {
			db.CreatorSpans[span.Output] = NewSpanSet()
		}
		db.CreatorSpans[span.Output].Add(span)

		db.resolvePendingResumeOutputs(span.Output, span.TraceID)
		db.resolvePendingLogs(span.Output, span.TraceID)
	}

	db.maybeResumeOutput(span)

	// finally, install the span if we don't already have it
	//
	// this dance is a little clumsy because we want to make sure parent spans
	// are inserted before their child spans, so we find-or-allocate the span but
	// aggressively initialize its parent span
	//
	// FIXME: refactor? can we keep some sort of flat map of spans an append
	// children to them instead of having the single big ordered list?
	db.Spans.Add(span)
	db.noteTestSpanUpdated(span)
}

func (db *DB) linkResumedOutput(span *Span, creator *Span) {
	changed := creator.ChildSpans.Add(span)
	creator.effectsViaLinks.Add(span)
	span.causesViaLinks.Add(creator)
	if changed {
		db.update(creator)
	}
}

func (db *DB) creatorSpanForDigestInTrace(dig string, traceID trace.TraceID) *Span {
	creators, ok := db.CreatorSpans[dig]
	if !ok {
		return nil
	}

	var best *Span
	for _, creator := range creators.Order {
		if creator.TraceID.TraceID != traceID {
			continue
		}
		if best == nil || creator.StartTime.After(best.StartTime) {
			best = creator
		}
	}

	return best
}

func (db *DB) maybeResumeOutput(span *Span) {
	if span.ResumeOutput == "" {
		return
	}

	creator := db.creatorSpanForDigestInTrace(span.ResumeOutput, span.TraceID.TraceID)
	if creator == nil {
		key := resumeOutputKey{
			TraceID: span.TraceID,
			Output:  span.ResumeOutput,
		}
		if db.pendingResumeOutputs[key] == nil {
			db.pendingResumeOutputs[key] = NewSpanSet()
		}
		db.pendingResumeOutputs[key].Add(span)
		return
	}

	db.linkResumedOutput(span, creator)
}

func (db *DB) resolvePendingResumeOutputs(output string, traceID TraceID) {
	key := resumeOutputKey{
		TraceID: traceID,
		Output:  output,
	}
	pending, ok := db.pendingResumeOutputs[key]
	if !ok {
		return
	}
	creator := db.creatorSpanForDigestInTrace(output, traceID.TraceID)
	if creator == nil {
		return
	}
	delete(db.pendingResumeOutputs, key)
	for _, span := range pending.Order {
		db.linkResumedOutput(span, creator)
		span.PropagateStatusToParentsAndLinks()
		db.update(span)
	}
}

func (db *DB) resolvePendingLogs(output string, traceID TraceID) {
	key := resumeOutputKey{
		TraceID: traceID,
		Output:  output,
	}
	pending, ok := db.pendingLogsByOutput[key]
	if !ok {
		return
	}
	creator := db.creatorSpanForDigestInTrace(output, traceID.TraceID)
	if creator == nil {
		return
	}
	delete(db.pendingLogsByOutput, key)
	for _, record := range pending {
		if creator.ID == db.PrimarySpan {
			db.PrimaryLogs[creator.ID] = append(db.PrimaryLogs[creator.ID], record)
		}
		db.initSpan(creator.ID).HasLogs = true
		db.resolvedLogsBySpan[creator.ID] = append(db.resolvedLogsBySpan[creator.ID], record)
	}
}

func (db *DB) DrainResolvedLogs(spanID SpanID) []sdklog.Record {
	logs := db.resolvedLogsBySpan[spanID]
	delete(db.resolvedLogsBySpan, spanID)
	return logs
}

func (db *DB) routeLog(record sdklog.Record) (SpanID, *resumeOutputKey) {
	fallback := SpanID{SpanID: record.SpanID()}

	var targetDig string
	record.WalkAttributes(func(kv otellog.KeyValue) bool {
		if kv.Key == telemetry.DagDigestAttr {
			targetDig = kv.Value.AsString()
			return false
		}
		return true
	})

	if targetDig == "" {
		return fallback, nil
	}

	if creator := db.creatorSpanForDigestInTrace(targetDig, record.TraceID()); creator != nil {
		return creator.ID, nil
	}

	key := resumeOutputKey{
		TraceID: TraceID{TraceID: record.TraceID()},
		Output:  targetDig,
	}
	return SpanID{}, &key
}

func (db *DB) LogTargetSpanID(record sdklog.Record) SpanID {
	spanID, _ := db.routeLog(record)
	return spanID
}

func (db *DB) HighLevelSpan(call *callpbv1.Call) *Span {
	return db.MostInterestingSpan(db.Simplify(call, false).Digest)
}

func (db *DB) MostInterestingSpan(dig string) *Span {
	var earliest *Span
	var earliestCached bool
	vs := make([]*Span, 0, len(db.Intervals[dig]))
	for _, span := range db.Intervals[dig] {
		vs = append(vs, span)
	}
	sort.Slice(vs, func(i, j int) bool {
		return vs[i].StartTime.Before(vs[j].StartTime)
	})
	for _, span := range vs {
		// a running vertex is always most interesting, and these are already in
		// order
		if span.IsRunningOrEffectsRunning() {
			return span
		}
		switch {
		case earliest == nil:
			// always show _something_
			earliest = span
			earliestCached = span.Cached
		case span.Cached:
			// don't allow a cached vertex to override a non-cached one
		case earliestCached:
			// unclear how this would happen, but non-cached versions are always more
			// interesting
			earliest = span
		case span.StartTime.Before(earliest.StartTime):
			// prefer the earliest active interval
			earliest = span
		}
	}
	return earliest
}

// func (db *DB) IsTransitiveDependency(dig, depDig string) bool {
// 	for _, v := range db.Intervals[dig] {
// 		for _, dig := range v.Inputs {
// 			if dig == depDig {
// 				return true
// 			}
// 			if db.IsTransitiveDependency(dig, depDig) {
// 				return true
// 			}
// 		}
// 		// assume they all have the same inputs
// 		return false
// 	}
// 	return false
// }

func (*DB) Close() error {
	return nil
}

func (db *DB) Call(dig string) *callpbv1.Call {
	// First, check if we already have the call cached
	if cached, ok := db.Calls[dig]; ok {
		return cached
	}

	// Next, try to decode from the call payload
	if callPayload, ok := db.CallPayloads[dig]; ok {
		var call callpbv1.Call
		if err := call.Decode(callPayload); err != nil {
			slog.Warn("failed to decode call", "err", err)
			// Cache nil so we don't keep retrying and spamming warnings
			// on every render cycle.
			db.Calls[dig] = nil
			return nil
		}
		// Cache the decoded call for future use
		db.Calls[dig] = &call
		return &call
	}

	// Finally, try to find the call through creator spans
	if creators, ok := db.CreatorSpans[dig]; ok {
		// Try each creator in order
		for _, creator := range creators.Order {
			if creatorCall := db.Call(creator.CallDigest); creatorCall != nil {
				return creatorCall
			}
		}
	}

	// No call found
	return nil
}

func (db *DB) MustCall(dig string) *callpbv1.Call {
	call := db.Call(dig)
	if call == nil {
		// Rather than blowing up, fill in an "easter egg."
		call = &callpbv1.Call{
			Field: "no",
			Type: &callpbv1.Type{
				NamedType: "Missing",
			},
			Args: []*callpbv1.Argument{
				{
					Name: "digest",
					Value: &callpbv1.Literal{
						Value: &callpbv1.Literal_String_{
							String_: dig,
						},
					},
				},
			},
			Digest: dig,
		}
	}
	return call
}

func (db *DB) Simplify(call *callpbv1.Call, force bool) *callpbv1.Call {
	creators, ok := db.CreatorSpans[call.Digest]
	if !ok {
		return call
	}

	for _, creator := range creators.Order {
		if creator.CallDigest == "" {
			continue
		}
		if creatorCall := db.Call(creator.CallDigest); creatorCall != nil {
			return creatorCall // TODO: re-simplify?
		}
	}

	return call
}

type WalkDecision int

const (
	WalkContinue WalkDecision = iota
	WalkSkip
	WalkPassthrough
	WalkStop
)

func WalkTree(tree []*TraceTree, f func(*TraceTree, int) WalkDecision) {
	var walk func([]*TraceTree, int)
	walk = func(rows []*TraceTree, depth int) {
		for _, row := range rows {
			switch f(row, depth) {
			case WalkContinue:
				walk(row.Children, depth+1)
			case WalkPassthrough:
				walk(row.Children, depth)
			case WalkSkip:
				continue
			case WalkStop:
				return
			}
		}
	}
	walk(tree, 0)
}

func (db *DB) FindResource(filter attribute.KeyValue) *resource.Resource {
	for _, res := range db.Resources {
		for _, kv := range res.Attributes() {
			if kv.Key == filter.Key && kv.Value == filter.Value {
				return res
			}
		}
	}
	return nil
}
