// Package trace provides a full-chain request tracing system.
// Each conversation request generates a Trace with multiple Spans
// capturing every intermediate step from message receipt to response delivery.
package trace

import (
	"sync"
	"time"
)

// Trace represents a single conversation request's full lifecycle.
type Trace struct {
	ID        string    `json:"id"`
	UserMsg   string    `json:"user_msg"`   // 用户说了什么（截取前80字）
	StartTime time.Time `json:"start_time"`
	EndTime   time.Time `json:"end_time,omitempty"`
	Status    string    `json:"status"`     // "running" | "success" | "error"
	Error     string    `json:"error,omitempty"`
	Spans     []Span    `json:"spans"`
}

// Span represents a single step within a trace.
type Span struct {
	Name     string        `json:"name"`      // "intent_analysis" | "llm_call" | etc.
	Label    string        `json:"label"`     // Display label
	StartMs  int64         `json:"start_ms"`  // Offset from trace start in ms
	Duration int64         `json:"duration_ms"`
	Status   string        `json:"status"`    // "running" | "success" | "error"
	Detail   string        `json:"detail"`    // Extra info (model name, token count, etc.)
	Children []Span        `json:"children,omitempty"` // Sub-spans
	Data     map[string]any `json:"data,omitempty"`
}

// Logger holds recent traces in a ring buffer.
type Logger struct {
	mu       sync.RWMutex
	traces   []*Trace
	max      int
	index    int
	running  map[string]*Trace // Currently running traces by ID
}

// NewLogger creates a trace logger with given capacity.
func NewLogger(capacity int) *Logger {
	return &Logger{
		traces:  make([]*Trace, capacity),
		max:     capacity,
		running: make(map[string]*Trace),
	}
}

// StartTrace begins a new trace for an incoming user message.
func (l *Logger) StartTrace(id, userMsg string) *Trace {
	t := &Trace{
		ID:        id,
		UserMsg:   truncateStr(userMsg, 80),
		StartTime: time.Now(),
		Status:    "running",
		Spans:     make([]Span, 0),
	}
	l.mu.Lock()
	l.running[id] = t
	l.mu.Unlock()
	return t
}

// AddSpan adds a span to a running trace.
func (l *Logger) AddSpan(traceID, name, label, status string, durationMs int64, detail string) {
	l.mu.Lock()
	defer l.mu.Unlock()

	t, ok := l.running[traceID]
	if !ok { return }

	span := Span{
		Name:      name,
		Label:     label,
		StartMs:   time.Since(t.StartTime).Milliseconds() - durationMs,
		Duration:  durationMs,
		Status:    status,
		Detail:    detail,
	}
	t.Spans = append(t.Spans, span)
}

// FinishTrace marks a trace as complete and stores it in the ring buffer.
func (l *Logger) FinishTrace(id, status, errMsg string) {
	l.mu.Lock()
	defer l.mu.Unlock()

	t, ok := l.running[id]
	if !ok { return }

	t.EndTime = time.Now()
	t.Status = status
	if errMsg != "" { t.Error = errMsg }
	delete(l.running, id)

	// Store in ring buffer
	l.traces[l.index%l.max] = t
	l.index++
}

// Recent returns the most recent N traces.
func (l *Logger) Recent(n int) []*Trace {
	l.mu.RLock()
	defer l.mu.RUnlock()

	if n > l.max { n = l.max }
	count := l.index
	if count > l.max { count = l.max }
	if n > count { n = count }

	result := make([]*Trace, n)
	for i := 0; i < n; i++ {
		idx := (l.index - 1 - i + l.max) % l.max
		if idx < l.max { result[n-1-i] = l.traces[idx] }
	}
	return result
}

// Running returns currently active traces.
func (l *Logger) Running() []*Trace {
	l.mu.RLock()
	defer l.mu.RUnlock()

	result := make([]*Trace, 0, len(l.running))
	for _, t := range l.running {
		result = append(result, t)
	}
	return result
}

// Clear removes all traces.
func (l *Logger) Clear() {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.traces = make([]*Trace, l.max)
	l.index = 0
	l.running = make(map[string]*Trace)
}

// SpanTimer is a convenience helper for timing a span.
type SpanTimer struct {
	logger   *Logger
	traceID  string
	name     string
	label    string
	start    time.Time
}

// StartSpan begins timing a span.
func (l *Logger) StartSpan(traceID, name, label string) *SpanTimer {
	return &SpanTimer{
		logger:  l,
		traceID: traceID,
		name:    name,
		label:   label,
		start:   time.Now(),
	}
}

// Done finishes the span and adds it to the trace.
func (st *SpanTimer) Done(status, detail string) {
	durationMs := time.Since(st.start).Milliseconds()
	st.logger.AddSpan(st.traceID, st.name, st.label, status, durationMs, detail)
}

// DoneMs finishes with a pre-computed duration.
func (st *SpanTimer) DoneMs(durationMs int64, status, detail string) {
	st.logger.AddSpan(st.traceID, st.name, st.label, status, durationMs, detail)
}

func truncateStr(s string, maxLen int) string {
	runes := []rune(s)
	if len(runes) <= maxLen { return s }
	return string(runes[:maxLen]) + "..."
}