// Package schemas defines the core schemas and types used by the Bifrost system.
package schemas

import (
	"context"
	"time"
)

// SpanHandle is an opaque handle to a span, implementation-specific.
// Different Tracer implementations can use their own concrete types.
type SpanHandle interface{}

// StreamAccumulatorResult contains the accumulated data from streaming chunks.
// This is the return type for tracer's streaming accumulation methods.
type StreamAccumulatorResult struct {
	RequestID             string                          // Request ID
	RequestedModel        string                          // Original model requested by the caller
	ResolvedModel         string                          // Actual model used by the provider (equals RequestedModel when no alias mapping exists)
	Provider              ModelProvider                   // Provider used
	Status                string                          // Status of the stream
	Latency               int64                           // Latency in milliseconds
	TimeToFirstToken      int64                           // Time to first token in milliseconds
	OutputMessage         *ChatMessage                    // Accumulated output message
	OutputMessages        []ResponsesMessage              // For responses API
	TokenUsage            *BifrostLLMUsage                // Token usage
	Cost                  *float64                        // Cost in dollars
	ErrorDetails          *BifrostError                   // Error details if any
	AudioOutput           *BifrostSpeechResponse          // For speech streaming
	TranscriptionOutput   *BifrostTranscriptionResponse   // For transcription streaming
	ImageGenerationOutput *BifrostImageGenerationResponse // For image generation streaming
	FinishReason          *string                         // Finish reason
	RawResponse           *string                         // Raw response
	RawRequest            interface{}                     // Raw request
}

// Tracer defines the interface for distributed tracing in Bifrost.
// Implementations can be injected via BifrostConfig to enable automatic instrumentation.
// The interface is designed to be minimal and implementation-agnostic.
type Tracer interface {
	// CreateTrace creates a new trace with optional parent ID and returns the trace ID.
	// The parentID can be extracted from W3C traceparent headers for distributed tracing.
	// The requestID is optional and can be used to identify the request.
	CreateTrace(parentID string, requestID ...string) string

	// EndTrace completes a trace and returns the trace data for observation/export.
	// After this call, the trace is removed from active tracking and returned for cleanup.
	// Returns nil if trace not found.
	EndTrace(traceID string) *Trace

	// StartSpan creates a new span as a child of the current span in context.
	// Returns updated context with new span and a handle for the span.
	// The context should be used for subsequent operations to maintain span hierarchy.
	StartSpan(ctx context.Context, name string, kind SpanKind) (context.Context, SpanHandle)

	// EndSpan completes a span with status and optional message.
	// Should be called when the operation represented by the span is complete.
	EndSpan(handle SpanHandle, status SpanStatus, statusMsg string)

	// SetAttribute sets an attribute on the span.
	// Attributes provide additional context about the operation.
	SetAttribute(handle SpanHandle, key string, value any)

	// AddEvent adds a timestamped event to the span.
	// Events represent discrete occurrences during the span's lifetime.
	AddEvent(handle SpanHandle, name string, attrs map[string]any)

	// PopulateLLMRequestAttributes populates all LLM-specific request attributes on the span.
	// This includes model parameters, input messages, temperature, max tokens, etc.
	PopulateLLMRequestAttributes(handle SpanHandle, req *BifrostRequest)

	// PopulateLLMResponseAttributes populates all LLM-specific response attributes on the span.
	// This includes output messages, tokens, usage stats, and error information if present.
	PopulateLLMResponseAttributes(ctx *BifrostContext, handle SpanHandle, resp *BifrostResponse, err *BifrostError)

	// StoreDeferredSpan stores a span handle for later completion (used for streaming requests).
	// The span handle is stored keyed by trace ID so it can be retrieved when the stream completes.
	StoreDeferredSpan(traceID string, handle SpanHandle)

	// GetDeferredSpanHandle retrieves a deferred span handle by trace ID.
	// Returns nil if no deferred span exists for the given trace ID.
	GetDeferredSpanHandle(traceID string) SpanHandle

	// ClearDeferredSpan removes the deferred span handle for a trace ID.
	// Should be called after the deferred span has been completed.
	ClearDeferredSpan(traceID string)

	// GetDeferredSpanID returns the span ID for the deferred span.
	// Returns empty string if no deferred span exists.
	GetDeferredSpanID(traceID string) string

	// AddStreamingChunk accumulates a streaming chunk for the deferred span.
	// Pass the full BifrostResponse to capture content, tool calls, reasoning, etc.
	// This is called for each streaming chunk to build up the complete response.
	AddStreamingChunk(traceID string, response *BifrostResponse)

	// GetAccumulatedChunks returns the accumulated response, TTFT, and chunk count for a deferred span.
	// The response is built from the streaming accumulator during the final ProcessStreamingChunk call.
	// Returns nil response if no plugin has called ProcessStreamingChunk (callers should nil-check).
	// Returns nil, 0, 0 if no accumulated data exists.
	GetAccumulatedChunks(traceID string) (response *BifrostResponse, ttftNs int64, chunkCount int)

	// CreateStreamAccumulator creates a new stream accumulator for the given trace ID.
	// This should be called at the start of a streaming request.
	CreateStreamAccumulator(traceID string, startTime time.Time)

	// CleanupStreamAccumulator removes the stream accumulator for the given trace ID.
	// This should be called after the streaming request is complete.
	CleanupStreamAccumulator(traceID string)

	// ProcessStreamingChunk processes a streaming chunk and accumulates it.
	// Returns the accumulated result. IsFinal will be true when the stream is complete.
	// This method is used by plugins to access accumulated streaming data.
	// The ctx parameter must contain the stream end indicator for proper final chunk detection.
	ProcessStreamingChunk(traceID string, isFinalChunk bool, result *BifrostResponse, err *BifrostError) *StreamAccumulatorResult

	// AttachPluginLogs appends plugin log entries to the trace identified by traceID.
	// Thread-safe. Should be called after plugin hooks complete, before trace completion.
	AttachPluginLogs(traceID string, logs []PluginLogEntry)

	// CompleteAndFlushTrace ends a trace, exports it to observability plugins, and
	// releases the trace resources. Used by transports that bypass normal HTTP trace completion.
	CompleteAndFlushTrace(traceID string)

	// Stop releases resources associated with the tracer.
	// Should be called during shutdown to stop background goroutines.
	Stop()
}

