// Package trace provides structured deep tracing for the Cyrene runtime.
//
// It records work as a tree of spans — one span per LLM call, tool call,
// orchestration node, state transition, sub-agent, memory access or guardrail —
// so the run can be rendered later as a flow graph and a chain-of-thought view
// (DevDocs/design/04-运行时.md §9). This package is a leaf: it depends only on
// the standard library and log/slog, and never imports the runtime packages
// (loop/agent) that consume it.
//
// Usage: create a Tracer over a Writer, then StartSpan a root and nested spans.
// StartSpan threads a parent/child chain through the context it returns, so
// runtime code only needs to pass that context along. A span is closed with
// Finish(status); finished spans are handed to the Writer (here JSONRecorder)
// for structured local logging.
package trace

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"log/slog"
	"sync"
	"time"
)

// spanKey is the unexported context key that carries the current *Span.
type spanKey struct{}

// Tracer creates, chains and completes spans. It is safe for concurrent use.
//
// A nil Writer is tolerated: spans are still created and chained, but nothing
// is emitted on Finish. Construct with New.
type Tracer struct {
	mu       sync.RWMutex
	w        Writer
	seq      uint64
	idPrefix string
}

// New builds a Tracer that emits completed spans to w. Pass a Writer such as a
// *JSONRecorder to persist structured logs; pass nil to trace without emitting.
func New(w Writer) *Tracer {
	return &Tracer{w: w, idPrefix: randPrefix()}
}

// StartSpan begins a span named name of the given kind and returns it alongside
// a context that carries the span as the current one. If ctx already carries a
// span, the new span is chained as its child (its ParentID is set), enabling
// nested tracing; a nil ctx is treated as context.Background().
func (t *Tracer) StartSpan(ctx context.Context, name string, kind SpanKind) (context.Context, *Span) {
	if ctx == nil {
		ctx = context.Background()
	}
	parent := SpanFromContext(ctx)
	span := t.newSpan(name, kind, parent)
	return context.WithValue(ctx, spanKey{}, span), span
}

// SpanFromContext returns the span carried by ctx, or nil when ctx holds none
// (including a nil ctx).
func SpanFromContext(ctx context.Context) *Span {
	if ctx == nil {
		return nil
	}
	sp, _ := ctx.Value(spanKey{}).(*Span)
	return sp
}

// newSpan allocates a span, assigns it a unique ID, links ParentID and records
// its start time.
func (t *Tracer) newSpan(name string, kind SpanKind, parent *Span) *Span {
	t.mu.Lock()
	t.seq++
	id := fmt.Sprintf("%s-%d", t.idPrefix, t.seq)
	t.mu.Unlock()

	sp := &Span{
		tracer:    t,
		ID:        id,
		Name:      name,
		Kind:      kind,
		StartedAt: time.Now(),
		Props:     make(map[string]any),
	}
	if parent != nil {
		sp.ParentID = parent.ID
	}
	return sp
}

// emit hands a finished span to the writer. Write failures are logged and never
// propagated, since Finish cannot return an error. It is a no-op for a nil
// writer.
func (t *Tracer) emit(rec SpanRecord) {
	if t.w == nil {
		return
	}
	if err := t.w.WriteSpan(rec); err != nil {
		slog.Error("trace: write span failed", "span_id", rec.ID, "error", err)
	}
}

// randPrefix returns a short unique prefix used in span IDs so spans generated
// by different Tracers do not collide. It falls back to "tr" if the system
// entropy source is unavailable.
func randPrefix() string {
	var b [8]byte
	if _, err := rand.Read(b[:]); err != nil {
		return "tr"
	}
	return hex.EncodeToString(b[:])
}