package runtimeevents

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"
)

// DataDecoder constructs an EventData from a raw JSON message for a given Kind.
// Register decoders via RegisterDataDecoder during init.
type DataDecoder func(kind Kind, raw json.RawMessage) (EventData, error)

var (
	decoderMu  sync.RWMutex
	decoderMap = make(map[Kind]DataDecoder)
)

func init() {
	// Session lifecycle
	for _, k := range []Kind{
		KindSessionStarting, KindSessionStarted,
		KindSessionSleeping, KindSessionStopped,
	} {
		RegisterDataDecoder(k, decodeLifecycleData)
	}
	RegisterDataDecoder(KindSessionStateChanged, nil)
	RegisterDataDecoder(KindAgentCatalogChanged, nil)

	// Wake
	for _, k := range []Kind{
		KindWakeReceived, KindWakeReplySent, KindWakeNoReply, KindWakeFailed,
	} {
		RegisterDataDecoder(k, decodeWakeData)
	}

	// Delivery
	for _, k := range []Kind{
		KindDeliveryAccepted, KindDeliveryDeferred, KindDeliveryPending,
		KindDeliveryAcked, KindDeliveryReplayed,
	} {
		RegisterDataDecoder(k, decodeLifecycleData)
	}

	// Runtime diagnostics
	for _, k := range []Kind{
		KindRuntimeStartFailed, KindRuntimeExitedUnexpectedly, KindRuntimeRetrying,
		KindRuntimeEventParseFailed, KindRuntimeResumeSafetyChanged,
	} {
		RegisterDataDecoder(k, decodeRuntimeDiagnosticData)
	}

	// Turn-level
	RegisterDataDecoder(KindTurnThinking, decodeTurnTextData)
	RegisterDataDecoder(KindTurnText, decodeTurnTextData)
	RegisterDataDecoder(KindTurnError, decodeTurnTextData)
	RegisterDataDecoder(KindTurnToolCall, decodeToolCallData)
	RegisterDataDecoder(KindTurnToolOutput, decodeToolCallData)
	RegisterDataDecoder(KindTurnSubagentActivity, decodeLifecycleData)
	RegisterDataDecoder(KindTurnSessionInit, nil) // no data payload
	RegisterDataDecoder(KindTurnEnded, nil)       // no data payload
	RegisterDataDecoder(KindTurnIdle, nil)        // no data payload
	RegisterDataDecoder(KindTurnCompacting, nil)  // no data payload
	RegisterDataDecoder(KindTurnCompacted, decodeLifecycleData)
	RegisterDataDecoder(KindPromptInput, decodePromptInputData)

	// Recovery
	RegisterDataDecoder(KindRecoveryPointPrepared, decodeRecoveryPointData)
	RegisterDataDecoder(KindRecoveryRestoreStarted, decodeLifecycleData)
	RegisterDataDecoder(KindRecoveryRestoreCompleted, decodeLifecycleData)
	RegisterDataDecoder(KindRecoveryGCCompleted, decodeLifecycleData)
	RegisterDataDecoder(KindContextWindowExceeded, decodeRuntimeDiagnosticData)

	// Usage & channel
	RegisterDataDecoder(KindUsageTokens, decodeUsageTokensData)
	RegisterDataDecoder(KindServerMessageReceived, decodeLifecycleData)
	RegisterDataDecoder(KindServerMessageSent, decodeLifecycleData)

	// Heartbeat
	RegisterDataDecoder(KindRuntimeHeartbeat, decodeHeartbeatData)
	RegisterDataDecoder(KindRuntimeHeartbeatStale, decodeHeartbeatData)

	// Bus diagnostics
	RegisterDataDecoder(KindBusSubscriberDropped, decodeRuntimeDiagnosticData)
	RegisterDataDecoder(KindBusSinkFailed, decodeRuntimeDiagnosticData)
}

// RegisterDataDecoder registers a DataDecoder for a given Kind. Pass nil to
// indicate that the kind carries no data payload (only envelope fields).
func RegisterDataDecoder(kind Kind, dec DataDecoder) {
	decoderMu.Lock()
	defer decoderMu.Unlock()
	decoderMap[kind] = dec
}

