package agent

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// ---------------------------------------------------------------------------
// Wake status helpers
// ---------------------------------------------------------------------------

// emitWakeStatus sends a wake status event via the onWakeStatus callback.
func (ap *Process) emitWakeStatus(ev WakeStatusEvent) {
	if strings.TrimSpace(ev.Status) == "" {
		return
	}
	if fn := ap.onWakeStatus; fn != nil {
		fn(ev)
	}
}

type wakeSnapshot struct {
	pendingWake          *pendingDelivery
	pendingWakeToken     string
	pendingDeliveryQueue []pendingDelivery
}

type deliveryAttemptSnapshot struct {
	wakeSnapshot
	isIdle              bool
	runtimeHandoffEpoch uint64
}

func deliveryStatusForWake(wake WakeInfo) pendingDeliveryStatus {
	if wake.DeliveryClass == DeliveryClassSignalOnly {
		return pendingDeliveryChecked
	}
	return pendingDeliveryAwaitingCheck
}

func (ap *Process) beginInitialWake(wake *WakeInfo) {
	if wake == nil || strings.TrimSpace(wake.MessageID) == "" {
		return
	}
	if wake.DeliveryClass == "" {
		wake.DeliveryClass = DeliveryClassVisibleReply
	}
	delivery := pendingDelivery{
		Wake:   *wake,
		Token:  generateWakeToken(),
		Status: deliveryStatusForWake(*wake),
	}
	ap.pendingWake = &delivery
	ap.pendingWakeToken = delivery.Token
	ap.currentTurnMessageID = strings.TrimSpace(wake.MessageID)
	ap.emitWakeStatus(WakeStatusEvent{
		MessageID: delivery.Wake.MessageID,
		WakeToken: delivery.Token,
		Target:    delivery.Wake.Target,
		Status:    "received",
		Detail:    "agent starting",
	})
}

func (ap *Process) deliveryAttemptSnapshotLocked() deliveryAttemptSnapshot {
	return deliveryAttemptSnapshot{
		wakeSnapshot:        ap.wakeSnapshotLocked(),
		isIdle:              ap.isIdle,
		runtimeHandoffEpoch: ap.runtimeHandoffEpoch,
	}
}

func clonePendingDeliveries(deliveries []pendingDelivery) []pendingDelivery {
	if len(deliveries) == 0 {
		return nil
	}
	return append([]pendingDelivery(nil), deliveries...)
}

func clonePendingDelivery(delivery *pendingDelivery) *pendingDelivery {
	if delivery == nil {
		return nil
	}
	copy := *delivery
	return &copy
}

func (ap *Process) wakeSnapshotLocked() wakeSnapshot {
	return wakeSnapshot{
		pendingWake:          clonePendingDelivery(ap.pendingWake),
		pendingWakeToken:     ap.pendingWakeToken,
		pendingDeliveryQueue: clonePendingDeliveries(ap.pendingDeliveryQueue),
	}
}

func (ap *Process) restoreWakeSnapshot(snapshot wakeSnapshot) {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	ap.restoreWakeSnapshotLocked(snapshot)
}

func (ap *Process) restoreWakeSnapshotLocked(snapshot wakeSnapshot) {
	ap.pendingWake = clonePendingDelivery(snapshot.pendingWake)
	ap.pendingWakeToken = snapshot.pendingWakeToken
	ap.pendingDeliveryQueue = clonePendingDeliveries(snapshot.pendingDeliveryQueue)
}

// restoreDeliveryAttempt atomically rolls wake and runtime phase back after a
// direct idle handoff failed. A newer successful handoff epoch wins and is
// never overwritten by an older RPC completion.
func (ap *Process) restoreDeliveryAttempt(snapshot deliveryAttemptSnapshot) bool {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	if ap.runtimeHandoffEpoch != snapshot.runtimeHandoffEpoch {
		return false
	}
	ap.restoreWakeSnapshotLocked(snapshot.wakeSnapshot)
	ap.isIdle = snapshot.isIdle
	return true
}