// NoOpTracer is a tracer that does nothing (default when tracing disabled).
// It satisfies the Tracer interface but performs no actual tracing operations.
type NoOpTracer struct{}

// CreateTrace returns an empty string (no trace created).
func (n *NoOpTracer) CreateTrace(_ string, _ ...string) string { return "" }

// EndTrace returns nil (no trace to end).
func (n *NoOpTracer) EndTrace(_ string) *Trace { return nil }

// StartSpan returns the context unchanged and a nil handle.
func (n *NoOpTracer) StartSpan(ctx context.Context, _ string, _ SpanKind) (context.Context, SpanHandle) {
	return ctx, nil
}

// EndSpan does nothing.
func (n *NoOpTracer) EndSpan(_ SpanHandle, _ SpanStatus, _ string) {}

// SetAttribute does nothing.
func (n *NoOpTracer) SetAttribute(_ SpanHandle, _ string, _ any) {}

// AddEvent does nothing.
func (n *NoOpTracer) AddEvent(_ SpanHandle, _ string, _ map[string]any) {}

// PopulateLLMRequestAttributes does nothing.
func (n *NoOpTracer) PopulateLLMRequestAttributes(_ SpanHandle, _ *BifrostRequest) {}

// PopulateLLMResponseAttributes does nothing.
func (n *NoOpTracer) PopulateLLMResponseAttributes(_ *BifrostContext, _ SpanHandle, _ *BifrostResponse, _ *BifrostError) {
}

// StoreDeferredSpan does nothing.
func (n *NoOpTracer) StoreDeferredSpan(_ string, _ SpanHandle) {}

// GetDeferredSpanHandle returns nil.
func (n *NoOpTracer) GetDeferredSpanHandle(_ string) SpanHandle { return nil }

// ClearDeferredSpan does nothing.
func (n *NoOpTracer) ClearDeferredSpan(_ string) {}

// GetDeferredSpanID returns empty string.
func (n *NoOpTracer) GetDeferredSpanID(_ string) string { return "" }

// AddStreamingChunk does nothing.
func (n *NoOpTracer) AddStreamingChunk(_ string, _ *BifrostResponse) {}

// GetAccumulatedChunks returns nil, 0, 0.
func (n *NoOpTracer) GetAccumulatedChunks(_ string) (*BifrostResponse, int64, int) { return nil, 0, 0 }

// CreateStreamAccumulator does nothing.
func (n *NoOpTracer) CreateStreamAccumulator(_ string, _ time.Time) {}

// CleanupStreamAccumulator does nothing.
func (n *NoOpTracer) CleanupStreamAccumulator(_ string) {}

// ProcessStreamingChunk returns nil.
func (n *NoOpTracer) ProcessStreamingChunk(_ string, _ bool, _ *BifrostResponse, _ *BifrostError) *StreamAccumulatorResult {
	return nil
}

// AttachPluginLogs does nothing.
func (n *NoOpTracer) AttachPluginLogs(_ string, _ []PluginLogEntry) {}

// CompleteAndFlushTrace does nothing.
func (n *NoOpTracer) CompleteAndFlushTrace(_ string) {}

// Stop does nothing.
func (n *NoOpTracer) Stop() {}

// DefaultTracer returns a no-op tracer for use when tracing is disabled.
func DefaultTracer() Tracer {
	return &NoOpTracer{}
}

// Ensure NoOpTracer implements Tracer at compile time
var _ Tracer = (*NoOpTracer)(nil)