package agent

import (
	"log/slog"
	"strings"
	"time"

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

// ---------------------------------------------------------------------------
// Compaction helpers
// ---------------------------------------------------------------------------

func (d *CodexAppServerDriver) queueCompaction(threadID, causeTurnID string) {
	trimmedThreadID := strings.TrimSpace(threadID)
	if trimmedThreadID == "" {
		return
	}
	causeTurnID = strings.TrimSpace(causeTurnID)

	d.mu.Lock()
	if !d.isPrimaryThreadEventLocked(trimmedThreadID) {
		d.mu.Unlock()
		return
	}
	if d.compactQueued && strings.TrimSpace(d.queuedCompactionCauseTurnID) != "" {
		causeTurnID = d.queuedCompactionCauseTurnID
	}
	d.compactQueued = true
	d.queuedCompactionCauseTurnID = causeTurnID
	generation := d.launchGeneration
	d.mu.Unlock()

	slog.Info("codex terminal context overflow queued emergency compaction",
		"driver", "codex-app-server",
		"launch_generation", generation,
		"thread_id", trimmedThreadID,
		"cause_turn_id", causeTurnID,
		"will_retry", false)
}

func (d *CodexAppServerDriver) prepareIdleCompaction(sessionID string) (string, bool) {
	trimmedSessionID := strings.TrimSpace(sessionID)
	if trimmedSessionID == "" {
		return "", false
	}

	d.mu.Lock()
	defer d.mu.Unlock()
	if primaryThreadID := strings.TrimSpace(d.threadID); primaryThreadID != "" && primaryThreadID != trimmedSessionID {
		return "", false
	}
	if d.threadStatus.Type != "idle" {
		return "", false
	}
	if !d.compactQueued {
		return "", false
	}
	d.compactQueued = false
	return trimmedSessionID, true
}

func (d *CodexAppServerDriver) restoreQueuedCompaction(threadID, causeTurnID string) {
	trimmedThreadID := strings.TrimSpace(threadID)
	if trimmedThreadID == "" {
		return
	}
	causeTurnID = strings.TrimSpace(causeTurnID)

	d.mu.Lock()
	defer d.mu.Unlock()
	if primaryThreadID := strings.TrimSpace(d.threadID); primaryThreadID != "" && primaryThreadID != trimmedThreadID {
		return
	}
	if d.compactionInProgress {
		causeTurnID = firstNonEmptyString(d.compactionCauseTurnID, causeTurnID)
		d.clearCompactionMutationLocked()
		if d.threadStatus.Type == "active" && strings.TrimSpace(d.activeTurnID) == "" {
			d.threadStatus = appServerThreadStatus{Type: "idle"}
			d.activeTurnNonSteerableKind = ""
		}
	}
	d.compactQueued = true
	d.queuedCompactionCauseTurnID = causeTurnID
}

func (d *CodexAppServerDriver) beginCompaction(threadID string) {
	trimmedThreadID := strings.TrimSpace(threadID)
	if trimmedThreadID == "" {
		return
	}

	d.mu.Lock()
	defer d.mu.Unlock()
	if primaryThreadID := strings.TrimSpace(d.threadID); primaryThreadID != "" && primaryThreadID != trimmedThreadID {
		return
	}
	d.beginCompactionLocked(
		trimmedThreadID,
		codexCompactionSourceAgentTeamEmergency,
		d.queuedCompactionCauseTurnID,
	)
	d.queuedCompactionCauseTurnID = ""
}

func (d *CodexAppServerDriver) beginCompactionLocked(threadID, source, causeTurnID string) {
	if d.threadStatus.Type == "idle" {
		d.threadStatus = appServerThreadStatus{Type: "active"}
		d.activeTurnID = ""
		d.activeTurnNonSteerableKind = ""
	}
	d.compactionInProgress = true
	d.compactionTurnID = ""
	d.compactionCauseTurnID = strings.TrimSpace(causeTurnID)
	d.compactionSource = strings.TrimSpace(source)
	d.compactionTerminalStatus = appServerThreadStatus{}
}

func (d *CodexAppServerDriver) clearCompactionMutationLocked() {
	d.compactionInProgress = false
	d.compactionTurnID = ""
	d.compactionCauseTurnID = ""
	d.compactionSource = ""
	d.compactionTerminalStatus = appServerThreadStatus{}
}

// proactiveCompactionDueLocked reports whether Codex has a current
// active-context signal at or above the shared proactive threshold. The caller
// must hold tokenUsageMu and mu, in that order.
func (d *CodexAppServerDriver) proactiveCompactionDueLocked(sessionID string) bool {
	if d.proactiveCompactionThreshold == 0 ||
		strings.TrimSpace(d.threadID) != strings.TrimSpace(sessionID) ||
		d.threadStatus.Type != "idle" ||
		strings.TrimSpace(d.activeTurnID) != "" ||
		d.compactQueued ||
		d.compactionInProgress {
		return false
	}
	usage := d.lastTokenUsageByThread[strings.TrimSpace(sessionID)]
	if nonNegativeInt64TokenCount(usage.Last.TotalTokens) < d.proactiveCompactionThreshold {
		return false
	}
	if !d.lastContextCompactionCompletedAt.IsZero() &&
		time.Since(d.lastContextCompactionCompletedAt) < proactiveIdleCompactionCooldown {
		return false
	}
	return true
}

func (d *CodexAppServerDriver) restoreRejectedProactiveCompaction(threadID string) bool {
	threadID = strings.TrimSpace(threadID)
	d.mu.Lock()
	defer d.mu.Unlock()
	if strings.TrimSpace(d.threadID) != threadID ||
		d.compactionSource != codexCompactionSourceLoopProactive ||
		strings.TrimSpace(d.compactionTurnID) != "" {
		return false
	}
	d.clearCompactionMutationLocked()
	if d.threadStatus.Type == "active" && strings.TrimSpace(d.activeTurnID) == "" {
		d.threadStatus = appServerThreadStatus{Type: "idle"}
		d.activeTurnNonSteerableKind = ""
	}
	return true
}

func (d *CodexAppServerDriver) handleContextCompactionItem(params appServerItemNotification, completed bool) {
	threadID := strings.TrimSpace(params.ThreadID)
	turnID := strings.TrimSpace(params.TurnID)
	itemID := strings.TrimSpace(params.Item.ID)

	d.mu.Lock()
	if !d.isPrimaryThreadEventLocked(threadID) {
		d.mu.Unlock()
		return
	}
	identity := codexContextCompactionIdentity(threadID, turnID, itemID)
	generation := d.launchGeneration
	if !completed {
		currentIdentity := codexContextCompactionIdentity(
			d.contextCompactionItemFence.threadID,
			d.contextCompactionItemFence.turnID,
			d.contextCompactionItemFence.itemID,
		)
		if identity != "" &&
			(identity == currentIdentity || identity == d.lastContextCompactionCompletion) {
			d.mu.Unlock()
			return
		}
		source, causeTurnID := d.contextCompactionSourceLocked(threadID, turnID)
		if identity == "" {
			source = codexCompactionSourceUnknown
		}
		before := d.lastTokenUsageByThread[threadID]
		beforeKnown := appServerThreadTokenUsageKnown(before)
		pairingConflict := currentIdentity != "" && currentIdentity != identity
		if identity != "" {
			d.contextCompactionItemFence = codexContextCompactionItemFence{
				threadID:        threadID,
				turnID:          turnID,
				itemID:          itemID,
				pairingConflict: pairingConflict,
				before:          before,
			}
		}
		d.mu.Unlock()

		slog.Info("codex context compaction started",
			"driver", "codex-app-server",
			"launch_generation", generation,
			"source", source,
			"thread_id", threadID,
			"turn_id", turnID,
			"item_id", itemID,
			"cause_turn_id", causeTurnID,
			"identity_complete", identity != "",
			"pairing_conflict", pairingConflict,
			"before_known", beforeKnown,
			"before_active_context_tokens", nonNegativeInt64TokenCount(before.Last.TotalTokens),
			"before_context_window_tokens", codexContextWindowTokens(before))
		return
	}

	if identity != "" && d.lastContextCompactionCompletion == identity {
		d.mu.Unlock()
		return
	}
	fenceIdentity := codexContextCompactionIdentity(
		d.contextCompactionItemFence.threadID,
		d.contextCompactionItemFence.turnID,
		d.contextCompactionItemFence.itemID,
	)
	paired := identity != "" && identity == fenceIdentity
	source := codexCompactionSourceUnknown
	var boundary codexContextCompactionBoundary
	if paired {
		causeTurnID := ""
		source, causeTurnID = d.contextCompactionSourceLocked(threadID, turnID)
		if d.contextCompactionItemFence.pairingConflict {
			source = codexCompactionSourceUnknown
		}
		boundary = codexContextCompactionBoundary{
			launchGeneration: generation,
			threadID:         threadID,
			turnID:           turnID,
			itemID:           itemID,
			source:           source,
			causeTurnID:      causeTurnID,
			before:           d.contextCompactionItemFence.before,
			after:            d.contextCompactionItemFence.during,
		}
		d.contextCompactionBoundary = boundary
		d.contextCompactionItemFence = codexContextCompactionItemFence{}
	} else {
		boundary = codexContextCompactionBoundary{
			launchGeneration: generation,
			threadID:         threadID,
			turnID:           turnID,
			itemID:           itemID,
			source:           source,
		}
	}
	if identity != "" {
		d.lastContextCompactionCompletion = identity
	}
	d.lastContextCompactionCompletedAt = time.Now()
	d.mu.Unlock()

	slog.Info("codex context compaction completed",
		"driver", "codex-app-server",
		"launch_generation", boundary.launchGeneration,
		"source", boundary.source,
		"thread_id", boundary.threadID,
		"turn_id", boundary.turnID,
		"item_id", boundary.itemID,
		"cause_turn_id", boundary.causeTurnID,
		"identity_complete", identity != "",
		"paired", paired,
		"before_known", appServerThreadTokenUsageKnown(boundary.before),
		"before_active_context_tokens", nonNegativeInt64TokenCount(boundary.before.Last.TotalTokens),
		"before_context_window_tokens", codexContextWindowTokens(boundary.before),
		"after_known", appServerThreadTokenUsageKnown(boundary.after),
		"after_active_context_tokens", nonNegativeInt64TokenCount(boundary.after.Last.TotalTokens),
		"after_context_window_tokens", codexContextWindowTokens(boundary.after))
}

func (d *CodexAppServerDriver) contextCompactionSourceLocked(threadID, turnID string) (string, string) {
	threadID = strings.TrimSpace(threadID)
	turnID = strings.TrimSpace(turnID)
	if threadID == "" || turnID == "" || strings.TrimSpace(d.threadID) != threadID {
		return codexCompactionSourceUnknown, ""
	}
	if d.compactionInProgress && strings.TrimSpace(d.compactionTurnID) == turnID {
		switch d.compactionSource {
		case codexCompactionSourceAgentTeamEmergency, codexCompactionSourceLoopProactive:
			return d.compactionSource, strings.TrimSpace(d.compactionCauseTurnID)
		default:
			return codexCompactionSourceUnknown, strings.TrimSpace(d.compactionCauseTurnID)
		}
	}
	if strings.TrimSpace(d.activeTurnID) == turnID {
		return codexCompactionSourceRuntimeNative, ""
	}
	return codexCompactionSourceUnknown, ""
}

func (d *CodexAppServerDriver) agentteamCompactionFailureCorrelation(
	threadID string,
	turnID string,
) (uint64, string, string, bool, bool) {
	threadID = strings.TrimSpace(threadID)
	turnID = strings.TrimSpace(turnID)
	d.mu.Lock()
	defer d.mu.Unlock()
	if threadID == "" ||
		turnID == "" ||
		strings.TrimSpace(d.threadID) != threadID ||
		!d.compactionInProgress ||
		strings.TrimSpace(d.compactionTurnID) != turnID {
		return 0, "", "", false, false
	}
	_, terminalAlreadyObserved := d.completedTurnIDs[turnID]
	return d.launchGeneration, d.compactionSource, strings.TrimSpace(d.compactionCauseTurnID), !terminalAlreadyObserved, true
}

func (d *CodexAppServerDriver) observeContextCompactionUsageLocked(
	threadID string,
	turnID string,
	usage appServerThreadTokenUsage,
) {
	fence := &d.contextCompactionItemFence
	if threadID == "" || turnID == "" ||
		threadID != fence.threadID || turnID != fence.turnID {
		return
	}
	fence.during = usage
}

func (d *CodexAppServerDriver) consumeContextCompactionBoundaryLocked(
	threadID string,
	usage appServerThreadTokenUsage,
) (bool, *codexContextCompactionBoundary) {
	if threadID == "" ||
		threadID != strings.TrimSpace(d.threadID) ||
		threadID != d.contextCompactionBoundary.threadID {
		return false, nil
	}
	boundary := d.contextCompactionBoundary
	var postCompactionBoundary *codexContextCompactionBoundary
	if !appServerThreadTokenUsageKnown(boundary.after) {
		boundary.after = usage
		postCompactionBoundary = &boundary
	}
	d.contextCompactionBoundary = codexContextCompactionBoundary{}
	return true, postCompactionBoundary
}

func logCodexContextCompactionUsage(boundary codexContextCompactionBoundary) {
	slog.Info("codex context compaction post-usage observed",
		"driver", "codex-app-server",
		"launch_generation", boundary.launchGeneration,
		"source", boundary.source,
		"thread_id", boundary.threadID,
		"turn_id", boundary.turnID,
		"item_id", boundary.itemID,
		"cause_turn_id", boundary.causeTurnID,
		"before_known", appServerThreadTokenUsageKnown(boundary.before),
		"before_active_context_tokens", nonNegativeInt64TokenCount(boundary.before.Last.TotalTokens),
		"before_context_window_tokens", codexContextWindowTokens(boundary.before),
		"after_known", appServerThreadTokenUsageKnown(boundary.after),
		"after_active_context_tokens", nonNegativeInt64TokenCount(boundary.after.Last.TotalTokens),
		"after_context_window_tokens", codexContextWindowTokens(boundary.after))
}

func codexContextCompactionCompletedEvent(
	threadID string,
	boundary codexContextCompactionBoundary,
) translatedEvent {
	threadID = strings.TrimSpace(threadID)
	var beforeTokens, afterTokens uint64
	var metadata map[string]any
	if boundary.threadID == threadID {
		beforeTokens = nonNegativeInt64TokenCount(boundary.before.Last.TotalTokens)
		afterTokens = nonNegativeInt64TokenCount(boundary.after.Last.TotalTokens)
		if source := strings.TrimSpace(boundary.source); source != "" {
			metadata = map[string]any{"context_compaction_source": source}
		}
	}
	return translatedEvent{
		Kind:      consts.ParsedEventKindRuntimeCompacted,
		SessionID: threadID,
		Message:   formatContextCompactionCompletedDetail(beforeTokens, afterTokens),
		Metadata:  metadata,
	}
}

func codexContextCompactionIdentity(threadID, turnID, itemID string) string {
	threadID = strings.TrimSpace(threadID)
	turnID = strings.TrimSpace(turnID)
	itemID = strings.TrimSpace(itemID)
	if threadID == "" || turnID == "" || itemID == "" {
		return ""
	}
	return threadID + "\x00" + turnID + "\x00" + itemID
}

func appServerThreadTokenUsageKnown(usage appServerThreadTokenUsage) bool {
	return usage.ModelContextWindow != nil ||
		usage.Last.InputTokens != 0 ||
		usage.Last.CachedInputTokens != 0 ||
		usage.Last.OutputTokens != 0 ||
		usage.Last.ReasoningOutputTokens != 0 ||
		usage.Last.TotalTokens != 0
}

func codexContextWindowTokens(usage appServerThreadTokenUsage) uint64 {
	if usage.ModelContextWindow == nil {
		return 0
	}
	return nonNegativeInt64TokenCount(*usage.ModelContextWindow)
}