// DecodeData looks up the registered DataDecoder for the given Kind and
// deserializes raw. Returns nil if the kind has no data payload. Returns a
// MapData wrapping the raw JSON as map[string]any if no decoder is registered.
func DecodeData(kind Kind, raw json.RawMessage) (EventData, error) {
	if len(raw) == 0 {
		return nil, nil
	}

	decoderMu.RLock()
	dec, ok := decoderMap[kind]
	decoderMu.RUnlock()

	if !ok {
		// Unknown kind — fall back to MapData. Try MapData wire format
		// first ({"map":{...}}), then raw map.
		var md MapData
		if err := json.Unmarshal(raw, &md); err == nil && len(md.Map) > 0 {
			md.kind = kind
			return &md, nil
		}
		var m map[string]any
		if err := json.Unmarshal(raw, &m); err != nil {
			return nil, fmt.Errorf("runtimeevents: decode unknown kind %q: %w", kind, err)
		}
		return NewMapData(kind, m), nil
	}

	if dec == nil {
		// Kind registered as "no data payload".
		return nil, nil
	}

	return dec(kind, unwrapMapDataRaw(raw))
}

// --- decoder implementations -----------------------------------------------

func unwrapMapDataRaw(raw json.RawMessage) json.RawMessage {
	var wrapped struct {
		Map json.RawMessage `json:"map"`
	}
	if err := json.Unmarshal(raw, &wrapped); err == nil && len(wrapped.Map) > 0 {
		return wrapped.Map
	}
	return raw
}

func decodeLifecycleData(kind Kind, raw json.RawMessage) (EventData, error) {
	var d struct {
		Activity   string         `json:"activity"`
		Detail     string         `json:"detail"`
		Text       string         `json:"text"`
		PromptKind string         `json:"prompt_kind"`
		EventType  string         `json:"event_type"`
		Metadata   map[string]any `json:"metadata"`
	}
	if err := json.Unmarshal(raw, &d); err != nil {
		return nil, err
	}
	return NewLifecycleInputData(kind, d.Activity, d.Detail, d.Text, d.PromptKind, d.EventType, d.Metadata), nil
}

func decodeWakeData(kind Kind, raw json.RawMessage) (EventData, error) {
	var d struct {
		Detail string `json:"detail"`
	}
	if err := json.Unmarshal(raw, &d); err != nil {
		return nil, err
	}
	return NewWakeData(kind, d.Detail), nil
}

func decodeRuntimeDiagnosticData(kind Kind, raw json.RawMessage) (EventData, error) {
	var d struct {
		EventClass string `json:"event_class"`
		Detail     string `json:"detail"`
	}
	if err := json.Unmarshal(raw, &d); err != nil {
		return nil, err
	}
	return NewRuntimeDiagnosticData(kind, d.EventClass, d.Detail), nil
}

func decodeTurnTextData(kind Kind, raw json.RawMessage) (EventData, error) {
	var d struct {
		Text      string         `json:"text"`
		Title     string         `json:"title"`
		SectionID string         `json:"section_id"`
		Metadata  map[string]any `json:"metadata"`
	}
	if err := json.Unmarshal(raw, &d); err != nil {
		return nil, err
	}
	out := NewTurnTextDataWithMeta(kind, d.Text, d.Title, d.SectionID)
	out.Metadata = cloneMetadata(d.Metadata)
	return out, nil
}

func decodeToolCallData(kind Kind, raw json.RawMessage) (EventData, error) {
	var d struct {
		ToolName   string         `json:"tool_name"`
		ToolInput  string         `json:"tool_input"`
		ToolCallID string         `json:"tool_call_id"`
		Metadata   map[string]any `json:"metadata"`
	}
	if err := json.Unmarshal(raw, &d); err != nil {
		return nil, err
	}
	return NewToolEventDataWithCallIDAndMetadata(kind, d.ToolName, d.ToolInput, d.ToolCallID, d.Metadata), nil
}

func decodePromptInputData(kind Kind, raw json.RawMessage) (EventData, error) {
	var d struct {
		Text       string `json:"text"`
		PromptKind string `json:"prompt_kind"`
		EventType  string `json:"event_type"`
	}
	if err := json.Unmarshal(raw, &d); err != nil {
		return nil, err
	}
	return NewPromptInputData(d.Text, d.PromptKind, d.EventType), nil
}

func decodeRecoveryPointData(kind Kind, raw json.RawMessage) (EventData, error) {
	var d struct {
		ThreadID    string `json:"thread_id"`
		Strategy    string `json:"strategy"`
		RecoveryRef string `json:"recovery_ref"`
	}
	if err := json.Unmarshal(raw, &d); err != nil {
		return nil, err
	}
	return NewRecoveryPointData(d.ThreadID, d.Strategy, d.RecoveryRef), nil
}

