package trace

import (
	"errors"
	"sync"
	"time"
)

// ErrSpanFinished is returned by Span.Record and Span.SetProp when called on a
// span that has already been closed by Finish. A closed span is immutable.
var ErrSpanFinished = errors.New("trace: span already finished")

// SpanKind classifies what a span represents so the M9 panel can group and color
// nodes by stage of the run. The values mirror the span kinds listed in
// DevDocs/design/04-运行时.md §9.
type SpanKind string

const (
	// SpanKindLLM is a single LLM invocation.
	SpanKindLLM SpanKind = "llm_call"
	// SpanKindTool is a tool/function call.
	SpanKindTool SpanKind = "tool_call"
	// SpanKindOrchestration is an orchestration node, e.g. a plan step.
	SpanKindOrchestration SpanKind = "orchestration"
	// SpanKindStateTransition is a runtime state transition.
	SpanKindStateTransition SpanKind = "state_transition"
	// SpanKindSubagent is a spawned sub-agent.
	SpanKindSubagent SpanKind = "subagent"
	// SpanKindMemory is a memory read or write.
	SpanKindMemory SpanKind = "memory"
	// SpanKindGuardrail is a guardrail input/output check.
	SpanKindGuardrail SpanKind = "guardrail"
)

// SpanStatus is the terminal outcome of a span.
type SpanStatus string

const (
	// SpanStatusOK marks a successfully completed span.
	SpanStatusOK SpanStatus = "ok"
	// SpanStatusError marks a failed span.
	SpanStatusError SpanStatus = "error"
)

// Event is a fine-grained record attached to a span, e.g. a token, retry or
// decision. Events enrich the chain-of-thought view.
type Event struct {
	// Time is when the event happened. Record sets it to now when left zero.
	Time time.Time `json:"time"`
	// Type is a short label identifying the event kind.
	Type string `json:"type"`
	// Props carries event-specific structured data.
	Props map[string]any `json:"props,omitempty"`
}

// Span is a single traced unit of work, safe for concurrent use.
//
// A span is created by (*Tracer).StartSpan, which also returns a context that
// carries the span as the current one. Runtime work then records detail with
// Record and SetProp, and closes the span once with Finish. ID, ParentID, Name
// and Kind are fixed at creation; Status, EndedAt and Duration are set when
// Finish runs. Reading a span's fields is safe after Finish returns.
type Span struct {
	mu     sync.RWMutex
	tracer *Tracer

	// ID is a span identifier unique across Tracers.
	ID string `json:"id"`
	// ParentID is the ID of the enclosing span, empty for a root span.
	ParentID string `json:"parent_id,omitempty"`
	// Name is a human-readable label for the work.
	Name string `json:"name"`
	// Kind classifies the span.
	Kind SpanKind `json:"kind"`
	// Status is the terminal outcome, valid after Finish.
	Status SpanStatus `json:"status"`
	// StartedAt is when the span began.
	StartedAt time.Time `json:"started_at"`
	// EndedAt is when Finish ran, zero until then.
	EndedAt time.Time `json:"ended_at,omitempty"`
	// Duration is EndedAt.Sub(StartedAt), zero until Finish.
	Duration time.Duration `json:"duration_ns"`
	// Props carries key/value metadata attached via SetProp.
	Props map[string]any `json:"props,omitempty"`
	// Events is the ordered list of events recorded via Record.
	Events []Event `json:"events,omitempty"`
}

// Record appends an event to the span. A zero Timestamp is stamped with now.
// Record returns ErrSpanFinished when the span has already been closed.
func (s *Span) Record(evt Event) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.EndedAt.IsZero() {
		return ErrSpanFinished
	}
	if evt.Time.IsZero() {
		evt.Time = time.Now()
	}
	s.Events = append(s.Events, evt)
	return nil
}

// SetProp attaches a key/value property to the span for structured output.
// SetProp returns ErrSpanFinished when the span has already been closed.
func (s *Span) SetProp(key string, value any) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.EndedAt.IsZero() {
		return ErrSpanFinished
	}
	s.Props[key] = value
	return nil
}

// Finish closes the span with the given status, stamping EndedAt and Duration
// and handing the completed span to the owning Tracer's Writer. Finish is
// idempotent: calls after the first are ignored.
func (s *Span) Finish(status SpanStatus) {
	s.mu.Lock()
	if !s.EndedAt.IsZero() {
		s.mu.Unlock()
		return
	}
	s.Status = status
	s.EndedAt = time.Now()
	s.Duration = s.EndedAt.Sub(s.StartedAt)
	rec := s.recordLocked()
	s.mu.Unlock()

	if s.tracer != nil {
		s.tracer.emit(rec)
	}
}

// recordLocked snapshots the span as a SpanRecord. The caller must hold s.mu.
func (s *Span) recordLocked() SpanRecord {
	rec := SpanRecord{
		ID:        s.ID,
		ParentID:  s.ParentID,
		Name:      s.Name,
		Kind:      s.Kind,
		Status:    s.Status,
		StartedAt: s.StartedAt,
		EndedAt:   s.EndedAt,
		Duration:  s.Duration,
	}
	if len(s.Props) > 0 {
		rec.Props = make(map[string]any, len(s.Props))
		for k, v := range s.Props {
			rec.Props[k] = v
		}
	}
	if len(s.Events) > 0 {
		rec.Events = make([]Event, len(s.Events))
		copy(rec.Events, s.Events)
	}
	return rec
}

// Snapshot returns an immutable copy of a finished span as a SpanRecord. It is
// only safe to call after Finish returns (a closed span is immutable), and it
// lets a consumer read the span's final state without holding its lock. Calling
// it on a non-finished span returns a record with a zero Status/EndedAt.
func (s *Span) Snapshot() SpanRecord {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.recordLocked()
}

// SpanRecord is an immutable, lock-free snapshot of a completed span, produced
// by Finish and passed to a Writer for persistence. It is safe to copy.
type SpanRecord struct {
	ID        string         `json:"id"`
	ParentID  string         `json:"parent_id,omitempty"`
	Name      string         `json:"name"`
	Kind      SpanKind       `json:"kind"`
	Status    SpanStatus     `json:"status"`
	StartedAt time.Time      `json:"started_at"`
	EndedAt   time.Time      `json:"ended_at,omitempty"`
	Duration  time.Duration  `json:"duration_ns"`
	Props     map[string]any `json:"props,omitempty"`
	Events    []Event        `json:"events,omitempty"`
}