package dispatcher

import (
	"encoding/json"
	"log/slog"
	"slices"
	"strings"
	"time"
)

type deferredDeliveryEntry struct {
	MessageID      string          `json:"messageId"`
	Seq            int             `json:"seq"`
	AcceptedAt     time.Time       `json:"acceptedAt"`
	Message        json.RawMessage `json:"message"`
	ReplayPending  bool            `json:"replayPending,omitempty"`
	ReplayAttempts int             `json:"replayAttempts,omitempty"`
}

func deferredMessageID(raw json.RawMessage) string {
	var payload struct {
		MessageID string `json:"message_id"`
	}
	if err := json.Unmarshal(raw, &payload); err != nil {
		return ""
	}
	return strings.TrimSpace(payload.MessageID)
}

func compareDeferredDeliveryEntries(a, b deferredDeliveryEntry) int {
	aSeq := a.Seq
	bSeq := b.Seq
	switch {
	case aSeq > 0 && bSeq > 0 && aSeq != bSeq:
		if aSeq < bSeq {
			return -1
		}
		return 1
	case !a.AcceptedAt.Equal(b.AcceptedAt):
		if a.AcceptedAt.Before(b.AcceptedAt) {
			return -1
		}
		return 1
	}

	aMessageID := strings.TrimSpace(a.MessageID)
	bMessageID := strings.TrimSpace(b.MessageID)
	switch {
	case aMessageID < bMessageID:
		return -1
	case aMessageID > bMessageID:
		return 1
	}

	aMessage := string(a.Message)
	bMessage := string(b.Message)
	switch {
	case aMessage < bMessage:
		return -1
	case aMessage > bMessage:
		return 1
	default:
		return 0
	}
}

func normalizeDeferredDeliveryOrder(entries []deferredDeliveryEntry) bool {
	if len(entries) < 2 || slices.IsSortedFunc(entries, compareDeferredDeliveryEntries) {
		return false
	}
	slices.SortStableFunc(entries, compareDeferredDeliveryEntries)
	return true
}

func loadDeferredDeliveries(workDir string) ([]deferredDeliveryEntry, error) {
	state, err := loadSessionResumeState(workDir)
	if err != nil {
		return nil, err
	}
	return state.DeferredDeliveries, nil
}

func appendDeferredDelivery(workDir string, seq int, raw json.RawMessage) (deferredDeliveryEntry, error) {
	var entry deferredDeliveryEntry
	err := updateSessionResumeState(workDir, func(state *sessionResumeState) error {
		entries := append([]deferredDeliveryEntry(nil), state.DeferredDeliveries...)

		entry = deferredDeliveryEntry{
			MessageID:  deferredMessageID(raw),
			Seq:        seq,
			AcceptedAt: time.Now().UTC(),
			Message:    slices.Clone(raw),
		}

		for _, existing := range entries {
			if strings.TrimSpace(existing.MessageID) != "" && existing.MessageID == entry.MessageID {
				entry = existing
				return nil
			}
		}

		entries = append(entries, entry)
		normalizeDeferredDeliveryOrder(entries)
		state.DeferredDeliveries = entries
		return nil
	})
	if err != nil {
		return deferredDeliveryEntry{}, err
	}
	return entry, nil
}

// ReconcileDeferredDeliveries removes confirmed deferred deliveries for an agent.
func (d *Dispatcher) ReconcileDeferredDeliveries(serverAgentID string, messageIDs []string) ([]string, error) {
	targetSession, ok := d.sessionManager.FirstByServerAgentID(serverAgentID)
	if !ok {
		return nil, nil
	}
	return d.reconcileDeferredDeliveriesForSession(targetSession, messageIDs)
}

// ReconcileDeferredDeliveriesForAgent removes confirmed deferred deliveries
// using the authenticated local agent/session lookup identity.
func (d *Dispatcher) ReconcileDeferredDeliveriesForAgent(agentID string, messageIDs []string) ([]string, error) {
	targetSession, ok := d.sessionManager.FindByLookupID(agentID)
	if !ok {
		return nil, nil
	}
	return d.reconcileDeferredDeliveriesForSession(targetSession, messageIDs)
}

func (d *Dispatcher) reconcileDeferredDeliveriesForSession(targetSession *interactiveSession, messageIDs []string) ([]string, error) {
	targetSession.mu.Lock()
	workDir := targetSession.workDir
	targetSession.mu.Unlock()

	if len(messageIDs) == 0 {
		return nil, nil
	}

	seen := make(map[string]struct{}, len(messageIDs))
	for _, messageID := range messageIDs {
		trimmed := strings.TrimSpace(messageID)
		if trimmed == "" {
			continue
		}
		seen[trimmed] = struct{}{}
	}
	if len(seen) == 0 {
		return nil, nil
	}

	suppressed := make([]string, 0, len(seen))
	err := updateSessionResumeState(workDir, func(state *sessionResumeState) error {
		if len(state.DeferredDeliveries) == 0 {
			return nil
		}

		removedAny := false
		nextEntries := make([]deferredDeliveryEntry, 0, len(state.DeferredDeliveries))
		for _, entry := range state.DeferredDeliveries {
			if _, ok := seen[strings.TrimSpace(entry.MessageID)]; ok {
				removedAny = true
				if entry.ReplayPending {
					suppressed = append(suppressed, entry.MessageID)
				}
				continue
			}
			nextEntries = append(nextEntries, entry)
		}
		if !removedAny {
			return nil
		}
		state.DeferredDeliveries = nextEntries
		return nil
	})
	if err != nil {
		return nil, err
	}
	return suppressed, nil
}

// AppendDeferredDelivery persists a deliver message for later replay when the
// agent's runtime does not support stdin delivery.
func (d *Dispatcher) AppendDeferredDelivery(serverAgentID string, seq int, raw json.RawMessage) error {
	targetSession, ok := d.sessionManager.FirstByServerAgentID(serverAgentID)
	if !ok {
		return nil
	}

	targetSession.mu.Lock()
	workDir := targetSession.workDir
	targetSession.mu.Unlock()

	_, err := appendDeferredDelivery(workDir, seq, raw)
	if err != nil {
		slog.Warn("append deferred delivery", "serverAgentID", serverAgentID, "error", err)
	}
	return err
}

// HasDeferredDeliveries reports whether an agent still has queued deferred deliveries.
func (d *Dispatcher) HasDeferredDeliveries(serverAgentID string) bool {
	targetSession, ok := d.sessionManager.FirstByServerAgentID(serverAgentID)
	if !ok {
		return false
	}

	targetSession.mu.Lock()
	workDir := targetSession.workDir
	targetSession.mu.Unlock()

	entries, err := loadDeferredDeliveries(workDir)
	if err != nil {
		slog.Debug("load deferred deliveries", "error", err)
		return false
	}
	return len(entries) > 0
}