func (ap *Process) beginWakeDeliveryForMessageLocked(msg types.InboxMessage, detail string) (WakeStatusEvent, bool) {
	return ap.beginWakeDeliveryWithTokenForMessageLocked(msg, "", detail)
}

func (ap *Process) beginWakeDeliveryWithTokenForMessageLocked(msg types.InboxMessage, wakeToken string, detail string) (WakeStatusEvent, bool) {
	messageID := strings.TrimSpace(msg.MessageID)
	if messageID == "" {
		return WakeStatusEvent{}, false
	}
	wakeToken = strings.TrimSpace(wakeToken)
	if wakeToken == "" {
		wakeToken = generateWakeToken()
	}
	wake := WakeInfo{
		MessageID:     messageID,
		Target:        strings.TrimSpace(msg.Target),
		TargetKey:     strings.TrimSpace(msg.TargetKey),
		TargetSeq:     msg.TargetSeq,
		MessageType:   strings.TrimSpace(msg.MessageType),
		DeliveryClass: DeliveryClassForMessage(msg),
	}
	delivery := pendingDelivery{
		Wake:   wake,
		Token:  wakeToken,
		Status: deliveryStatusForWake(wake),
	}
	if ap.pendingDeliveryIndexLocked("", messageID, "") >= 0 {
		return WakeStatusEvent{}, false
	}
	// Capacity is reserved before runtime admission. Never drop an already
	// accepted/acked body here: an older binary or a concurrent correlated
	// echo may legitimately convert a pending owner after the nominal cap.
	if ap.pendingWake == nil {
		ap.activatePendingDeliveryLocked(delivery)
	} else {
		ap.pendingDeliveryQueue = append(ap.pendingDeliveryQueue, delivery)
	}
	return wakeStatusForDelivery(delivery, "received", "", detail), true
}

func (ap *Process) activePendingDeliveryLocked() (pendingDelivery, bool) {
	if ap.pendingWake == nil {
		return pendingDelivery{}, false
	}
	return *ap.pendingWake, true
}

func (ap *Process) activatePendingDeliveryLocked(delivery pendingDelivery) {
	if delivery.Wake.DeliveryClass == "" {
		delivery.Wake.DeliveryClass = DeliveryClassVisibleReply
	}
	if delivery.Status != pendingDeliveryChecked {
		delivery.Status = deliveryStatusForWake(delivery.Wake)
	}
	ap.pendingWake = clonePendingDelivery(&delivery)
	ap.currentTurnMessageID = strings.TrimSpace(delivery.Wake.MessageID)
	ap.pendingWakeToken = strings.TrimSpace(delivery.Token)
	ap.runtimeErrorDetail = ""
	ap.runtimeErrorEventEmitted = false
	ap.runtimeErrorWakeFailed = false
	ap.transientRuntimeRetry = transientRuntimeRetryState{}
	ap.transientRuntimeRetryPending = false
	ap.transientRuntimeRetryAwaitingTurnEnd = false
	ap.transientRuntimeRetryRedelivered = false
	ap.transientRuntimeRetryToolCallObserved = false
}

func (ap *Process) pendingDeliveryCountLocked() int {
	if ap.pendingWake == nil {
		return 0
	}
	return 1 + len(ap.pendingDeliveryQueue)
}

// discardUnackedPendingDeliveryLocked releases the reply-obligation
// reservation created before a correlated runtime mutation. Authoritative
// Absent proves that mutation was never acknowledged, so the same durable message can
// re-enter admission without being blocked by its stale pre-ack row.
func (ap *Process) discardUnackedPendingDeliveryLocked(wakeToken, messageID string) {
	index := ap.pendingDeliveryIndexLocked(wakeToken, messageID, "")
	if index < 0 {
		return
	}
	ap.removePendingDeliveryAtLocked(index)
}