func decodeUsageTokensData(kind Kind, raw json.RawMessage) (EventData, error) {
	var d struct {
		APIKey              string `json:"api_key"`
		InputTokens         uint64 `json:"input_tokens"`
		Input               uint64 `json:"input"`
		CachedInputTokens   uint64 `json:"cached_input_tokens"`
		Cached              uint64 `json:"cached"`
		OutputTokens        uint64 `json:"output_tokens"`
		Output              uint64 `json:"output"`
		ActiveContextTokens uint64 `json:"active_context_tokens"`
		ContextWindowTokens uint64 `json:"context_window_tokens"`
		DurationNS          *int64 `json:"duration"`
		DurationMS          *int64 `json:"duration_ms"`
	}
	if err := json.Unmarshal(raw, &d); err != nil {
		return nil, err
	}
	inputTokens := d.InputTokens
	if inputTokens == 0 {
		inputTokens = d.Input
	}
	cachedInputTokens := d.CachedInputTokens
	if cachedInputTokens == 0 {
		cachedInputTokens = d.Cached
	}
	outputTokens := d.OutputTokens
	if outputTokens == 0 {
		outputTokens = d.Output
	}
	var duration time.Duration
	switch {
	case d.DurationNS != nil:
		duration = time.Duration(*d.DurationNS)
	case d.DurationMS != nil:
		duration = time.Duration(*d.DurationMS) * time.Millisecond
	}
	data := NewUsageTokensData(
		d.APIKey,
		inputTokens,
		cachedInputTokens,
		outputTokens,
		duration,
	)
	data.ActiveContextTokens = d.ActiveContextTokens
	data.ContextWindowTokens = d.ContextWindowTokens
	return data, nil
}

func decodeHeartbeatData(kind Kind, raw json.RawMessage) (EventData, error) {
	var d HeartbeatData
	if err := json.Unmarshal(raw, &d); err != nil {
		return nil, err
	}
	return &d, nil
}

// --- JSON marshal/unmarshal for Event --------------------------------------

type eventJSON struct {
	ID       string          `json:"id"`
	Time     time.Time       `json:"time"`
	Kind     Kind            `json:"kind"`
	Severity Severity        `json:"severity,omitempty"`
	Context  EventContext    `json:"context,omitempty"`
	Trace    EventTrace      `json:"trace,omitempty"`
	Data     json.RawMessage `json:"data,omitempty"`
	Payload  json.RawMessage `json:"payload,omitempty"`
}

// MarshalJSON implements custom JSON marshaling for Event. It serializes
// EventData into the "data" field using the concrete type's JSON tags.
func (ev Event) MarshalJSON() ([]byte, error) {
	ej := eventJSON{
		ID:       ev.ID,
		Time:     ev.Time,
		Kind:     ev.Kind,
		Severity: ev.Severity,
		Context:  ev.Context,
		Trace:    ev.Trace,
	}

	if ev.Data != nil {
		if raw, ok := ev.Data.(*RawData); ok {
			ej.Data = raw.Raw
		} else {
			dataBytes, err := json.Marshal(ev.Data)
			if err != nil {
				return nil, fmt.Errorf("runtimeevents: marshal data for kind %q: %w", ev.Kind, err)
			}
			ej.Data = dataBytes
		}
	}

	return json.Marshal(ej)
}

// UnmarshalJSON implements custom JSON unmarshaling for Event. It uses the
// DataDecoder registry to reconstruct the correct concrete EventData type.
func (ev *Event) UnmarshalJSON(b []byte) error {
	var ej eventJSON
	if err := json.Unmarshal(b, &ej); err != nil {
		return err
	}

	ev.ID = ej.ID
	ev.Time = ej.Time
	ev.Kind = ej.Kind
	ev.Severity = ej.Severity
	ev.Context = ej.Context
	ev.Trace = ej.Trace

	rawData := ej.Data
	if len(rawData) == 0 {
		rawData = ej.Payload
	}
	if len(rawData) > 0 {
		data, err := DecodeData(ej.Kind, rawData)
		if err != nil {
			// Fall back to RawData so we don't lose the event entirely.
			ev.Data = &RawData{Kind: ej.Kind, Raw: rawData}
			return nil
		}
		ev.Data = data
	}

	return nil
}