package agentstate

import "sync"

type ApplyResult string

const (
	ApplyAccepted             ApplyResult = "accepted"
	ApplyIgnored              ApplyResult = "ignored"
	ApplyNeedsResync          ApplyResult = "needs_resync"
	ApplyIgnoredUntilSnapshot ApplyResult = "ignored_until_snapshot"
)

type AcceptedState struct {
	AgentID          string
	SessionEpoch     uint64
	Revision         uint64
	AwaitingSnapshot bool
	Execution        AgentExecutionState
	Deliveries       map[string]WakeDeliveryState
	Diagnostics      map[string]AgentDiagnostic
}

type Consumer struct {
	mu sync.Mutex

	state    AcceptedState
	accepted bool
}

func NewConsumer() *Consumer {
	return &Consumer{}
}

func (consumer *Consumer) Apply(envelope AgentStateEnvelopeV2) ApplyResult {
	if err := envelope.Validate(); err != nil {
		return ApplyIgnored
	}
	consumer.mu.Lock()
	defer consumer.mu.Unlock()

	if envelope.Snapshot != nil {
		return consumer.applySnapshotLocked(*envelope.Snapshot)
	}
	return consumer.applyDeltaLocked(*envelope.Delta)
}

func (consumer *Consumer) applySnapshotLocked(snapshot AgentStateSnapshotV2) ApplyResult {
	if consumer.accepted {
		switch {
		case snapshot.SessionEpoch < consumer.state.SessionEpoch:
			return ApplyIgnored
		case snapshot.SessionEpoch == consumer.state.SessionEpoch &&
			snapshot.Revision <= consumer.state.Revision:
			return ApplyIgnored
		}
	}
	next := AcceptedState{
		AgentID:          snapshot.AgentID,
		SessionEpoch:     snapshot.SessionEpoch,
		Revision:         snapshot.Revision,
		AwaitingSnapshot: false,
		Execution:        snapshot.Execution,
		Deliveries:       make(map[string]WakeDeliveryState, len(snapshot.Deliveries)),
		Diagnostics:      make(map[string]AgentDiagnostic, len(snapshot.Diagnostics)),
	}
	for _, delivery := range snapshot.Deliveries {
		next.Deliveries[delivery.LocalWakeID] = cloneWakeDelivery(delivery)
	}
	for _, diagnostic := range snapshot.Diagnostics {
		next.Diagnostics[diagnostic.DiagnosticID] = cloneDiagnostic(diagnostic)
	}
	consumer.state = next
	consumer.accepted = true
	return ApplyAccepted
}

func (consumer *Consumer) applyDeltaLocked(delta AgentStateDeltaV2) ApplyResult {
	if !consumer.accepted {
		consumer.state.AwaitingSnapshot = true
		return ApplyNeedsResync
	}
	if delta.AgentID != consumer.state.AgentID {
		return ApplyIgnored
	}
	switch {
	case delta.SessionEpoch < consumer.state.SessionEpoch:
		return ApplyIgnored
	case delta.SessionEpoch > consumer.state.SessionEpoch:
		consumer.state.AwaitingSnapshot = true
		return ApplyNeedsResync
	case consumer.state.AwaitingSnapshot:
		return ApplyIgnoredUntilSnapshot
	case delta.Revision <= consumer.state.Revision:
		return ApplyIgnored
	case delta.BaseRevision != consumer.state.Revision:
		consumer.state.AwaitingSnapshot = true
		return ApplyNeedsResync
	}

	next := cloneAcceptedState(consumer.state)
	if delta.Execution != nil {
		next.Execution = *delta.Execution
	}
	for _, delivery := range delta.DeliveryUpserts {
		next.Deliveries[delivery.LocalWakeID] = cloneWakeDelivery(delivery)
	}
	for _, localWakeID := range delta.DeliveryRemovals {
		delete(next.Deliveries, localWakeID)
	}
	for _, diagnostic := range delta.DiagnosticUpserts {
		next.Diagnostics[diagnostic.DiagnosticID] = cloneDiagnostic(diagnostic)
	}
	for _, diagnosticID := range delta.DiagnosticRemovals {
		delete(next.Diagnostics, diagnosticID)
	}
	next.Revision = delta.Revision
	consumer.state = next
	return ApplyAccepted
}

func (consumer *Consumer) State() (AcceptedState, bool) {
	consumer.mu.Lock()
	defer consumer.mu.Unlock()
	if !consumer.accepted {
		return AcceptedState{}, false
	}
	return cloneAcceptedState(consumer.state), true
}

func (consumer *Consumer) AwaitingSnapshot() bool {
	consumer.mu.Lock()
	defer consumer.mu.Unlock()
	return consumer.state.AwaitingSnapshot
}

func cloneAcceptedState(state AcceptedState) AcceptedState {
	state.Deliveries = cloneDeliveryMap(state.Deliveries)
	state.Diagnostics = cloneDiagnosticMap(state.Diagnostics)
	return state
}

func cloneDeliveryMap(deliveries map[string]WakeDeliveryState) map[string]WakeDeliveryState {
	if deliveries == nil {
		return nil
	}
	result := make(map[string]WakeDeliveryState, len(deliveries))
	for id, delivery := range deliveries {
		result[id] = cloneWakeDelivery(delivery)
	}
	return result
}

func cloneDiagnosticMap(diagnostics map[string]AgentDiagnostic) map[string]AgentDiagnostic {
	if diagnostics == nil {
		return nil
	}
	result := make(map[string]AgentDiagnostic, len(diagnostics))
	for id, diagnostic := range diagnostics {
		result[id] = cloneDiagnostic(diagnostic)
	}
	return result
}

func cloneDiagnostic(diagnostic AgentDiagnostic) AgentDiagnostic {
	diagnostic.LocalWakeIDs = append([]string(nil), diagnostic.LocalWakeIDs...)
	return diagnostic
}