// pendingWakeLifecycleOwnerCountLocked counts unique reply obligations already
// accepted by the runtime: acked deliveries plus correlated pending inputs
// that have not produced a row yet. Queued daemon owners are deliberately not
// counted; they have not crossed the runtime admission boundary.
func (ap *Process) pendingWakeLifecycleOwnerCountLocked() int {
	count := ap.pendingDeliveryCountLocked()
	for _, pending := range ap.pendingRuntimeInputs {
		if pending.reclaimed {
			continue
		}
		messageID := strings.TrimSpace(pending.Message.MessageID)
		if messageID == "" {
			messageID = strings.TrimSpace(pending.Input.ID)
		}
		if messageID == "" || ap.pendingDeliveryIndexLocked("", messageID, "") >= 0 {
			continue
		}
		count++
	}
	return count
}

func (ap *Process) canAcceptWakeLifecycleOwnerLocked(messageID string) bool {
	messageID = strings.TrimSpace(messageID)
	if messageID != "" && ap.pendingDeliveryIndexLocked("", messageID, "") >= 0 {
		return true
	}
	return ap.pendingWakeLifecycleOwnerCountLocked() < processInboxMaxMessages
}

// pendingDeliveryIndexLocked returns 0 for the active FIFO head and i+1 for a
// queued delivery. Every non-empty criterion must match; selector matches
// either the exact target or triggering message id.
func (ap *Process) pendingDeliveryIndexLocked(token, selector, target string) int {
	token = strings.TrimSpace(token)
	selector = strings.TrimSpace(selector)
	target = strings.TrimSpace(target)
	matches := func(delivery pendingDelivery) bool {
		if token != "" && strings.TrimSpace(delivery.Token) != token {
			return false
		}
		if target != "" && strings.TrimSpace(delivery.Wake.Target) != target {
			return false
		}
		if selector != "" && strings.TrimSpace(delivery.Wake.Target) != selector && strings.TrimSpace(delivery.Wake.MessageID) != selector {
			return false
		}
		return true
	}
	if active, ok := ap.activePendingDeliveryLocked(); ok && matches(active) {
		return 0
	}
	for i, delivery := range ap.pendingDeliveryQueue {
		if matches(delivery) {
			return i + 1
		}
	}
	return -1
}

func (ap *Process) pendingDeliveryAtLocked(index int) (pendingDelivery, bool) {
	if index == 0 {
		return ap.activePendingDeliveryLocked()
	}
	i := index - 1
	if i < 0 || i >= len(ap.pendingDeliveryQueue) {
		return pendingDelivery{}, false
	}
	return ap.pendingDeliveryQueue[i], true
}

func (ap *Process) pendingDeliveriesLocked() []pendingDelivery {
	deliveries := make([]pendingDelivery, 0, ap.pendingDeliveryCountLocked())
	if active, ok := ap.activePendingDeliveryLocked(); ok {
		deliveries = append(deliveries, active)
	}
	return append(deliveries, ap.pendingDeliveryQueue...)
}

func (ap *Process) pendingDeliveryInfosLocked() []WakeInfo {
	deliveries := ap.pendingDeliveriesLocked()
	wakes := make([]WakeInfo, 0, len(deliveries))
	for _, delivery := range deliveries {
		wakes = append(wakes, delivery.Wake)
	}
	return wakes
}

// pendingWakeNeedsFinalCheckLocked reports whether any ordinary visible-reply
// obligation is awaiting its one final delivery check. Signal-only rows are
// already in the checked status and never enter this path.
func (ap *Process) pendingWakeNeedsFinalCheckLocked() bool {
	if ap.pendingWake != nil &&
		ap.pendingWake.Wake.DeliveryClass != DeliveryClassSignalOnly &&
		ap.pendingWake.Status == pendingDeliveryAwaitingCheck {
		return true
	}
	for _, delivery := range ap.pendingDeliveryQueue {
		if delivery.Wake.DeliveryClass != DeliveryClassSignalOnly &&
			delivery.Status == pendingDeliveryAwaitingCheck {
			return true
		}
	}
	return false
}

func (ap *Process) markPendingDeliveriesCheckedLocked() {
	if ap.pendingWake != nil && ap.pendingWake.Wake.DeliveryClass != DeliveryClassSignalOnly {
		ap.pendingWake.Status = pendingDeliveryChecked
	}
	for index := range ap.pendingDeliveryQueue {
		if ap.pendingDeliveryQueue[index].Wake.DeliveryClass != DeliveryClassSignalOnly {
			ap.pendingDeliveryQueue[index].Status = pendingDeliveryChecked
		}
	}
}

