package dispatcher

import (
	"context"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeevents"
	"github.com/tidbcloud/AgentTeam/desktop/internal/metrics"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// TurnStateObservation is an immutable same-generation observation of an
// agent's real turn execution state and the RuntimeSessionID that owns it. Both
// values are read together under a single session lock so the pair can never
// straddle a generation replacement: a torn read would publish old idle with a
// new session, which the Hub's new-session fence would wrongly accept as the
// current generation's authoritative idle.
type TurnStateObservation struct {
	TurnState string
	SessionID string
}

// AgentTeamReporterSink adapts normalized runtime events to the existing AgentTeam channel
// EventReporter contract. It keeps the daemon-internal EventBus independent of
// the websocket-facing types.AgentEvent wire model.
type AgentTeamReporterSink struct {
	reporterForServer func(serverID string) types.EventReporter
	// turnStateResolver returns one immutable TurnStateObservation: the
	// daemon-derived real Agent turn execution state (active/queued/idle) and the
	// RuntimeSessionID that owns it, read from a single locked SessionSnapshot. A
	// zero TurnState means unknown (no session) and suppresses emission. The pair
	// is emitted on every TurnStateEvent so the server Hub session fence can drop
	// a stale turn-state from a replaced session (recovery observation fencing),
	// never old idle + new session.
	turnStateResolver func(serverID, agentID string) TurnStateObservation
	mu                sync.Mutex
	lastActivity      map[string]reportedActivity
}

type reportedActivity struct {
	activity string
	detail   string
}

// NewAgentTeamReporterSink creates the AgentTeam server connection event sink. The
// optional turnStateResolver publishes the authoritative active/queued/idle
// signal and the owning RuntimeSessionID as one immutable observation; when nil,
// no turn-state events are emitted (old daemon / disabled).
func NewAgentTeamReporterSink(reporterForServer func(serverID string) types.EventReporter, turnStateResolver func(serverID, agentID string) TurnStateObservation) *AgentTeamReporterSink {
	return &AgentTeamReporterSink{
		reporterForServer: reporterForServer,
		turnStateResolver: turnStateResolver,
		lastActivity:      make(map[string]reportedActivity),
	}
}

// Subscribe attaches the sink to bus and starts its forwarding agentteam.
func (s *AgentTeamReporterSink) Subscribe(ctx context.Context, bus runtimeevents.Bus) runtimeevents.Subscription {
	sub := bus.Subscribe("agentteam_reporter", nil, 1024, runtimeevents.DropOldest)
	go runRuntimeEventSink(ctx, sub, s.Handle)
	return sub
}

// Handle maps one normalized runtime event into zero or more AgentTeam server events.
func (s *AgentTeamReporterSink) Handle(ev runtimeevents.Event) {
	if s == nil || s.reporterForServer == nil {
		return
	}
	reporter := s.reporterForServer(ev.Context.ServerID)
	if reporter == nil {
		return
	}
	isTerminal := ev.Kind == runtimeevents.KindSessionStopped
	// A session-terminal event maps below to session-clearing agent:session,
	// inactive agent:status, and offline agent:activity frames. Resolve the
	// stopped event's TurnStateObservation EXACTLY ONCE and reuse that same
	// immutable (state, session) pair as the single dying-generation identity for
	// EVERY terminal projection (idle turn-state, inactive/offline telemetry, and
	// the empty session clear), so a generation replacement between them can never
	// tear (old idle + new session, or a clear/telemetry fenced by the wrong
	// generation).
	terminalObs := TurnStateObservation{}
	if isTerminal {
		terminalObs = s.observationFor(ev)
		// Ordering is essential for a matching current-generation stop: publish
		// the terminal idle turn-state FIRST, then the fenced inactive/offline
		// telemetry, and the fenced empty session clear LAST. If the clear were
		// first, the subsequently arriving matching telemetry would see no current
		// session and be incorrectly dropped by the session fence. For a stale
		// (replaced-generation) stop, every one of these frames carries the dying
		// session identity and is rejected against the new session.
		s.emitTurnStateObservation(ev, reporter, terminalObs)
		s.emitTerminalLifecycleFrames(ev, reporter, terminalObs)
		return
	}
	for _, agentEvent := range mapRuntimeEventToAgentEvents(ev) {
		if activity, ok := agentEvent.(types.ActivityEvent); ok && s.skipDuplicateActivity(ev.Kind, ev.Context.ServerID, activity) {
			continue
		}
		reporter.Emit(agentEvent)
	}
	s.emitTurnStateIfRelevant(ev, reporter)
}

// emitTerminalLifecycleFrames publishes the terminal inactive/offline telemetry
// and the empty session clear for a session-terminal event. All frames carry the
// SAME dying-generation ExpectedSessionID so the server applies them only while
// its current agent session still matches; the session clear is emitted LAST so
// a matching stop still tears down the session after the telemetry lands, while
// a stale stop is rejected wholesale.
func (s *AgentTeamReporterSink) emitTerminalLifecycleFrames(ev runtimeevents.Event, reporter types.EventReporter, obs TurnStateObservation) {
	if s == nil || reporter == nil {
		return
	}
	var clearEvent *types.SessionEvent
	for _, agentEvent := range mapRuntimeEventToAgentEvents(ev) {
		switch typed := agentEvent.(type) {
		case types.SessionEvent:
			if strings.TrimSpace(typed.SessionID) != "" {
				reporter.Emit(agentEvent)
				continue
			}
			// Hold the empty session clear for LAST.
			copy := typed
			copy.ExpectedSessionID = obs.SessionID
			clearEvent = &copy
		case types.StatusEvent:
			copy := typed
			copy.ExpectedSessionID = obs.SessionID
			reporter.Emit(copy)
		case types.ActivityEvent:
			copy := typed
			copy.ExpectedSessionID = obs.SessionID
			reporter.Emit(copy)
		default:
			reporter.Emit(agentEvent)
		}
	}
	if clearEvent != nil {
		reporter.Emit(*clearEvent)
	}
}

// observationFor resolves the immutable (turn state, owning RuntimeSessionID)
// observation for one event through the combined resolver, or a zero observation
// when the event cannot change turn state, the agent is unknown, or the resolver
// is absent. One combined resolver call reads both values from the SAME locked
// SessionSnapshot, so a generation replacement committed between two independent
// reads can never produce a torn old-idle + new-session pair.
func (s *AgentTeamReporterSink) observationFor(ev runtimeevents.Event) TurnStateObservation {
	if s == nil || s.turnStateResolver == nil {
		return TurnStateObservation{}
	}
	agentID := runtimeEventAgentID(ev)
	if agentID == "" {
		return TurnStateObservation{}
	}
	return s.turnStateResolver(ev.Context.ServerID, agentID)
}

// emitTurnStateIfRelevant publishes the daemon-derived real turn execution state
// at the events that can change active/queued/idle: session start/stop, delivery
// pending/acked, and turn end/idle/start. The state is resolved live from the
// Process so it reflects the current real queue, not a presence guess.
func (s *AgentTeamReporterSink) emitTurnStateIfRelevant(ev runtimeevents.Event, reporter types.EventReporter) {
	if s == nil || s.turnStateResolver == nil || reporter == nil {
		return
	}
	if !s.turnStateEventRelevant(ev.Kind) {
		return
	}
	s.emitTurnStateObservation(ev, reporter, s.observationFor(ev))
}

// emitTurnStateObservation emits the given immutable observation as a
// TurnStateEvent when it carries a real state. A zero TurnState (unknown/no
// session) suppresses emission.
func (s *AgentTeamReporterSink) emitTurnStateObservation(ev runtimeevents.Event, reporter types.EventReporter, obs TurnStateObservation) {
	if s == nil || reporter == nil {
		return
	}
	agentID := runtimeEventAgentID(ev)
	if agentID == "" {
		return
	}
	if strings.TrimSpace(obs.TurnState) == "" {
		return
	}
	reporter.Emit(types.TurnStateEvent{AgentID: agentID, TurnState: obs.TurnState, SessionID: obs.SessionID})
}

func (s *AgentTeamReporterSink) turnStateEventRelevant(kind runtimeevents.Kind) bool {
	switch kind {
	case runtimeevents.KindSessionStarting,
		runtimeevents.KindSessionStarted,
		runtimeevents.KindSessionSleeping,
		runtimeevents.KindSessionStopped,
		runtimeevents.KindDeliveryPending,
		runtimeevents.KindDeliveryAcked,
		runtimeevents.KindTurnEnded,
		runtimeevents.KindTurnIdle,
		runtimeevents.KindTurnThinking,
		runtimeevents.KindTurnText,
		runtimeevents.KindTurnToolCall,
		runtimeevents.KindTurnCompacting:
		return true
	default:
		return false
	}
}

func (s *AgentTeamReporterSink) skipDuplicateActivity(kind runtimeevents.Kind, serverID string, ev types.ActivityEvent) bool {
	if s == nil {
		return false
	}
	if kind != runtimeevents.KindTurnThinking && kind != runtimeevents.KindTurnText {
		return false
	}
	agentID := strings.TrimSpace(ev.AgentID)
	if agentID == "" {
		return false
	}
	key := strings.TrimSpace(serverID) + "\x00" + agentID
	next := reportedActivity{activity: strings.TrimSpace(ev.Activity), detail: strings.TrimSpace(ev.Detail)}

	s.mu.Lock()
	defer s.mu.Unlock()
	if s.lastActivity == nil {
		s.lastActivity = make(map[string]reportedActivity)
	}
	if s.lastActivity[key] == next {
		return true
	}
	s.lastActivity[key] = next
	return false
}

// TrajectorySink adapts normalized turn/activity events to the dashboard
// trajectory hub used by SSE subscribers.
type TrajectorySink struct {
	hub *trajectoryHub
}

// NewTrajectorySink creates the dashboard trajectory adapter.
func NewTrajectorySink(hub *trajectoryHub) *TrajectorySink {
	return &TrajectorySink{hub: hub}
}

// Subscribe attaches the sink to turn events on bus.
func (s *TrajectorySink) Subscribe(ctx context.Context, bus runtimeevents.Bus) runtimeevents.Subscription {
	sub := bus.Subscribe("trajectory", runtimeevents.FilterKindPrefix("turn."), 1024, runtimeevents.DropOldest)
	go runRuntimeEventSink(ctx, sub, s.Handle)
	return sub
}

// Handle appends trajectory rows for one normalized runtime event.
func (s *TrajectorySink) Handle(ev runtimeevents.Event) {
	if s == nil || s.hub == nil {
		return
	}
	entries := runtimeEventTrajectoryEntries(ev)
	if len(entries) == 0 {
		return
	}
	agentID := runtimeEventAgentID(ev)
	if agentID == "" {
		return
	}
	s.hub.BroadcastAt(agentID, entries, runtimeEventUnixMillis(ev))
}

// MetricsSink adapts normalized usage/channel events to existing metrics
// collectors. The record functions are injectable so mapper tests do not rely
// on global singletons.
type MetricsSink struct {
	recordRequest  func(apiKey, model string, inputTokens, cachedInputTokens, outputTokens uint64, duration time.Duration)
	recordReceived func(serverID, agentID string)
	recordSent     func(serverID, agentID string)
}

// NewMetricsSink creates a metrics adapter backed by the global metrics
// singletons.
func NewMetricsSink() *MetricsSink {
	return &MetricsSink{
		recordRequest: func(apiKey, model string, inputTokens, cachedInputTokens, outputTokens uint64, duration time.Duration) {
			collector := metrics.DefaultRequestRecorder()
			if collector == nil {
				return
			}
			collector.RecordRequest(apiKey, model, inputTokens, cachedInputTokens, outputTokens, duration)
		},
		recordReceived: func(serverID, agentID string) {
			collector := metrics.DefaultServerMessageCollector()
			if collector == nil {
				return
			}
			collector.RecordReceivedForAgent(serverID, agentID)
		},
		recordSent: func(serverID, agentID string) {
			collector := metrics.DefaultServerMessageCollector()
			if collector == nil {
				return
			}
			collector.RecordSentForAgent(serverID, agentID)
		},
	}
}

// Subscribe attaches the sink to usage and server message events on bus.
func (s *MetricsSink) Subscribe(ctx context.Context, bus runtimeevents.Bus) runtimeevents.Subscription {
	filter := runtimeevents.FilterOr(
		runtimeevents.FilterKind(runtimeevents.KindUsageTokens),
		runtimeevents.FilterKind(runtimeevents.KindServerMessageReceived),
		runtimeevents.FilterKind(runtimeevents.KindServerMessageSent),
	)
	sub := bus.Subscribe("metrics", filter, 1024, runtimeevents.DropOldest)
	go runRuntimeEventSink(ctx, sub, s.Handle)
	return sub
}

// Handle records metrics represented by one normalized runtime event.
func (s *MetricsSink) Handle(ev runtimeevents.Event) {
	if s == nil {
		return
	}
	switch ev.Kind {
	case runtimeevents.KindUsageTokens:
		if s.recordRequest == nil {
			return
		}
		if data, ok := runtimeevents.DataAs[*runtimeevents.UsageTokensData](ev); ok {
			model := firstNonEmpty(ev.Context.Model, payloadString(ev, "model"))
			if data.APIKey == "" || model == "" {
				return
			}
			s.recordRequest(data.APIKey, model, data.InputTokens, data.CachedInputTokens, data.OutputTokens, data.Duration)
			return
		}
		apiKey := payloadString(ev, "api_key")
		model := firstNonEmpty(ev.Context.Model, payloadString(ev, "model"))
		if apiKey == "" || model == "" {
			return
		}
		s.recordRequest(
			apiKey,
			model,
			payloadUint64(ev, "input_tokens", "input"),
			payloadUint64(ev, "cached_input_tokens", "cached"),
			payloadUint64(ev, "output_tokens", "output"),
			payloadDuration(ev, "duration", "duration_ms"),
		)
	case runtimeevents.KindServerMessageReceived:
		if s.recordReceived != nil {
			s.recordReceived(ev.Context.ServerID, runtimeEventAgentID(ev))
		}
	case runtimeevents.KindServerMessageSent:
		if s.recordSent != nil {
			s.recordSent(ev.Context.ServerID, runtimeEventAgentID(ev))
		}
	}
}

func runRuntimeEventSink(ctx context.Context, sub runtimeevents.Subscription, handle func(runtimeevents.Event)) {
	if sub == nil || handle == nil {
		return
	}
	for {
		select {
		case <-ctx.Done():
			sub.Unsubscribe()
			return
		case ev, ok := <-sub.Events():
			if !ok {
				return
			}
			handle(ev)
		}
	}
}