package trace

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"
)

type span struct {
	name string
	dur  time.Duration
}

type Trace struct {
	mu    sync.Mutex
	spans []span
}

type ctxKey struct{}

func New(ctx context.Context) (context.Context, *Trace) {
	t := &Trace{}
	return context.WithValue(ctx, ctxKey{}, t), t
}

func Start(ctx context.Context, name string) func() {
	start := time.Now()
	return func() {
		Add(ctx, name, time.Since(start))
	}
}

func Add(ctx context.Context, name string, dur time.Duration) {
	t, _ := ctx.Value(ctxKey{}).(*Trace)
	if t == nil {
		return
	}
	t.mu.Lock()
	t.spans = append(t.spans, span{name: name, dur: dur})
	t.mu.Unlock()
}

func (t *Trace) Header(total time.Duration) string {
	t.mu.Lock()
	defer t.mu.Unlock()
	parts := []string{fmt.Sprintf("total;dur=%.1f", ms(total))}
	for _, s := range t.spans {
		parts = append(parts, fmt.Sprintf("%s;dur=%.1f", token(s.name), ms(s.dur)))
	}
	return strings.Join(parts, ", ")
}

func (t *Trace) String() string {
	t.mu.Lock()
	defer t.mu.Unlock()
	parts := []string{}
	for _, s := range t.spans {
		parts = append(parts, fmt.Sprintf("%s %s", s.name, s.dur.Round(time.Millisecond)))
	}
	return strings.Join(parts, ", ")
}

func ms(d time.Duration) float64 {
	return float64(d.Microseconds()) / 1000
}

func token(name string) string {
	return strings.ReplaceAll(name, " ", "-")
}