// generateWakeToken returns a random 32-character hex string used to
// correlate wake lifecycle events. It is a delivery correlation value, not a
// model-visible routing selector.
func generateWakeToken() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return fmt.Sprintf("%x", time.Now().UnixNano())
	}
	return hex.EncodeToString(b)
}

// PrepareVisibleReply returns a non-persistent delivery correlation token for
// an explicit reaction binding. It does not mutate delivery state or participate
// in ordinary target-frontier settlement; server freshness responses settle those
// obligations instead.
func (ap *Process) PrepareVisibleReply(selectors ...string) string {
	selector := ""
	for _, candidate := range selectors {
		if selector = strings.TrimSpace(candidate); selector != "" {
			break
		}
	}
	ap.mu.Lock()
	defer ap.mu.Unlock()
	index := 0
	if selector != "" {
		index = ap.pendingDeliveryIndexLocked("", selector, "")
	}
	delivery, ok := ap.pendingDeliveryAtLocked(index)
	if !ok || delivery.Wake.DeliveryClass == DeliveryClassSignalOnly {
		return ""
	}
	return delivery.Token
}

// ConfirmVisibleReplyForTarget settles every acked, reply-eligible delivery
// for one resolved target whose target sequence is covered by the server's
// SeenUpToSeq frontier. It deliberately ignores wake tokens: those tokens
// identify individual delivery rows, while the server freshness response is
// the authoritative coverage boundary for the target.
func (ap *Process) ConfirmVisibleReplyForTarget(targetKey string, seenUpToSeq int64, replyMessageID string) bool {
	targetKey = strings.TrimSpace(targetKey)
	if targetKey == "" || seenUpToSeq <= 0 {
		return false
	}

	ap.mu.Lock()
	var events []WakeStatusEvent
	for {
		index := ap.pendingDeliveryIndexByTargetLocked(targetKey, seenUpToSeq)
		if index < 0 {
			break
		}
		event := ap.confirmWakeReplyAtLocked(index, replyMessageID, "confirmed")
		if strings.TrimSpace(event.MessageID) == "" {
			break
		}
		events = append(events, event)
	}
	// A successful send may cover only the older part of a target's acked
	// frontier. Newer same-target rows that had already been checked must be
	// eligible for a fresh check at the next boundary.
	ap.resetPendingDeliveryChecksForTargetLocked(targetKey, seenUpToSeq)
	ap.mu.Unlock()

	for _, event := range events {
		ap.emitWakeStatus(event)
	}
	if len(events) > 0 {
		ap.scheduleInboxRetry(true)
	}
	return len(events) > 0
}

func (ap *Process) pendingDeliveryIndexByTargetLocked(targetKey string, seenUpToSeq int64) int {
	matches := func(delivery pendingDelivery) bool {
		return delivery.Wake.DeliveryClass != DeliveryClassSignalOnly &&
			strings.TrimSpace(delivery.Wake.TargetKey) == targetKey &&
			delivery.Wake.TargetSeq > 0 && delivery.Wake.TargetSeq <= seenUpToSeq
	}
	if active, ok := ap.activePendingDeliveryLocked(); ok && matches(active) {
		return 0
	}
	for index, delivery := range ap.pendingDeliveryQueue {
		if matches(delivery) {
			return index + 1
		}
	}
	return -1
}

func (ap *Process) resetPendingDeliveryChecksForTargetLocked(targetKey string, seenUpToSeq int64) {
	reset := func(delivery *pendingDelivery) {
		if delivery == nil || delivery.Wake.DeliveryClass == DeliveryClassSignalOnly {
			return
		}
		if strings.TrimSpace(delivery.Wake.TargetKey) == targetKey && delivery.Wake.TargetSeq > seenUpToSeq {
			delivery.Status = pendingDeliveryAwaitingCheck
		}
	}
	reset(ap.pendingWake)
	for index := range ap.pendingDeliveryQueue {
		reset(&ap.pendingDeliveryQueue[index])
	}
}

// ConfirmReactionReply marks a wake as answered when the agent reacts to the
// exact message that triggered the current wake.
func (ap *Process) ConfirmReactionReply(wakeToken, messageID string) bool {
	ap.mu.Lock()
	index := ap.pendingDeliveryIndexLocked(wakeToken, messageID, "")
	delivery, ok := ap.pendingDeliveryAtLocked(index)
	if !ok || strings.TrimSpace(delivery.Wake.MessageID) != strings.TrimSpace(messageID) || delivery.Wake.DeliveryClass == DeliveryClassSignalOnly {
		ap.mu.Unlock()
		return false
	}
	event := ap.confirmWakeReplyAtLocked(index, "", "reaction confirmed")
	ap.mu.Unlock()

	ap.emitWakeStatus(event)
	ap.scheduleInboxRetry(true)
	return true
}

func (ap *Process) confirmWakeReplyAtLocked(index int, replyMessageID string, detail string) WakeStatusEvent {
	delivery, ok := ap.pendingDeliveryAtLocked(index)
	if !ok {
		return WakeStatusEvent{}
	}
	return ap.removePendingDeliveryWithStatusLocked(index, consts.WakeStatusReplySent, replyMessageID, detail, delivery)
}

func (ap *Process) removePendingDeliveryAtLocked(index int) {
	if index == 0 {
		if len(ap.pendingDeliveryQueue) > 0 {
			next := ap.pendingDeliveryQueue[0]
			ap.pendingDeliveryQueue = append([]pendingDelivery(nil), ap.pendingDeliveryQueue[1:]...)
			ap.activatePendingDeliveryLocked(next)
		} else {
			ap.pendingWake = nil
			ap.pendingWakeToken = ""
		}
		return
	}
	queueIndex := index - 1
	if queueIndex >= 0 && queueIndex < len(ap.pendingDeliveryQueue) {
		ap.pendingDeliveryQueue = append(ap.pendingDeliveryQueue[:queueIndex], ap.pendingDeliveryQueue[queueIndex+1:]...)
	}
}

func (ap *Process) removePendingDeliveryWithStatusLocked(index int, status, replyMessageID, detail string, delivery pendingDelivery) WakeStatusEvent {
	event := wakeStatusForDelivery(delivery, status, replyMessageID, detail)
	ap.removePendingDeliveryAtLocked(index)
	return event
}

func (ap *Process) terminalWakeStatusLocked(status string, replyMessageID string, detail string) WakeStatusEvent {
	delivery, ok := ap.activePendingDeliveryLocked()
	if !ok {
		return WakeStatusEvent{}
	}
	return ap.removePendingDeliveryWithStatusLocked(0, status, replyMessageID, detail, delivery)
}

func wakeStatusForDelivery(delivery pendingDelivery, status, replyMessageID, detail string) WakeStatusEvent {
	return WakeStatusEvent{
		MessageID:      delivery.Wake.MessageID,
		WakeToken:      delivery.Token,
		Target:         delivery.Wake.Target,
		Status:         status,
		ReplyMessageID: replyMessageID,
		Detail:         strings.TrimSpace(detail),
	}
}

func (ap *Process) terminalAllWakeStatusesLocked(status, detail string) []WakeStatusEvent {
	deliveries := ap.pendingDeliveriesLocked()
	if len(deliveries) == 0 {
		return nil
	}
	events := make([]WakeStatusEvent, 0, len(deliveries))
	for _, delivery := range deliveries {
		rowDetail := strings.TrimSpace(detail)
		if rowDetail == "" && status == consts.WakeStatusCompletedWithoutVisibleReply {
			if delivery.Wake.DeliveryClass == DeliveryClassSignalOnly {
				rowDetail = "delivery_class=signal_only"
			} else {
				rowDetail = "reason=no_reply"
			}
		}
		events = append(events, wakeStatusForDelivery(delivery, status, "", rowDetail))
	}
	ap.pendingWake = nil
	ap.pendingWakeToken = ""
	ap.pendingDeliveryQueue = nil
	return events
}